"""
Test script to process experimental rotation images using the feature
tracking methods.
"""
#%%-----------------------------------------------------------------------------
# setup
#-------------------------------------------------------------------------------
import numpy as np
import matplotlib.pyplot as plt
plt.ion()

import cv2 as cv
import lib_feature_tracking as lib
import time

##Load data
imgs = np.load('z rotation.npy')

##Setup range of rotations
angles = np.arange(0,360.5,0.5)
angles[ angles>180 ] = angles[ angles>180]-360 #flip angles to negative

##Other setup 
roi = (0,0,512,512)     #x0,y0, w,h, window to use

#%%-----------------------------------------------------------------------------
# Setup matcher methods
#-------------------------------------------------------------------------------
ncc = lib.NCC()
orb = lib.ORB(patchSize=64, fastThreshold=36, distance_threshold=64 )
brisk = lib.BRISK(fastThreshold=36, patternScale=1.8, distance_threshold=50 )
usurf = lib.USURF(threshold=50, distance_threshold=0.5, norm=cv.NORM_L1 )
surf = lib.SURF(threshold=50, 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()

#fixed reference image
img0  = imgs[0, roi[1]:roi[1]+roi[3],roi[0]:roi[0]+roi[2]]
Nmax = 721

##Processing loop
tic=time.time()
for n in range(0, Nmax):
    ang = angles[n]
    img1 = imgs[n, roi[1]:roi[1]+roi[3],roi[0]:roi[0]+roi[2]]
    #print 'doing angle', ang

    ##Do ORB
    Ax,Ay,theta = orb.calc_shift_and_rotation(img0, img1, adj_fast_threshold=True)
    #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)
    #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)
    #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)
    #print '\tsurf: {0:0.2f}{1:0.2f}{1:0.2f}'.format(Ax,Ay,theta)
toc=time.time()
print 'Time: ',(toc-tic)/(n+1)

#make results arrays
Axs_orb = np.array(orb.Axs)
Ays_orb = np.array(orb.Ays)
thetas_orb = np.array(orb.thetas) 
Nfms_orb = np.array(orb.Nfms)

Axs_brisk = np.array(brisk.Axs)
Ays_brisk = np.array(brisk.Ays)
thetas_brisk = np.array(brisk.thetas)  
Nfms_brisk = np.array(brisk.Nfms)  

Axs_usurf = np.array(usurf.Axs)
Ays_usurf = np.array(usurf.Ays)
thetas_usurf = np.array(usurf.thetas)  
Nfms_usurf = np.array(usurf.Nfms)  

Axs_surf = np.array(surf.Axs)
Ays_surf = np.array(surf.Ays)
thetas_surf = np.array(surf.thetas) 
Nfms_surf = np.array(surf.Nfms)


#%%-----------------------------------------------------------------------------
# theta plot
#-------------------------------------------------------------------------------
f = plt.figure(figsize=(6,6))

ax1 = f.add_axes( [0.12,0.70,0.8,0.25])
ax1.plot(angles[0:Nmax], thetas_orb, '.b', label='ORB', mew=0, ms=6)
ax1.plot(angles[0:Nmax], thetas_brisk, 'xg', label='BRISK',mew=1)
ax1.plot(angles[0:Nmax], thetas_surf, '^c', label='SURF', mew=0,ms=8)
ax1.plot(angles[0:Nmax], thetas_usurf, '*r', label='USURF', mew=0,ms=8)
ax1.plot(angles[0:Nmax], angles[0:Nmax], ':k', label='Applied rotation')
ax1.set_ylabel('rotation [degrees]')
ax1.set_xlim(-50,50)
ax1.set_ylim(-50,50)
ax1.yaxis.labelpad = 7
ax1.hlines(0,-180,180,'k', '--')
plt.legend(ncol=2, loc=(0.02,0.55), numpoints = 1, scatterpoints=1, 
            labelspacing=0.00,columnspacing=1.20, frameon=False)

#error
ax2 = f.add_axes( [0.12,0.40,0.8,0.25])
ax2.plot(angles[0:Nmax], thetas_orb-angles[0:Nmax], '.b', label='ORB', mew=0, ms=6)
ax2.plot(angles[0:Nmax], thetas_brisk-angles[0:Nmax], 'xg', label='BRISK', mew=1)
ax2.plot(angles[0:Nmax], thetas_surf-angles[0:Nmax], '^c', label='SURF', mew=0, ms=8)
ax2.plot(angles[0:Nmax], thetas_usurf-angles[0:Nmax], '*r', label='USURF', mew=0, ms=8)

ax2.set_ylim(-0.25,0.25)
ax2.set_xlim(-50,50)
ax2.set_ylabel('Angle error [degrees]')
#ax2.set_xlabel('Rotation angle [degrees]')
ax2.hlines(0,-180,180, 'k', '--')
ax2.yaxis.labelpad = +5


#Number of features
ax3 = f.add_axes( [0.12,0.10,0.8,0.25])
ax3.plot(angles[0:Nmax], Nfms_orb, '.b', label='ORB', mew=0, ms=6)
ax3.plot(angles[0:Nmax], Nfms_brisk, 'xg', label='BRISK', mew=1)
ax3.plot(angles[0:Nmax], Nfms_usurf, '^c', label='USURF', mew=0, ms=8)
ax3.plot(angles[0:Nmax], Nfms_surf, '*r', label='SURF', mew=0, ms=8)
ax3.set_ylim(0,400)
ax3.set_yticks( (0,100,200,300,400))
ax3.set_xlim(-50,50)
ax3.set_ylabel('Number of feature matches')
ax3.set_xlabel('Rotation angle [degrees]')
ax3.yaxis.labelpad = +10


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

f.savefig('..\\figures\\rotation experimental.pdf', dpi=600, transparent=True)

#%%-----------------------------------------------------------------------------
# shifts of top-left of image.
#-------------------------------------------------------------------------------
f = plt.figure( figsize=(6,3))
ax1 = f.add_axes( (0.13,0.13,0.8,0.8))
ax1.plot( angles[0:Nmax], Axs_orb, '.-', label='ORB')
ax1.plot( angles[0:Nmax], Axs_brisk, '.-', label='BRISK')
ax1.plot( angles[0:Nmax], Axs_usurf, '.-', label='USURF')
ax1.plot( angles[0:Nmax], Axs_surf, '.-', label='SURF')