"""
Test of feature descriptor and matching calculation times
"""

#%%-----------------------------------------------------------------------------
# imports and load data
#-------------------------------------------------------------------------------
import numpy as np
import time
import cv2
import random 

#load test speckle patterns
imgs = np.load("random positions.npy")


#Find features using the fast detector
threshold = 255
nonmaxSuppression = True
fast = cv2.FastFeatureDetector_create( threshold, nonmaxSuppression, cv2.FastFeatureDetector_TYPE_9_16)

N = 500 #target number of features
kps_list = []

for img in imgs:
    
    #find approx threshold for first image
    kps = fast.detect(img)
    while len(kps) < N:
        threshold = threshold - 1
        fast.setThreshold(threshold)
        kps = fast.detect(img)
    print 'using threshold: ',threshold, 'Number of features: ', len(kps)
    print 'threshold/mean', threshold/img.mean()
    
    kps = kps[0:500]
    kps_list.append(kps)
    
    #reset threshold to ensure N features are found next frame
    threshold = 255
    fast.setThreshold(threshold)


## create BFMatcher objects
l1_matcher = cv2.BFMatcher(cv2.NORM_L1, crossCheck=True)
l2_matcher = cv2.BFMatcher(cv2.NORM_L2, crossCheck=True)
hamming_matcher = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)

#%%-----------------------------------------------------------------------------
# U-SURF - Difference of Gaussians detector
#-------------------------------------------------------------------------------
threshold =  2000 #KK used varibale threshold 300 - 50000
nOctaves = 1        #KK =1 no image rescaling 4 by default
nOctaveLayers = 1 #KK used three default is 1 the minimum?
extended = False
upright = True
usurf = cv2.xfeatures2d.SURF_create(threshold, nOctaves, nOctaveLayers ,extended, upright)


#%%-----------------------------------------------------------------------------
# U-SURF descriptor speed test
#-------------------------------------------------------------------------------

#Loop over all the images/kps and calculate the dsecriptor.
tic = time.time()
for n in range(0, len(kps_list)):
    img = imgs[n]
    kps = kps_list[n]
    #compute descriptors
    kps, des = usurf.compute(img, kps)
toc =time.time()
t = (toc-tic)/(n+1) * 1000 
print 'U-SURF: ', t, 'ms per frame'   

#%%-----------------------------------------------------------------------------
# U-SURF Matching speed test
#-------------------------------------------------------------------------------
des0_list = []  #list of descriptors
des1_list = []  #shuffled list of descriptors

for n in range(0, len(kps_list)):
    img = imgs[n]
    kps = kps_list[n]
    kps, des = usurf.compute(img, kps)
    
    des0_list.append(des)
    
    des1 = des.copy()
    random.shuffle( des1)
    des1_list.append( des1 )


des0_list = des0_list * 10
des1_list = des1_list * 10

#L1 norm
tic = time.time()
for des0, des1 in zip(des0_list, des1_list):
    matches = l1_matcher.match(des0,des1)
toc = time.time()
t_usurf_l1 = (toc-tic)/len(des0_list)
print '     L1_norm', t_usurf_l1*1000, 'ms'
 
#L2 norm
tic = time.time()
for des0, des1 in zip(des0_list, des1_list):
    matches = l1_matcher.match(des0,des1)
toc = time.time()
t_usurf_l2 = (toc-tic)/len(des0_list)
print '     L2_norm', t_usurf_l2*1000, 'ms'

#%%-----------------------------------------------------------------------------
# SURF - Difference of Gaussians detector
#-------------------------------------------------------------------------------
threshold =  2000 #KK used varibale threshold 300 - 50000
nOctaves = 1        #KK =1 no image rescaling 4 by default
nOctaveLayers = 1 #KK used three default is 1 the minimum?
extended = False
upright = False
usurf = cv2.xfeatures2d.SURF_create(threshold, nOctaves, nOctaveLayers ,extended, upright)

#%%-----------------------------------------------------------------------------
# SURF descriptor speed test
#-------------------------------------------------------------------------------

#Loop over all the images/kps and calculate the dsecriptor.
tic = time.time()
for n in range(0, len(kps_list)):
    img = imgs[n]
    kps = kps_list[n]
    #compute descriptors
    kps, des = usurf.compute(img, kps)
toc =time.time()
t = (toc-tic)/(n+1) * 1000 
print 'SURF: ', t, 'ms per frame'   

#%%-----------------------------------------------------------------------------
# SURF - Matching speed test
#-------------------------------------------------------------------------------
des0_list = []  #list of descriptors
des1_list = []  #shuffled list of descriptors

for n in range(0, len(kps_list)):
    img = imgs[n]
    kps = kps_list[n]
    kps, des = surf.compute(img, kps)
    
    des0_list.append(des)
    
    des1 = des.copy()
    random.shuffle( des1)
    des1_list.append( des1 )

#make a bit longer for timeing test
des0_list = des0_list * 10
des1_list = des1_list * 10

#L1 norm
tic = time.time()
for des0, des1 in zip(des0_list, des1_list):
    matches = l1_matcher.match(des0,des1)
toc = time.time()
t_surf_l1 = (toc-tic)/len(des0_list)
print '     L1_norm', t_surf_l1*1000, 'ms'
 
#L2 norm
tic = time.time()
for des0, des1 in zip(des0_list, des1_list):
    matches = l1_matcher.match(des0,des1)
toc = time.time()
t_surf_l2 = (toc-tic)/len(des0_list)
print '     L2_norm', t_surf_l2*1000, 'ms'

#%%-----------------------------------------------------------------------------
# ORB / rotated brief
#-------------------------------------------------------------------------------
nFeatures= 500
scaleFactor = 1.2   #pyramid scale factor
nLevels = 1       #pyramid levels 1 as no scale changes
edgeThreshold = 0 #size of border to not detect in
firstLevel = 0  #zero from opencv docs
#for descriptors
WTA_K = 2       #he number of points that produce each element of the oriented BRIEF descriptor. The default value 2 means the BRIEF where we take a random point pair and compare their brightnesses, so we get 0/1 response (index of brightest)
scoreType = cv2.ORB_HARRIS_SCORE #or ORB_FAST_SCORE
patchSize = 31  #size of patch used to find BRIEF dscriptor
fastThreshold = 50

orb = cv2.ORB_create(nFeatures, scaleFactor, nLevels, edgeThreshold, firstLevel, WTA_K, scoreType, patchSize, fastThreshold)
#work around for orb bug - this needs to be called first
cv2.ocl.setUseOpenCL(False)

#%%-----------------------------------------------------------------------------
# ORB descriptor speed test
#-------------------------------------------------------------------------------

#Loop over all the images/kps and calculate the dsecriptor.
tic = time.time()
for n in range(0, len(kps_list)):
    img = imgs[n]
    kps = kps_list[n]
    #compute descriptors
    kps, des = orb.compute(img, kps)
toc =time.time()
t = (toc-tic)/(n+1) * 1000 
print 'ORB: ', t, 'ms per frame' 

#%%-----------------------------------------------------------------------------
# ORB Matching speed test
#-------------------------------------------------------------------------------
des0_list = []  #list of descriptors
des1_list = []  #shuffled list of descriptors

for n in range(0, len(kps_list)):
    img = imgs[n]
    kps = kps_list[n]
    kps, des = orb.compute(img, kps)
    
    des0_list.append(des)
    
    des1 = des.copy()
    random.shuffle( des1)
    des1_list.append( des1 )


des0_list = des0_list * 10
des1_list = des1_list * 10

#l1 and l2 do not make sense with binary descriptors
 
#Hamming
tic = time.time()
for des0, des1 in zip(des0_list, des1_list):
    matches = hamming_matcher.match(des0,des1)
toc = time.time()
t_orb_hamming = (toc-tic)/len(des0_list)
print '     L2_hamming', t_orb_hamming*1000, 'ms'


#%%-----------------------------------------------------------------------------
# BRISK
#-------------------------------------------------------------------------------
octaves = 0         #detection octaves. Use 0 to do single scale.
patternScale = 1.8  #apply this scale to the pattern used for sampling the neighbourhood of a keypoint.
                    #in the paper the pattern is approximately 1 = +-18px, 1.8=+-32
threshold = 91     #FAST threshold
brisk = cv2.BRISK_create(threshold, octaves, patternScale)


#%%-----------------------------------------------------------------------------
# BRISK descriptor speed test
#-------------------------------------------------------------------------------

#Loop over all the images/kps and calculate the dsecriptor.
tic = time.time()
for n in range(0, len(kps_list)):
    img = imgs[n]
    kps = kps_list[n]
    #compute descriptors
    kps, des = brisk.compute(img, kps)
toc =time.time()
t = (toc-tic)/(n+1) * 1000 
print 'BRISK: ', t, 'ms per frame' 


#%%-----------------------------------------------------------------------------
# BRISK Matching speed test
#-------------------------------------------------------------------------------
des0_list = []  #list of descriptors
des1_list = []  #shuffled list of descriptors

for n in range(0, len(kps_list)):
    img = imgs[n]
    kps = kps_list[n]
    kps, des = brisk.compute(img, kps)
    
    des0_list.append(des)
    
    des1 = des.copy()
    random.shuffle( des1)
    des1_list.append( des1 )


des0_list = des0_list * 10
des1_list = des1_list * 10

#l1 and l2 do not make sense with binary descriptors
 
#Hamming
tic = time.time()
for des0, des1 in zip(des0_list, des1_list):
    matches = hamming_matcher.match(des0,des1)
toc = time.time()
t_orb_hamming = (toc-tic)/len(des0_list)
print '     L2_hamming', t_orb_hamming*1000, 'ms'