#!/usr/bin/env python
from __future__ import division, print_function
import os
from errno import EEXIST
import sys
import json
import time
import glob
import warnings
import threading
from shutil import copyfile
import exifread
import csv
from PIL import Image
from collections import defaultdict
from scipy.optimize import linear_sum_assignment
from datetime import datetime
import numpy as np
import cv2
import StringIO
from PIL import Image as PILImage
import matplotlib.pyplot as plt

# ROS imports
from rosbag import Bag
from tf.transformations import quaternion_multiply, quaternion_matrix, \
    quaternion_from_euler, euler_from_quaternion, quaternion_from_matrix, \
    quaternion_inverse, euler_matrix, quaternion_slerp, random_quaternion
from cv_bridge import CvBridge


bridge = CvBridge()

bag_fname = '/home/user/adapt_ws/data/2021-09-08-14-30-50.bag'
imagery_topic = '/image_raw'
imu_topic = '/an_device/Imu'
gps_topic = '/an_device/NavSatFix'

filter_status_topic = '/an_device/FilterStatus'
ins_system_status_topic = '/an_device/SystemStatus'
out_dir = os.path.splitext(bag_fname)[0]


image_dir = '%s/images' % out_dir
try:
    os.makedirs(image_dir)
except (IOError, OSError):
    pass


if False:
    # Print out the full list of topics.
    topics = set()
    for bag_fnames in bag_fname:
        with Bag(bag_fname, 'r') as ib:
            for topic, msg, t in ib:
                topics.add(topic)
    print('Full list of topics:', topics)


msgs = {}
images =  {}
system_status = {}
filter_status = {}
topics = set()
with Bag(bag_fname, 'r') as ib:
    for topic, msg, t in ib:
        topics.add(topic)

        if topic == ins_system_status_topic:
            system_status[t.to_sec()] = msg
            continue

        if topic == filter_status_topic:
            filter_status[t.to_sec()] = msg
            continue

        try:
            t = msg.header.stamp.to_sec()
        except AttributeError:
            continue

        if topic == imagery_topic:
            continue
            if hasattr(msg, 'format'):
                sio = StringIO.StringIO(msg.data)
                im = PILImage.open( sio )
                image = np.array( im )
            else:
                image = bridge.imgmsg_to_cv2(msg, "rgb8")

            frame_id = msg.header.frame_id

            images[t] = image[-100:, : ,:]

            fname = '%s/%i.jpg' % (image_dir, int(np.round(t*1000000)))

            cv2.imwrite(fname, image[:, :, ::-1])
        else:
            if topic not in msgs:
                msgs[topic] = {}

            msgs[topic][t] = msg


image_times = np.array(sorted(list(images.keys())))
t0 = image_times.min()
img1 = None
optical_flow = []
for t in image_times:
    img2 = images[t]
    if img1 is None:
        img1 = img2
        optical_flow.append(0)
        continue

    w = 99
    XY = np.meshgrid(np.arange(w//2, img1.shape[1] - w//2, w),
                     np.arange(w//2, img1.shape[0] - w//2, w))
    pts = np.vstack([XY[0].ravel(), XY[1].ravel()]).T.astype(np.float32)

    # calculate optical flow
    lk_params = dict(winSize  = (w,w),
                     maxLevel = 2,
                     criteria = (cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10, 0.03))
    pts2, st, err = cv2.calcOpticalFlowPyrLK(img1, img2, pts, None,
                                             **lk_params)

    delta = np.sqrt(np.mean((pts - pts2)**2))
    print('Calculating optical flow:', delta)
    optical_flow.append(delta)
    img1 = img2

optical_flow = np.array(optical_flow)
plt.plot(image_times[:len(optical_flow)] - t0, optical_flow/optical_flow.max(),
                     'ro', label='Optical Flow')


msgs_ = msgs[imu_topic]
state_times = np.array(sorted(list(msgs_.keys())))
angular_speeds = []
linear_acceleration = []
last_quat = None
orientation_change = []
orientation_cov = []
for t in state_times:
    v = msgs_[t].angular_velocity
    angular_speeds.append(np.linalg.norm([v.x, v.y, v.z]))
    a = msgs_[t].linear_acceleration
    linear_acceleration.append(np.linalg.norm([a.x, a.y, a.z]))

    orientation_cov.append(sum(msgs_[t].orientation_covariance))
    quat = np.array([msgs_[t].orientation.x, msgs_[t].orientation.y,
                     msgs_[t].orientation.z, msgs_[t].orientation.w])

    if last_quat is None:
        orientation_change.append(0)
        last_quat = quat
        continue

    dq = quaternion_multiply(last_quat, quaternion_inverse(quat))
    orientation_change.append(np.arccos(min([dq[3], 1]))*2)
    last_quat = quat


angular_speeds = np.array(angular_speeds)
linear_acceleration = np.array(linear_acceleration)
orientation_cov = np.array(orientation_cov)
orientation_change = np.array(orientation_change)

plt.rc('axes', linewidth=4)
plt.rcParams.update({'font.size': 22})
plt.figure(num=None, figsize=(20, 16), dpi=80)
plt.plot(image_times[:len(optical_flow)] - t0, optical_flow/optical_flow.max(),
                     'r-', label='Optical Flow')
plt.plot(state_times - t0, angular_speeds/angular_speeds.max(), 'b-',
         label='Angular Speed')
plt.plot(state_times - t0, linear_acceleration/linear_acceleration.max(), 'g-',
         label='Linear Accel')
plt.plot(state_times - t0, orientation_change/orientation_change.max(), 'c-',
         label='Orientation Change')
plt.plot(state_times - t0, orientation_cov/orientation_cov.max(), 'm-',
         label='Orientation Cov')
plt.xlabel('Seconds', fontsize=16)
plt.legend(fontsize=16)


plt.rc('axes', linewidth=4)
plt.rcParams.update({'font.size': 22})
plt.figure(num=None, figsize=(20, 16), dpi=80)
plt.plot(image_times[:len(optical_flow)] - t0, optical_flow/optical_flow.max(),
                     'ro', label='Optical Flow')
angular_acell = np.abs(np.diff(angular_speeds))
plt.plot(state_times[1:] - t0, angular_acell/angular_acell.max(), 'bo',
         label='Angular Acceleration')
plt.xlabel('Seconds', fontsize=40)
plt.ylabel('Magnitude of Change', fontsize=40)
plt.legend(fontsize=40)
