#!/usr/bin/env python
from __future__ import division, print_function
import os
import sys
import numpy as np
from bisect import bisect_left, bisect_right


# ROS imports
try:
    from rosbag import Bag
    HAS_ROS = True
except ImportError:
    HAS_ROS = False

try:
    import matplotlib.pyplot as plt
    HAS_MATPLOTLIB = True
except ImportError:
    HAS_MATPLOTLIB = False

from colmap_processing.platform_pose import PlatformPoseProvider
from colmap_processing.geo_conversions import enu_to_llh, llh_to_enu, \
    ned_quat_to_enu_quat, quat_std_to_ypr_std


def get_nearest_element(sorted_array, val):
    # bisect_left will return the index for the exact match if one exists.
    # Otherwise, it will return the index for the smallest val that is greater
    # than 'val' or the length of sorted_array if none are greater than 'val'.

    # ind points to a time that is >= time.
    ind = bisect_left(sorted_array, val)

    if ind == len(sorted_array):
        return sorted_array[-1]
    elif ind == 0:
        return sorted_array[0]
    elif abs(sorted_array[ind - 1] - val) < abs(sorted_array[ind] - val):
        return sorted_array[ind - 1]
    else:
        return sorted_array[ind]


class ANINS(PlatformPoseProvider):
    """Model for Advanced Navigation intertial navigation system.
    """
    def __init__(self, lat0=None, lon0=None, h0=None, max_length=np.inf):
        """

        """
        self.lat0 = lat0
        self.lon0 = lon0
        self.h0 = h0
        self.max_length = max_length

    @classmethod
    def from_ros_bag(cls, bag_fname, lat0=None, lon0=None, h0=None,
                     filter_status_topic='/an_device/FilterStatus',
                     imu_topic='/an_device/Imu',
                     navsatfix_topic='/an_device/NavSatFix',
                     odometry_topic='/an_device/Odometry',
                     raw_gnss_topic='/an_device/RawGNSS',
                     raw_imu_topic='/an_device/RawImu',
                     sys_status_topic='/an_device/SystemStatus',
                     twist_topic='/an_device/Twist',
                     ptp_health_ins_topic='/ptp_health/ins',
                     ptp_health_nmea_topic='/ptp_health/nmea',
                     ptp_health_pps_topic='/ptp_health/pps',
                     ptp_health_ptp_topic='/ptp_health/ptp'):
        """Create instance from dictionary of ROS messages index by topic.

        :param bag_fname: File path to ROS bag.
        :type msgs: str
        """
        self = cls()

        filter_status = {}
        self.imu_msgs = {}
        self.nav_sat_fix_msgs = {}
        self.odometry_msgs = {}
        self.raw_gnss_msgs = {}
        self.raw_imu_msgs = {}
        self.sys_status_msgs = {}
        self.twist_msgs = {}
        self.ptp_health_in_msgs = {}
        self.ptp_health_nmea_msgs = {}
        self.ptp_health_pps_msgs = {}
        self.ptp_health_ptp_msgs = {}
        topics = set()
        with Bag(bag_fname, 'r') as ib:
            for topic, msg, t in ib:
                try:
                    # For many messages, the header time corresponds to an
                    # accurate time traceable to the source.
                    t = msg.header.stamp.to_sec()
                except AttributeError:
                    # Use the time it was added to the bag.
                    t = t.to_sec()

                topics.add(topic)
                if topic == filter_status_topic:
                    m = msg.message
                    orientation_fiter_init = '0. Orientation Filter Initialised.' in m
                    navigation_filter_init = '1. Navigation Filter Initialised.' in m
                    heading_initialized = '2. Heading Initialised.' in m
                    utc_time_initialized = '3. UTC Time Initialised.' in m
                    internal_gnss_enabled = '9. Internal GNSS Enabled.' in m
                    magnetic_heading = '10. Magnetic Heading Active.' in m
                    velocity_heading = '11. Velocity Heading Enabled.' in m
                    atmospheric_altitude = '12. Atmospheric Altitude Enabled.' in m
                    filter_status[t] = [orientation_fiter_init,
                                        navigation_filter_init,
                                        heading_initialized,
                                        utc_time_initialized,
                                        internal_gnss_enabled,
                                        magnetic_heading,
                                        velocity_heading,
                                        atmospheric_altitude]
                elif topic == imu_topic:
                    self.imu_msgs[t] = msg
                elif topic == navsatfix_topic:
                    self.nav_sat_fix_msgs[t] = msg
                elif topic == odometry_topic:
                    self.odometry_msgs[t] = msg
                elif topic == raw_gnss_topic:
                    self.raw_gnss_msgs[t] = msg
                elif topic == raw_imu_topic:
                    self.raw_imu_msgs[t] = msg
                elif topic == sys_status_topic:
                    self.sys_status_msgs[t] = msg
                elif topic == twist_topic:
                    self.twist_msgs[t] = msg
                elif topic == ptp_health_ins_topic:
                    self.ptp_health_in_msgs[t] = msg
                elif topic == ptp_health_nmea_topic:
                    self.ptp_health_nmea_msgs[t] = msg
                elif topic == ptp_health_pps_topic:
                    self.ptp_health_pps_msgs[t] = msg
                elif topic == ptp_health_ptp_topic:
                    self.ptp_health_ptp_msgs[t] = msg

        print('Found topics:', topics)

        self.filter_status_times = np.array(sorted(list(filter_status.keys())))
        self.filter_status = np.array([filter_status[t]
                                       for t in self.filter_status_times],
                                       dtype=bool).T

        self.odometry_times = sorted(list(self.odometry_msgs.keys()))
        self.raw_imu_times = sorted(list(self.raw_imu_msgs.keys()))

        if lat0 is None:
            lats = [self.odometry_msgs[t].pose.pose.position.y
                    for t in self.odometry_msgs]
            lat0 = np.median(lats)

        if lon0 is None:
            lons = [self.odometry_msgs[t].pose.pose.position.x
                    for t in self.odometry_msgs]
            lon0 = np.median(lons)

        if h0 is None:
            hs = [self.odometry_msgs[t].pose.pose.position.z
                    for t in self.odometry_msgs]
            h0 = np.min(hs)

        self.lat0 = lat0
        self.lon0 = lon0
        self.h0 = h0

        return self

    def plot_ptp_status(self):
        def explode(msgs):
            msgs_out = {}
            for t in msgs:
                if msgs[t].message not in msgs_out:
                    msgs_out[msgs[t].message] = []

                msgs_out[msgs[t].message].append(t)

            return msgs_out

        msgs_out = explode(self.ptp_health_in_msgs)
        msgs_out.update(explode(self.ptp_health_nmea_msgs))
        msgs_out.update(explode(self.ptp_health_pps_msgs))
        msgs_out.update(explode(self.ptp_health_ptp_msgs))

        labels = sorted(list(msgs_out.keys()))

        plt.figure(num=None, figsize=(15.3, 10.7), dpi=80)
        plt.rc('font', **{'size': 20})
        plt.rc('axes', linewidth=4)
        for i in range(len(msgs_out)):
            ts = msgs_out[labels[i]]
            plt.plot(ts, np.ones(len(ts))*i, 'o', label=labels[i])

        plt.xlabel('Time (s)', fontsize=40)
        plt.yticks(range(len(labels)), labels)
        #plt.legend(loc=0, fontsize=30)
        plt.tight_layout()

    @property
    def times(self):
        return np.array(self.odometry_times)

    def plot_filter_status(self):
        labels = ['Orientation Filter Initialised',
                  'Navigation Filter Initialised',
                  'Heading Initialised',
                  'UTC Time Initialised',
                  'Internal GNSS Enabled',
                  'Magnetic Heading Active',
                  'Velocity Heading Enabled',
                  'Atmospheric Altitude Enabled']

        plt.figure(num=None, figsize=(15.3, 10.7), dpi=80)
        plt.rc('font', **{'size': 20})
        plt.rc('axes', linewidth=4)
        for i in range(len(self.filter_status)):
            ts = self.filter_status_times[self.filter_status[i]]
            plt.plot(ts, np.ones(len(ts))*i, 'o', label=labels[i])

        plt.xlabel('Time (s)', fontsize=40)
        plt.yticks(range(len(self.filter_status)), labels)
        #plt.legend(loc=0, fontsize=30)
        plt.tight_layout()

    def get_odometry_msg(self, t, time_thresh=0.1):
        """
        :param t: Time at which to query the INS state (time in seconds since
            Unix epoch).
        :type t: float

        :param time_thresh: Time (seconds) threshold such that if the nearest
            state is further in time to the provide time, then return None.
        :type time_thresh: float

        :return: ROS Odometry message.
        :rtype: nav_msgs.Odometry

        """
        nearest_t = get_nearest_element(self.odometry_times, t)

        if abs(nearest_t - t) > time_thresh:
            return None

        msg = self.odometry_msgs[nearest_t]
        return msg

    def get_raw_imu_msg(self, t, time_thresh=0.1):
        """
        :param t: Time at which to query the INS state (time in seconds since
            Unix epoch).
        :type t: float

        :param time_thresh: Time (seconds) threshold such that if the nearest
            state is further in time to the provide time, then return None.
        :type time_thresh: float

        :return: ROS Odometry message.
        :rtype: nav_msgs.Odometry

        """
        nearest_t = get_nearest_element(self.raw_imu_times, t)

        if abs(nearest_t - t) > time_thresh:
            return None

        msg = self.raw_imu_msgs[nearest_t]
        return msg

    def get_raw_imu(self):
        data = []
        for t in self.raw_imu_msgs:
            msg = self.raw_imu_msgs[t]
            data.append([t, msg.accelerometer.x, msg.accelerometer.y,
                         msg.accelerometer.z, msg.gyroscope.x, msg.gyroscope.y,
                         msg.gyroscope.z, msg.magnetometer.x,
                         msg.magnetometer.y, msg.magnetometer.z, msg.imuTempC,
                         msg.pressurePascals, msg.pressureTempC])

        data = np.array(data)
        ind = np.argsort(data[:, 0])
        data = data[ind].T
        return data

    def get_raw_imu_at_time(self, t, time_thresh=0.1):
        """
        :param t: Time at which to query the INS state (time in seconds since
            Unix epoch).
        :type t: float

        :param time_thresh: Time (seconds) threshold such that if the nearest
            state is further in time to the provide time, then return None.
        :type time_thresh: float

        :return: List with the first element the 3-array acceleration (m/s^2),
            the second element 3-array angular rotation vector (rad/s), the
            third element magnetometer field direction, the fourth element IMU
            temperature (C), fifth element atomspheric pressure (Pa), and sixth
            element the ambient temperature (C).
        :rtype: None | list

        """
        msg = self.get_raw_imu_msg(t, time_thresh)

        if msg is None:
            return

        return [[msg.accelerometer.x, msg.accelerometer.y, msg.accelerometer.z],
                [msg.gyroscope.x, msg.gyroscope.y, msg.gyroscope.z],
                [msg.magnetometer.x, msg.magnetometer.y, msg.magnetometer.z],
                 msg.imuTempC, msg.pressurePascals, msg.pressureTempC]

    def get_imu_msg(self, t, time_thresh=0.1):
        """
        :param t: Time at which to query the INS state (time in seconds since
            Unix epoch).
        :type t: float

        :param time_thresh: Time (seconds) threshold such that if the nearest
            state is further in time to the provide time, then return None.
        :type time_thresh: float

        :return: ROS sensor_msgs/Imu.
        :rtype: nav_msgs.Odometry

        """
        nearest_t = get_nearest_element(self.raw_imu_times, t)

        if abs(nearest_t - t) > time_thresh:
            return None

        msg = self.raw_imu_msgs[nearest_t]
        return msg

    def get_imu(self):
        data = []
        for t in self.imu_msgs:
            msg = self.imu_msgs[t]
            data.append([t, msg.linear_acceleration.x,
                         msg.linear_acceleration.y,
                         msg.linear_acceleration.z,
                         msg.angular_velocity.x,
                         msg.angular_velocity.y,
                         msg.angular_velocity.z,])

        data = np.array(data)
        ind = np.argsort(data[:, 0])
        data = data[ind].T
        return data

    def get_imu_at_time(self, t, time_thresh=0.1):
        """
        :param t: Time at which to query the INS state (time in seconds since
            Unix epoch).
        :type t: float

        :param time_thresh: Time (seconds) threshold such that if the nearest
            state is further in time to the provide time, then return None.
        :type time_thresh: float

        :return: List with the first element the 3-array acceleration (m/s^2),
            the second element 3-array angular rotation vector (rad/s), the
            third element magnetometer field direction, the fourth element IMU
            temperature (C), fifth element atomspheric pressure (Pa), and sixth
            element the ambient temperature (C).
        :rtype: None | list

        """
        msg = self.get_raw_imu_msg(t, time_thresh)

        if msg is None:
            return

        return [[msg.accelerometer.x, msg.accelerometer.y, msg.accelerometer.z],
                [msg.gyroscope.x, msg.gyroscope.y, msg.gyroscope.z],
                [msg.magnetometer.x, msg.magnetometer.y, msg.magnetometer.z],
                 msg.imuTempC, msg.pressurePascals, msg.pressureTempC]

    def pose(self, t, time_thresh=0.1, in_enu=True, include_cov=False):
        """
        :param t: Time at which to query the INS state (time in seconds since
            Unix epoch).
        :type t: float

        :param time_thresh: Time (seconds) threshold such that if the nearest
            state is further in time to the provide time, then return None.
        :type time_thresh: float

        :param in_enu: If True, position will be returned in meters east/north/
            up relative to the origin at self.lat0, self.lon0, self.h0.
            Otherwise, position will be in latitude (degrees), longitude
            (degrees), height above WGS84 ellipsoid (meters).
        :type in_enu: bool

        :return: List with the first element the 3-array position (see pos) and
            the second element the 4-array orientation quaternion (see quat).
            If covariance is included, it will all return a 6x6 covariance
            matrix only populated along the diagonal. Elements along the
            diagonal encode the variance in easting (m^2), northing (m^2), up
            (m^2), roll (rad^2), pitch (rad^2), yaw (rad^2).
        :rtype: None | list

        """
        msg = self.get_odometry_msg(t, time_thresh)

        if msg is None:
            return

        pos = [msg.pose.pose.position.y, msg.pose.pose.position.x,
               msg.pose.pose.position.z]

        if in_enu:
            pos = llh_to_enu(pos[0], pos[1], pos[2], self.lat0, self.lon0,
                             self.h0)

        pos = np.array(pos)

        ned_quat = [msg.pose.pose.orientation.x, msg.pose.pose.orientation.y,
                    msg.pose.pose.orientation.z, msg.pose.pose.orientation.w]
        ned_quat = np.array(ned_quat)
        quat = ned_quat_to_enu_quat(ned_quat)

        if include_cov:
            cov = np.zeros((6, 6))

            # Variance in easting (m^2).
            cov[0, 0] = msg.pose.covariance[0]

            # Variance in northing (m^2).
            cov[1, 1] = msg.pose.covariance[7]

            # Variance in up (m^2).
            cov[2, 2] = msg.pose.covariance[14]

            # Calculate heading, pitch, roll uncertainty.
            # Packet 27, Fields 2, 3, 4, Anpp Packet 27 Field 1 is for the
            # W axis, which is not requested.
            qx_var = msg.pose.covariance[21]
            qy_var = msg.pose.covariance[28]
            qz_var = msg.pose.covariance[35]

            # It is fine to pass in varience instead of standard deviation
            # (i.e., sqrt(variance)) and then interpret the output as variance.
            heading_var, pitch_var, roll_var = quat_std_to_ypr_std(ned_quat,
                                                                   qx_var,
                                                                   qy_var,
                                                                   qz_var)

            cov[3, 3] = heading_var
            cov[4, 4] = pitch_var
            cov[5, 5] = roll_var

            return pos, quat, cov
        else:
            return pos, quat

    def plot_odometry(self, show_in_enu=False):
        data = []
        for t in self.odometry_msgs:
            msg = self.odometry_msgs[t]
            data.append([t, msg.pose.pose.position.y, msg.pose.pose.position.x,
                         msg.pose.pose.position.z, msg.pose.pose.orientation.x,
                         msg.pose.pose.orientation.y,
                         msg.pose.pose.orientation.z,
                         msg.pose.pose.orientation.w])

        data = np.array(data)
        ind = np.argsort(data[:, 0])
        data = data[ind]

        plt.figure(num=None, figsize=(15.3, 10.7), dpi=80)
        plt.rc('font', **{'size': 20})
        plt.rc('axes', linewidth=4)
        plt.plot(data[:, 1], data[:, 2])
        plt.xlabel('Longitude (deg)', fontsize=40)
        plt.ylabel('Latitude (deg)', fontsize=40)
        plt.tight_layout()

        plt.figure(num=None, figsize=(15.3, 10.7), dpi=80)
        plt.rc('font', **{'size': 20})
        plt.rc('axes', linewidth=4)
        plt.plot(data[:, 1], data[:, 3])
        plt.xlabel('Longitude (deg)', fontsize=40)
        plt.ylabel('Height Above WGS84 (m)', fontsize=40)
        plt.tight_layout()

        plt.figure(num=None, figsize=(15.3, 10.7), dpi=80)
        plt.rc('font', **{'size': 20})
        plt.rc('axes', linewidth=4)
        plt.plot(data[:, 2], data[:, 3])
        plt.xlabel('Latitude (deg)', fontsize=40)
        plt.ylabel('Height Above WGS84 (m)', fontsize=40)
        plt.tight_layout()

        plt.figure(num=None, figsize=(15.3, 10.7), dpi=80)
        plt.rc('font', **{'size': 20})
        plt.rc('axes', linewidth=4)
        plt.plot(data[:, 0], data[:, 3])
        plt.xlabel('Time (s)', fontsize=40)
        plt.ylabel('Height Above WGS84 (m)', fontsize=40)
        plt.tight_layout()

        if show_in_enu:
            enu = [llh_to_enu(d[1], d[2], d[3], self.lat0, self.lon0, self.h0)
                   for d in data]
            enu = np.array(enu).T

            plt.figure(num=None, figsize=(15.3, 10.7), dpi=80)
            plt.rc('font', **{'size': 20})
            plt.rc('axes', linewidth=4)
            plt.plot(enu[0], enu[1])
            plt.xlabel('Easting (m)', fontsize=40)
            plt.ylabel('Northing (m)', fontsize=40)
            plt.tight_layout()

            plt.figure(num=None, figsize=(15.3, 10.7), dpi=80)
            plt.rc('font', **{'size': 20})
            plt.rc('axes', linewidth=4)
            plt.plot(enu[0], enu[2])
            plt.xlabel('Easting (m)', fontsize=40)
            plt.ylabel('Height (m)', fontsize=40)
            plt.tight_layout()

            plt.figure(num=None, figsize=(15.3, 10.7), dpi=80)
            plt.rc('font', **{'size': 20})
            plt.rc('axes', linewidth=4)
            plt.plot(enu[1], enu[2])
            plt.xlabel('Northing (m)', fontsize=40)
            plt.ylabel('Height (m)', fontsize=40)
            plt.tight_layout()

    def plot_raw_imu(self):
        data = []
        for t in self.raw_imu_msgs:
            msg = self.raw_imu_msgs[t]
            data.append([t, msg.accelerometer.x, msg.accelerometer.y,
                         msg.accelerometer.z, msg.gyroscope.x, msg.gyroscope.y,
                         msg.gyroscope.z, msg.magnetometer.x,
                         msg.magnetometer.y, msg.magnetometer.z, msg.imuTempC,
                         msg.pressurePascals, msg.pressureTempC])

        data = np.array(data)
        ind = np.argsort(data[:, 0])
        data = data[ind].T

        plt.figure(num=None, figsize=(15.3, 10.7), dpi=80)
        plt.rc('font', **{'size': 20})
        plt.rc('axes', linewidth=4)
        plt.subplot(311)
        plt.plot(data[0], data[1], '.')
        plt.ylabel('X-Accel (m/s^2)')
        plt.subplot(312)
        plt.plot(data[0], data[2], '.')
        plt.ylabel('Y-Accel (m/s^2)')
        plt.subplot(313)
        plt.plot(data[0], data[3], '.')
        plt.ylabel('Z-Accel (m/s^2)')
        plt.tight_layout()

        plt.figure(num=None, figsize=(15.3, 10.7), dpi=80);
        plt.subplot(311)
        plt.plot(data[0], data[4], '.')
        plt.ylabel('X-Angular Vel (rad/s)')
        plt.subplot(312)
        plt.plot(data[0], data[5], '.')
        plt.ylabel('Y-Angular Vel (rad/s)')
        plt.subplot(313)
        plt.plot(data[0], data[6], '.')
        plt.ylabel('Z-Angular Vel (rad/s)')
        plt.tight_layout()

        plt.figure(num=None, figsize=(15.3, 10.7), dpi=80);
        plt.subplot(311)
        plt.plot(data[0], data[7], '.')
        plt.ylabel('X-Mag (Ga)')
        plt.subplot(312)
        plt.plot(data[0], data[8], '.')
        plt.ylabel('Y-Mag (Ga)')
        plt.subplot(313)
        plt.plot(data[0], data[9], '.')
        plt.ylabel('Z-Mag (Ga)')
        plt.tight_layout()

        plt.figure(num=None, figsize=(15.3, 10.7), dpi=80);
        plt.subplot(311)
        plt.plot(data[0], data[10], '.')
        plt.ylabel('IMU Temp (C)')
        plt.subplot(312)
        plt.plot(data[0], data[11], '.')
        plt.ylabel('Ambient Pressure (Pa)')
        plt.subplot(313)
        plt.plot(data[0], data[12], '.')
        plt.ylabel('Ambient Temp (C)')
        plt.tight_layout()

    def print_origin(self):
        print('Latiude of ENU coordinate system:', self.lat0, 'degrees')
        print('Longitude of ENU coordinate system:', self.lon0, 'degrees')
        print('Height above the WGS84 ellipsoid of the ENU coordinate system:',
              self.h0, 'meters')


