#  This library is free software; you can redistribute it and/or
#  modify it under the terms of the GNU General Public
#  License as published by the Free Software Foundation; either
#  version 3.0 of the License, or (at your option) any later version.
#
#  The library is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
#  General Public License for more details.
#
# (c) Sam Burden, UC Berkeley, 2012
#
# based on viconreader (c) Shai Revzen, U Penn, 2010

import struct
import socket 
import sys
import time

DEFAULT_HOST = "localhost"  
DEFAULT_PORT = 43949

# enforce network byte order
pack   = lambda fmt,*vals : struct.pack( '!'+fmt, *vals)
unpack = lambda fmt,str : struct.unpack( '!'+fmt, str)

  
class OptiReader(object):
  """
  OptiReader instances provide access to the real-time data stream 
  made available by opti.exe
  
  Instances of OptiReader can send QUERY, START and STOP messages and parse 
  responses into lists of attributes and blocks of doubles (encoded as strings).
  
  OptiReader is used from the commandline to log tracking data to a file.
  
  Typical use is:
  >>> R = OptiReader()
  >>> # get list of names from QUERY Response
  >>> names = R.connect() 
  >>> # R.stream() is a generator returning packet payload each .next()
  >>> for pkt,_ in zip(R.stream(),xrange(100)): 
  >>>   dat = R.unpack("%dd" % (len(pkt) / R.DATLEN),pkt)
  >>>   print
  >>>   for nm,val in zip(name,dat):
  >>>     print nm,val
  >>> R.stop()
  """
  QUERY = 1
  INFO  = 2
  START = 3
  STOP  = 4
  DATA  = 5

  # number of bytes per unit of data
  DATLEN = 4
  # number of bytes per command
  CMDLEN = 4

  if struct.calcsize('L') == CMDLEN:
    CMD = 'L'
  elif struct.calcsize('I') == CMDLEN:
    CMD = 'I'
  else:
    raise OSError,'Cannot find %d byte unsigned int for your OS' % CMDLEN

  def __init__(self):
    self.sock = None
    self.push = ''

  def _get( self, fmt ):
    "Read data from socket based on format string, and parse it accordingly"
    N = struct.calcsize(fmt)
    # Start reading from push-back buffer
    buf = self.push[:min(len(self.push),N)]
    self.push = self.push[:len(buf)]
    while len(buf)<N:
      buf += self.sock.recv(N-len(buf))
    return unpack(fmt,buf)

  def _parseInfo( self ):
    "Parse an INFO packet, starting with byte after the header"
    N = self._get( self.__class__.CMD )[0]
    lst = []
    for _ in xrange(N):
      L = self._get( self.__class__.CMD )[0]
      lst.append(self._get("%ds" % L)[0])
    return lst
  
  def _parseData( self ):
    "Parse a DATA packet, starting with byte after the header"
    N = self._get( self.__class__.CMD )[0]
    return self._get("%ds" % (N * self.__class__.DATLEN))[0]
  
  def _parse( self ):
    "Parse an incoming packet"    
    hdr = self._get( self.__class__.CMD )[0]
    # parse data
    if hdr == self.__class__.DATA:
      return (hdr, self._parseData())
    # parse info
    elif hdr == self.__class__.INFO:
      return (hdr, self._parseInfo())

    # Failed -- need to resync
    self.push = pack( self.__class__.CMD, hdr)
    self._resync()
  
  def _resync(self):
    raise ValueError,"Lost synchronization on socket"

  def _cmd( self, hdr ):
    "Command encoded as a header of length CMDLEN"
    self.sock.send(pack( self.__class__.CMD, hdr ))
  
  def connect( self, host=DEFAULT_HOST, port=DEFAULT_PORT ):
    # Connect the socket
    self.sock = socket.socket( socket.AF_INET, socket.SOCK_STREAM )
    self.sock.connect((host,port))
    # Send a query
    self._cmd( self.__class__.QUERY )
    # Loop until a query response is received
    while True:
      hdr,names = self._parse()
      if hdr == self.__class__.INFO:
        return names
          
  def stream( self, fps=1e9 ):
    """
    Generator producing a stream of data packet payloads

    INPUTS
      fps -- expected fps, for throwing away stale data. 
        If not set, all data will be returned. Set to the expected vicon 
        frames-per-second data capture rate; typically 120
    """
    self._cmd( self.__class__.START )
    thresh = 0.5/fps
    while True:
      # Make sure we enter data reading loop at least once
      last = time.time() + 1e6
      while time.time() < thresh + last:
        last = time.time()        
        hdr,data = self._parse()
      if hdr == self.__class__.DATA:
        yield data

  def stop( self ):
    "Tell opti to stop streaming"
    self._cmd( self.__class__.STOP )
  
  def close( self ):
    "Close connection to opti"
    self.sock.close()
    self.sock = None

if __name__=="__main__":
  hosts = dict( 
    local = "localhost",
    winlocal = "127.0.0.2",
    olin =  "10.33.25.15",
    DEFAULT = "localhost" 
  )
  fn = None
  host = hosts['DEFAULT']
  port = DEFAULT_PORT
  argv = []; argv[:] = sys.argv[1:]
  err = ''
  while argv:
    arg = argv.pop(0).lower()
    if arg in ['-h','--host']:
      arg = argv.pop(0)
      host = hosts.get( arg.lower(), arg )
      continue
    if arg in ['-p','--port']:
      port = int(argv.pop(0))
      continue
    if arg in ['--help','--usage']:
      err = ' '
      break
    if arg in ['-a','--aliases']:
      for k,v in hosts.iteritems():
        print "%-16s : %s" % (k,v)
      sys.exit(0)
    if arg[0] != '-':
      if fn is not None:
        err = 'Only one filename allowed on commandline; got "%s" and "%s"' % (fn,arg)
        break
      fn = arg
      continue
    err = 'Unknown commandline argument "%s"' % arg
    break
  if fn is None:
    err = "Must specify filename on the commandline"
  if err:
    sys.stderr.write("%s\n" % err)
    sys.stderr.write("""
Usage: %s {options} filename

  Creates filename.dcr with the field names and filename.dat with the raw
  data stored as floats.
  
  Options:
   -h <hostname>, --host <hostname> 
        Specify opti.exe hostname where hostname is an IP address, FQDN or alias 

   -p <port>, --port <port> 
        Specify opti.exe port to connect to 

   -a, --aliases
        List the builtin host aliases
   
   --help, --usage
        Print this message
  """ % sys.argv[0])
    sys.exit(5)
  # Filename user requested
  print "Connecting to %s:%d" % (host,port)   
  R = OptiReader()
  names = R.connect(host=host, port=port)
  # Write descriptor
  dcr = open(fn+".dcr", "w")
  dcr.write("\n".join(names)+"\n")
  dcr.close()
  # Write data
  L = 0
  N = 0
  dat = []
  try: # for catching ctrl-c termination
    print "Recording to %s.dcr, %s.raw" % (fn,fn)   
    for pkt in R.stream():
      if not L: # First packet sets the length we'll expect 
        L = len(pkt)
      elif L != len(pkt):
        print "\nUnexpected packet of length %d instead of %d" % (len(pkt),L)
        continue
      #P dat.write(pkt)
      dat.append(pkt)
      N += 1
      if (N%50)==0:
        print "\r%5d" % N,
    # ---
  except KeyboardInterrupt, ex:
    #P dat.close()
    print "Writing data -- WAIT !!!"
    f = open(fn+".dat", "w")
    f.write("".join(dat))
    f.close()
    R.stop()
    R.close()
    print "\ndone."