#!/usr/bin/env python
"""
Library handling projection operations of a standard camera model.

Note: the image coordiante system has its origin at the center of the top left
pixel.

"""
from __future__ import division, print_function, absolute_import
import numpy as np
import threading
import simplekml
import time
import bisect

# ROS imports
from tf.transformations import euler_from_quaternion, quaternion_from_euler, \
    quaternion_slerp

# ADAPT imports
from sensor_models.nav_conversions import enu_to_llh, llh_to_enu, \
    ned_quat_to_enu_quat
from sensor_models.utilities import get_2d_ellipsoid_from_3d_cov, \
    get_ellipse2d_pts, draw_moving_coordinate_system


lock = threading.Lock()


class NavStateProvider(object):
    """Camera manager object.

    This object supports requests for the navigation coordinate system state,
    at a particular time, relative to a local east/north/up coordinate system.

    Attributes:
    :param lat0: Latitude of the origin (deg).
    :type lat0: float

    :param lon0: Longitude of the origin (deg).
    :type lon0: float

    :param h0: Height above the WGS84 ellipsoid of the origin (meters).
    :type h0: float

    """
    def __init__(self):
        raise NotImplementedError

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

        :return: Position of the navigation coordinate system relative to a
            local level east/north/up coordinate system.
        :rtype: 3-pos

        """
        raise NotImplementedError

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

        :return: Quaternion (qx,qy,qz,qw) specifying the orientatio of the
            navigation coordinate system relative to a local level
            east/north/up (ENU) coordinate system. The quaternion represent a
            coordinate system rotation from ENU to the navigation coordinate
            system.
        :rtype: 4-array

        """
        raise NotImplementedError

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

        :return: List with the first element the 3-array position (see pos) and
            the second element the 4-array orientation quaternion (see quat).
        :rtype: list

        """
        raise NotImplementedError

    def __str__(self):
        return str(type(self))

    def __repr__(self):
        return self.__str__()


class NavStateFixed(NavStateProvider):
    def __init__(self, pos=np.array([0, 0, 0]),
                 quat=np.array([1/np.sqrt(2), 1/np.sqrt(2), 0, 0])):
        self._pos = pos
        self._quat = quat

    def pos(self, t):
        """See NavState documentation.

        """
        return self._pos

    def quat(self, t):
        """See NavState documentation.

        """
        return self._quat

    def pose(self, t):
        """See NavState documentation.

        """
        return [self._pos, self._quat]


class NavStateINS(NavStateProvider):
    def __init__(self, nav_states, lat0=None, lon0=None, h0=None):
        """
        :param nav_states: Each row encodes the reported navigation state from
            one time. The first column is time in UTC seconds. Columns 2-4 are
            the associated WGS84 geodetic latitude (degrees), longitude
            (degrees), and height above the WGS84 ellipsoid. Columns 5-8 are
            the associated orientation quaternions.
        :type nav_states: N x 8

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

        self._nav_states = nav_states

        if lat0 is not None:
            self.lat0 = lat0
        else:
            self.lat0 = np.median(nav_states[:, 1])

        if lon0 is not None:
            self.lon0 = lon0
        else:
            self.lon0 = np.median(nav_states[:, 2])

        if h0 is not None:
            self.h0 = h0
        else:
            self.h0 = np.min(nav_states[:, 3])

    def ins_state(self, t):
        """Interpolated INS state at time t.

        """
        ind = bisect.bisect(self._nav_states[:, 0], t)
        if ind == len(self._nav_states):
            ind -= 1
        elif ind == 0:
            ind += 1

        t1 = self._nav_states[ind - 1, 0]
        t2 = self._nav_states[ind, 0]

        if np.abs(t - t1) < 1e-4:
            return self._nav_states[ind - 1, 1:4], self._nav_states[ind - 1, 4:]
        elif np.abs(t - t2) < 1e-4:
            return self._nav_states[ind, 1:4], self._nav_states[ind, 4:]

        llh1 = self._nav_states[ind - 1, 1:4]
        quat1 = self._nav_states[ind - 1, 4:]
        llh2 = self._nav_states[ind, 1:4]
        quat2 = self._nav_states[ind, 4:]

        weight = (t - t1) / (t2 - t1)
        quat = quaternion_slerp(quat1, quat2, weight, spin=0,
                                shortestpath=True)

        # Interpolate position.
        llh = llh1*(1 - weight) + llh2*weight

        return llh, quat

    @property
    def times(self):
        """Return times at which pose is exactly specified.

        """
        return self._nav_states[:,0]

    def pos(self, t):
        """See NavState documentation.

        """
        llh = self.ins_state(t)[0]
        pos = llh_to_enu(llh[0], llh[1], llh[2], self.lat0, self.lon0, self.h0)
        return pos

    def quat(self, t):
        """See NavState documentation.

        """
        return self.ins_state(t)[1]

    def pose(self, t):
        """See NavState documentation.

        """
        llh, quat = self.ins_state(t)
        pos = llh_to_enu(llh[0], llh[1], llh[2], self.lat0, self.lon0, self.h0)
        return pos, quat


def trajectory_to_kmz(filename, times, lats, lons, alts, pos_covs=None,
                      min_dt=0.25, name='flight_path'):
    """Write navigation trajectory to kmz file to view in Google Earth.

    :param lat: Latitude of the detection (decimal degrees).
    :type lat: float

    :param lon: Longitude of the detection (decimal degrees).
    :type lon: float

    :param h: Height the detection above the WGS84 ellipsoid (m).
    :type h: float

    :param t: Time of the detection (seconds since Unix epoch).
    :type t: float

    :param cov: Covariance matrix describing the uncertainty in a local
        level east/north/up coordinate system centered at the detection
        (meters^2).
    :type cov: 3x3 numpy.ndarray

    :param camera_llh: Camera's latitude, longitude, and height.
    :type camera_llh: 3-array

    """
    kml = simplekml.Kml(name=name, open=1)
    assert np.all(np.diff(times) > 0)

    last_t = -np.inf
    for i in range(len(times)):
        lat = lats[i]
        lon = lons[i]
        h = alts[i]
        t = times[i]

        if t - last_t < min_dt:
            continue
        else:
            last_t = t

        if pos_covs is not None:
            cov = np.identity(3)
        else:
            cov = pos_covs[i]

        axes = get_2d_ellipsoid_from_3d_cov(cov)
        pts = get_ellipse2d_pts(np.zeros(2), axes, 30, plot_results=False)
        boundary = []

        # Use a linear approximation of the enu -> llh conversion.
        dlon_de = enu_to_llh(1, 0, 0, lat, lon, h, in_degrees=True)[1]
        dlon_de = dlon_de - lon
        dlat_dn = enu_to_llh(0, 1, 0, lat, lon, h, in_degrees=True)[0]
        dlat_dn = dlat_dn - lat

        lats_ = lat + dlat_dn*pts[1]
        lons_ = lon + dlon_de*pts[0]
        hs = np.ones(len(lons))*h
        boundary = zip(lons_, lats_, hs)

        pol = kml.newpolygon(name='ins')
        pol.outerboundaryis = boundary
        pol.style.linestyle.color = simplekml.Color.red
        pol.style.linestyle.width = 2
        pol.style.polystyle.color = simplekml.Color.changealphaint(100, simplekml.Color.red)
        when = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(t))
        pol.timestamp = simplekml.TimeStamp(when)

    kml.savekmz(filename)
