#!/usr/bin/env python
# coding: utf-8

# # Testing Data
# 
# blah ...

# In[1]:


get_ipython().run_line_magic('matplotlib', 'inline')
from matplotlib import pyplot as plt
import numpy as np


# In[16]:


from collections import namedtuple
import time
import pickle
import cv2

Data_ts = namedtuple('Data_ts', 'data timestamp')


# # Data Formats
# 
# Data was saved to a `dict` with each key being a `list` full of `namedtuples`. The `Data_ts` hold the data plus a time stamp since the start of the data.
# 
# - `imu`: accel, mag, gyros
# - `camera`: image bytes
# - `lidar`: array of (angle, range,) for each point

# In[3]:


get_ipython().run_line_magic('ls', '-alh')


# In[4]:


data = pickle.load( open( "data.pickle", "rb" ) )


# In[5]:


# data['imu'][0]


# In[6]:


# accel = [x[0][0] for x in data['imu']]
# mags = [x[0][1] for x in data['imu']]
# gyros = [x[0][2] for x in data['imu']]
# imutime = [x[1] for x in data['imu']]


# In[7]:


# plt.figure()
# plt.plot(accel);
# plt.grid(True);
# plt.title("Acceration [g's]")

# plt.figure()
# plt.plot(mags);
# plt.grid(True);
# plt.title("Magnetometer [uT's']")

# plt.figure()
# plt.plot(gyros);
# plt.grid(True);
# plt.title("Gyroscopes [deg/sec]");


# In[8]:


imgs = [np.frombuffer(x[0], dtype=np.uint8).reshape((480,640)) for x in data['camera']]
itime = [x[1] for x in data['camera']]


# In[9]:


plt.imshow(imgs[0], cmap='gray');


# In[10]:


plt.imshow(imgs[900], cmap='gray');


# In[13]:


def featureDetection():
    thresh = dict(threshold=25, nonmaxSuppression=True);
    fast = cv2.FastFeatureDetector_create(**thresh)
    return fast


# In[18]:


def getAbsoluteScale(f0, f1):
      x_pre, y_pre, z_pre = f0
      x    , y    , z     = f1
      scale = np.sqrt((x-x_pre)**2 + (y-y_pre)**2 + (z-z_pre)**2)
      return x, y, z, scale

def featureTracking(img_1, img_2, p1):

    lk_params = dict( winSize  = (21,21),
                      maxLevel = 3,
                      criteria = (cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 30, 0.01))

    p2, st, err = cv2.calcOpticalFlowPyrLK(img_1, img_2, p1, None, **lk_params)
    st = st.reshape(st.shape[0])
    ##find good one
    p1 = p1[st==1]
    p2 = p2[st==1]

    return p1,p2


# In[54]:


def run(images):
    #initialization
#     ground_truth =getTruePose()

    img_1 = images[0]
    img_2 = images[0]

    if len(img_1) == 3:
        gray_1 = cv2.cvtColor(img_1, cv2.COLOR_BGR2GRAY)
        gray_2 = cv2.cvtColor(img_2, cv2.COLOR_BGR2GRAY)
    else:
        gray_1 = img_1
        gray_2 = img_2

    #find the detector
    detector = featureDetection()
    kp1      = detector.detect(img_1)
    p1       = np.array([kp.pt for kp in kp1],dtype='float32')
    p1, p2   = featureTracking(gray_1, gray_2, p1)

    #Camera parameters
    fc = 718.8560
    pp = (640/2, 480/2)
#     K  = getK()

    E, mask = cv2.findEssentialMat(p2, p1, fc, pp, cv2.RANSAC,0.999,1.0);
    _, R, t, mask = cv2.recoverPose(E, p2, p1,focal=fc, pp = pp);

    #initialize some parameters
    MAX_FRAME     = 500
    MIN_NUM_FEAT  = 150

    preFeature = p2
    preImage   = gray_2

    R_f = R
    t_f = t

    maxError = 0
    ret_pos = []
    
    for numFrame in range(2, MAX_FRAME):

        if numFrame % 20 == 0:
            print(numFrame)

        if (len(preFeature) < MIN_NUM_FEAT):
            feature   = detector.detect(preImage)
            preFeature = np.array([ele.pt for ele in feature],dtype='float32')
            print(">> features found: ", len(preFeature))
            if len(preFeature) < MIN_NUM_FEAT:
                continue

        curImage_c = images[numFrame]

        if len(curImage_c) == 3:
              curImage = cv2.cvtColor(currImage_c, cv2.COLOR_BGR2GRAY)
        else:
              curImage = curImage_c

        kp1 = detector.detect(curImage);
        preFeature, curFeature = featureTracking(preImage, curImage, preFeature)
        E, mask = cv2.findEssentialMat(curFeature, preFeature, fc, pp, cv2.RANSAC,0.999,1.0);
        
#         print(E)
        
        _, R, t, mask = cv2.recoverPose(E, curFeature, preFeature, focal=fc, pp = pp);

#         truth_x, truth_y, truth_z, absolute_scale = getAbsoluteScale(
#             ground_truth[numFrame-1], ground_truth[numFrame])
        
#         if numFrame % 20 == 0:
#             print('scale', absolute_scale)

        absolute_scale = 1.0
            
        if absolute_scale > 0.1:
            t_f = t_f + absolute_scale*R_f.dot(t)
            R_f = R.dot(R_f)
        else:
            print("crap ... bad scale:", absolute_scale)

        preImage = curImage
        preFeature = curFeature

#         ret_pos.append((t_f[0], t_f[2],))
        ret_pos.append(t_f)

    return ret_pos


# In[55]:


pos = run(imgs)


# In[51]:


def getK():
    return   np.array([[7.188560000000e+02, 0, 6.071928000000e+02],
              [0, 7.188560000000e+02, 1.852157000000e+02],
              [0, 0, 1]])


# In[56]:


plt.plot(pos)


# In[57]:


pos


# In[46]:


aa=np.array([1,2,3])


# In[59]:


aa.transpose()


# In[48]:


aa[::2]


# In[ ]: