#  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 viconparser (c) Shai Revzen, U Penn, 2010

import re
from copy import copy
from numpy import all,asarray,dtype,float32,frombuffer,fromfile,nan,newaxis,nonzero,unique
from numpy.linalg import norm

OBJ = re.compile("dat='[xyar]', cam=(\d+), obj=(\d+), .*")
FPS = re.compile('fps=(\d+), .*')

def find(ar):
    "find indices where ar is python true"
    res, = nonzero(asarray(ar,bool).flat)
    return res

def iwindow( N, p, mapping=copy):
  """Iterator returning a list consisting of a sliding window 
  of N consecutive items from iterator p. 
  By default, a copy of each sliding window is returned. 
  Setting mapping to some other callable that expects a list
  allows other functions of the sliding window to be computed 
  without an unnecessary copy operation.
  
  >>> [ copy.copy(x) for x in VP.iwindow(2,xrange(5)) ]
  [[0, 1], [1, 2], [2, 3], [3, 4]]  

  WARNING: iwindow returns the *SAME* list object over again, so
  >>> [ x for x in VP.iwindow(2,xrange(5)) ]
  [[4], [4], [4], [4]]  
  """
  W = []
  for item in p:
    W.append(mapping(item))
    if len(W)==N:
      yield W
      W.pop(0)
        
class OptiParser(object):
  """OptiParser parses the output of OptiReader-s into numpy data
  
  Typical usage reading from a file:
  >>> P = OptiParser()
  >>> P.load("run1")
  >>> print abs( P.x + 1j* P.y )
  
  Typical use for a live stream:
  >>> P = OptiParser()
  >>> R = OptiReader()
  >>> P.useInfo( R.connect() )
  >>> for now in P.parseStream( R.stream() ):
  ...   print "%6.2f %g,%g" % (now, mean(P.x,0), mean(P.y,0) )
  
  The OptiParse creates the following self.attributes:
    fps -- frames per second data rate of the mocap
    t -- N -- timestamps (seconds)
    x,y,a,r -- N x M -- x, y, area, radius of M objects over N times
  """
  def __init__( self, src = None ):
    self.ind = None
  
  def useInfo( self, dcr ):
    """parse the descriptors from info packet into an indexing scheme"""
    # Number of columns
    self.cols = len(dcr)
    # Find time data
    fm = [ FPS.match(x) for x in dcr ]
    i_t = find(fm)
    # Extract fps
    self.fps = int( fm[i_t].group(1) )
    #
    # Find xyar data
    fm = [ OBJ.match(x) for x in dcr ]
    i_xyar = find(fm)
    self.cams = set( [f.group(1) for f in fm if f] )
    self.objs = set( [f.group(2) for f in fm if f] )
    #
    self.ind = dict(
      t = i_t,
      xyar = i_xyar
      )

  def parseStream( self, idat ):
    """Parse data from a stream (python generator). 
    P.parseStream() updates the values in P IN PLACE and yields timestamp.

    If your stream returns data that should not be used in place, use a generator
    comprehension to copy it, e.g. parseStream( ( copy.copy(x) for x in idat ) )
    """
    # loop over input and use sliding window data for updates
    for dat in idat:
      self.parseData( dat )
      yield self.t[-1]
      
  def parseData( self, dat ):
    """Parse a data blob into useful data
       dat -- data block as array of floats or some buffer with float32-s
       
       data is used *IN PLACE* without any copy operations
    """
    # If a string --> parse into floating point numbers
    if (type(dat)==str):
      d = frombuffer( dat, float32 )
      d.shape = (d.size / self.cols, self.cols)
    else: # else --> must be something that becomes an array
      d = asarray(dat,float32)
      # If rank 1 array --> reshape it
      if d.ndim != 2:
        # Reshape
        d.shape = (d.size / self.cols, self.cols)
    # Create all columns
    for key,val in self.ind.iteritems():
      col = d[:,val]
      setattr( self, key, col )
    self.xyar.shape = (self.xyar.shape[0],self.xyar.shape[1]/4,4)
    self.x = self.xyar[...,0]
    self.y = self.xyar[...,1]
    self.a = self.xyar[...,2]
    self.r = self.xyar[...,3]

  def load( self, filename, dt=dtype('float32') ):
    """Load OptiReader data from filename.dcr and filename.dat"""
    # Load columns descriptors
    f = open(filename+".dcr",'r')
    dcr = f.readlines()
    f.close()
    # Load data
    f = open(filename+".dat",'rb')
    dat = fromfile( f, dtype=dt )
    try:
      dat.shape = ( dat.size / len(dcr), len(dcr) )
    except ValueError:
      print "dat.size = %d not divisibile by len(dcr) = %d -- removing %d datum" % (dat.size,len(dcr),dat.size - (dat.size/len(dcr))*len(dcr))
      dat = dat[:(dat.size/len(dcr)) * len(dcr)]
      dat.shape = ( dat.size / len(dcr), len(dcr) )

    self.useInfo( dcr ); self.dcr = dcr
    self.parseData( dat ); self.dat = dat