"""
Test script to compare feature tracking with NCC using experimental results 
"""
#%%-----------------------------------------------------------------------------
# setup
#-------------------------------------------------------------------------------
import numpy as np
import matplotlib.pyplot as plt
plt.ion()
#set math text style
import matplotlib as mpl
mpl.rcParams['mathtext.default'] = 'regular'

import time
import cv2 as cv        #opencv v3.1.0
import lib_feature_tracking as lib

##Load data
imgs = np.load("linear move direction=0.0 speed=5.0mms-1.npy")


##Other setup 
win = (128,128,256,256)     #window size to use (x0,y0,w,h)

#%%-----------------------------------------------------------------------------
# Setup matcher methods
#-------------------------------------------------------------------------------
ncc = lib.NCC()
orb = lib.ORB(patchSize=64, fastThreshold=20, distance_threshold=64 )
brisk = lib.BRISK(fastThreshold=18, patternScale=1.8, distance_threshold=50 )
usurf = lib.USURF(threshold=40, distance_threshold=0.5, norm=cv.NORM_L1 )
surf = lib.SURF(threshold=40, distance_threshold=0.5, norm=cv.NORM_L1 )


#%%-----------------------------------------------------------------------------
# process multiple speckle patterns and average
#-------------------------------------------------------------------------------
#clear all stored results
ncc.reset()
orb.reset()
brisk.reset()
usurf.reset()
surf.reset()

#set initial reference image
img0 = imgs[0,win[1]:win[1]+win[3], win[0]:win[0]+win[2]]
N = 200 #max matches to use

tic = time.clock()
for n in xrange(1, imgs.shape[0],1):
    
    img1 = imgs[n,win[1]:win[1]+win[3], win[0]:win[0]+win[2]]
    
    ##Do NCC processsing
    Ax,Ay = ncc.calc_shift(img0,img1)
    #print '\tncc: {0:0.2f}{1:0.2f}'.format(Ax,Ay)
    
    ##Do ORB
    Ax,Ay,theta = orb.calc_shift_and_rotation(img0, img1, adj_fast_threshold=False, maxMatches=N)
    #print '\torb: {0:0.2f}{1:0.2f}{1:0.2f}'.format(Ax,Ay,theta)
    
    ##Do BRISK
    Ax,Ay,theta = brisk.calc_shift_and_rotation(img0, img1, maxMatches=N)
    #print '\tbrisk: {0:0.2f}{1:0.2f}{1:0.2f}'.format(Ax,Ay,theta)
    
    ##Do USURF
    Ax,Ay,theta = usurf.calc_shift_and_rotation(img0, img1, maxMatches=N)
    #print '\tusurf: {0:0.2f}{1:0.2f}{1:0.2f}'.format(Ax,Ay,theta)
    
    ##Do SURF
    Ax,Ay,theta = surf.calc_shift_and_rotation(img0, img1, maxMatches=N)
    #print '\tsurf: {0:0.2f}{1:0.2f}{1:0.2f}'.format(Ax,Ay,theta)

    #set new reference
    img0 = img1
    
toc = time.clock()
dt = (toc-tic)/(n+1)
print dt*1000, 'ms per measurment'
    

#make results arrays
Axs_ncc = np.array(ncc.Axs)
Ays_ncc = np.array(ncc.Ays)

Axs_orb = np.array(orb.Axs)
Ays_orb = np.array(orb.Ays)

Axs_brisk = np.array(brisk.Axs)
Ays_brisk = np.array(brisk.Ays)

Axs_usurf = np.array(usurf.Axs)
Ays_usurf = np.array(usurf.Ays)

Axs_surf = np.array(surf.Axs)
Ays_surf = np.array(surf.Ays)

ts = np.arange(0, imgs.shape[0])*(1/500.0) # 500fps

#%%-----------------------------------------------------------------------------
# Comparision plot - orb vs ncc
#-------------------------------------------------------------------------------
f = plt.figure(figsize=(6,4))

ax1 = f.add_axes( [0.1,0.55,0.8,0.4])

ax1.plot(Axs_ncc, '-r', label='$A_x$ shift [ncc]', ms=4, lw=2 )
ax1.plot(Ays_ncc, '--b', label='$A_y$ shift [ncc]', ms=4, lw=2)

ax1.plot(Axs_orb, 'xr', label='$A_x$ shift [ORB]', ms=4, lw=1)
ax1.plot(Ays_orb, 'xb', label='$A_y$ shift [ORB]',ms=4, lw=1)

ax1.set_ylim(-12,12)
ax1.set_ylabel('shift [pixels]')
ax1.yaxis.labelpad = 8
plt.legend(loc=8,ncol=2, columnspacing=0.5, numpoints=1,frameon=False)

#error
ax2 = f.add_axes( [0.1,0.1,0.8,0.4])
errx = Axs_orb-Axs_ncc
erry = Ays_orb-Ays_ncc

ax2.plot( errx, '-r', label='error $d_x = A_x - Ax_ncc$')
ax2.plot( erry, '--b', label='error $d_y = A_y - Ay_ncc$', lw=2)

ax2.set_ylabel('difference [pixels]')
ax2.set_xlabel('frame number')

#set font sizes
texts = f.findobj(plt.Text)
res = plt.setp(texts, fontsize=8)

#f.savefig('ORB vs NCC.pdf', dpi=300, transparent=True)

#%%-----------------------------------------------------------------------------
# Comparision plot - brisk vs ncc
#-------------------------------------------------------------------------------
f = plt.figure(figsize=(6,4))

ax1 = f.add_axes( [0.1,0.55,0.8,0.4])

ax1.plot(Axs_ncc, '-r', label='$A_x$ shift [ncc]', ms=4, lw=2 )
ax1.plot(Ays_ncc, '--b', label='$A_y$ shift [ncc]', ms=4, lw=2)

ax1.plot(Axs_brisk, 'xr', label='$A_x$ shift [ORB]', ms=4, lw=1)
ax1.plot(Ays_brisk, 'xb', label='$A_y$ shift [ORB]',ms=4, lw=1)

ax1.set_ylim(-12,12)
ax1.set_ylabel('shift [pixels]')
ax1.yaxis.labelpad = 8
plt.legend(loc=8,ncol=2, columnspacing=0.5, numpoints=1,frameon=False)

#error
ax2 = f.add_axes( [0.1,0.1,0.8,0.4])
errx = Axs_brisk-Axs_ncc
erry = Ays_brisk-Ays_ncc

ax2.plot( errx, '-r', label='error $d_x = A_x - Ax_ncc$')
ax2.plot( erry, '--b', label='error $d_y = A_y - Ay_ncc$', lw=2)

ax2.set_ylabel('difference [pixels]')
ax2.set_xlabel('frame number')

#set font sizes
texts = f.findobj(plt.Text)
res = plt.setp(texts, fontsize=8)

#f.savefig('BRISK vs NCC.pdf', dpi=300, transparent=True)

#%%-----------------------------------------------------------------------------
# Comparision plot - USURF vs ncc
#-------------------------------------------------------------------------------
f = plt.figure(figsize=(6,4.0))

ax1 = f.add_axes( [0.1,0.55,0.8,0.4])

ax1.plot(Axs_ncc, '-r', label='$A_x$ [NCC]', ms=4, lw=2 )
ax1.plot(Ays_ncc, '--b', label='$A_y$ [NCC]', ms=4, lw=2)

ax1.plot(Axs_usurf, 'xr', label='$A_x$ [USURF]', ms=4, lw=1)
ax1.plot(Ays_usurf, 'xb', label='$A_y$ [USURF]',ms=4, lw=1)

ax1.set_ylim(-15,10)
ax1.set_ylabel('shift [pixels]')
#ax1.set_yticks((-10,-5,0,5,10))
ax1.set_yticklabels( ('-15.0', '-10.0','-5.0','0.0','5.0','10.0'))
ax1.yaxis.labelpad = 8
plt.legend(loc=8,ncol=2, columnspacing=0.5, numpoints=1,frameon=False)

#error
ax2 = f.add_axes( [0.1,0.1,0.8,0.4])
errx = Axs_usurf-Axs_ncc
erry = Ays_usurf-Ays_ncc

ax2.plot( errx, '-r', label='Difference $A_x$',lw=2)
ax2.plot( erry, '--b', label='Difference $A_y$', lw=2)
ax2.set_ylim(-0.2,0.2)

ax2.set_ylabel('difference [pixels]')
ax2.set_xlabel('frame number')

plt.legend(loc=8,ncol=2,frameon=False)

#set font sizes
texts = f.findobj(plt.Text)
res = plt.setp(texts, fontsize=8)

#f.savefig('USURF vs NCC.pdf', dpi=300, transparent=True)

#%%-----------------------------------------------------------------------------
# Comparision plot - SURF vs ncc
#-------------------------------------------------------------------------------
f = plt.figure(figsize=(6,3.5))

ax1 = f.add_axes( [0.1,0.55,0.8,0.4])

ax1.plot(Axs_ncc, '-r', label='$A_x$ shift [ncc]', ms=4, lw=2 )
ax1.plot(Ays_ncc, '--b', label='$A_y$ shift [ncc]', ms=4, lw=2)

ax1.plot(Axs_surf, 'xr', label='$A_x$ shift [ORB]', ms=4, lw=1)
ax1.plot(Ays_surf, 'xb', label='$A_y$ shift [ORB]',ms=4, lw=1)

ax1.set_ylim(-7,12)
ax1.set_ylabel('shift [pixels]')
ax1.set_yticks((-5,0,5,10))
ax1.set_yticklabels( ('-5.0','0.0','5.0','10.0'))
ax1.yaxis.labelpad = 8
plt.legend(loc=8,ncol=2, columnspacing=0.5, numpoints=1,frameon=False)

#error
ax2 = f.add_axes( [0.1,0.1,0.8,0.4])
errx = Axs_surf-Axs_ncc
erry = Ays_surf-Ays_ncc

ax2.plot( errx, '-r', label='error $d_x = A_x - Ax_ncc$')
ax2.plot( erry, '--b', label='error $d_y = A_y - Ay_ncc$', lw=2)

ax2.set_ylabel('difference [pixels]')
ax2.set_xlabel('frame number')

#set font sizes
texts = f.findobj(plt.Text)
res = plt.setp(texts, fontsize=8)

#f.savefig('SURF vs NCC.pdf', dpi=300, transparent=True)