#!/usr/bin/env python

import sys

import matplotlib.mlab as pylab
import matplotlib.pyplot as plt
import numpy as np

import cv
import utils
import cvutils
import ubc_utils
import moving

# use something like getopt to manage arguments if necessary

if len(sys.argv) < 3:
    print('Usage: {0} <video-filename> <n-objects>'.format(sys.argv[0]))
    sys.exit()

if sys.argv[1].endswith('.avi'):
    videoFilenamePrefix = utils.removeExtension(sys.argv[1],'.')
else:
    videoFilenamePrefix = sys.argv[1]

objectNum = int(sys.argv[2])

objects = ubc_utils.loadTrajectories(videoFilenamePrefix+'-objects.txt', objectNum+1)
obj = objects[objectNum]
features = ubc_utils.loadTrajectories(videoFilenamePrefix+'-features.txt', max(obj.featureNumbers)+1)
h = np.loadtxt(videoFilenamePrefix+'-homography.txt')

invh = cvutils.invertHomography(h)

def computeObject(features, homography, timeInterval = None):
    if not timeInterval:
        raise Exception('not implemented') # compute from the features
    
    yCoordinates = -np.ones((len(features),int(timeInterval.length())))
    for i,f in enumerate(features):
        traj = f.getPositions().asArray()
        imgTraj = cvutils.projectArray(homography, traj)
        yCoordinates[i,f.getFirstInstant()-timeInterval.first:f.getLastInstant()+1-timeInterval.first] = imgTraj[1,:]

    indices = np.argmax(yCoordinates,0)
    newTraj = moving.Trajectory()
    for j,idx in enumerate(indices):
        newTraj.addPosition(features[idx].getPositionAtInstant(j+timeInterval.first))
        #newVelocities.addPosition(features[obj.featureNumbers[idx]].getVelocityAtInstant(j+obj.getFirstInstant()))

    return newTraj

# TODO version median

def kalmanFilter(positions, velocities, processNoiseCov, measurementNoiseCov):
    kalman=cv.CreateKalman(6, 4)
    kalman.transition_matrix[0,2]=1
    kalman.transition_matrix[0,4]=1./2
    kalman.transition_matrix[1,3]=1
    kalman.transition_matrix[1,5]=1./2
    kalman.transition_matrix[2,4]=1
    kalman.transition_matrix[3,5]=1

    cv.SetIdentity(kalman.measurement_matrix, 1.)
    cv.SetIdentity(kalman.process_noise_cov, processNoiseCov)
    cv.SetIdentity(kalman.measurement_noise_cov, measurementNoiseCov)
    cv.SetIdentity(kalman.error_cov_post, 1.)

    p = positions[0]
    v = velocities[0]
    v2 = velocities[2]
    a = (v2-v).multiply(0.5)
    kalman.state_post[0,0]=p.x
    kalman.state_post[1,0]=p.y
    kalman.state_post[2,0]=v.x
    kalman.state_post[3,0]=v.y
    kalman.state_post[4,0]=a.x
    kalman.state_post[5,0]=a.y
    
    filteredPositions = moving.Trajectory()
    filteredVelocities = moving.Trajectory()
    measurement = cv.CreateMat(4,1,cv.CV_32FC1)
    for i in xrange(positions.length()):
        cv.KalmanPredict(kalman) # no control
        p = positions[i]
        v = velocities[i]
        measurement[0,0] = p.x
        measurement[1,0] = p.y
        measurement[2,0] = v.x
        measurement[3,0] = v.y
        cv.KalmanCorrect(kalman, measurement)
        filteredPositions.addPositionXY(kalman.state_post[0,0], kalman.state_post[1,0])
        filteredVelocities.addPositionXY(kalman.state_post[2,0], kalman.state_post[3,0])

    return (filteredPositions, filteredVelocities)

(filteredPositions, filteredVelocities) = kalmanFilter(newTraj, obj.getVelocities(), 1e-5, 1e-1)

plt.clf()
obj.draw('rx')
for fnum in obj.featureNumbers: features[fnum].draw()
newTraj.draw('bx')
plt.axis('equal')
