import numpy as np
from scipy.spatial.transform import Rotation
import matplotlib.pyplot as plt

# from rotorpy.utils.axes3ds import Axes3Ds
from rotorpy.utils.animate import animate

import os

"""
Functions for showing the results from the simulator.

"""

class Plotter():

    def __init__(self, results, world):

        (self.time, self.x, self.x_des, self.v, 
        self.v_des, self.q, self.q_des, self.w, 
        self.s, self.s_des, self.M, self.T, self.wind,
        self.accel, self.gyro, self.accel_gt,
        self.x_mc, self.v_mc, self.q_mc, self.w_mc, 
        self.filter_state, self.covariance, self.sd) = self.unpack_results(results)

        self.R = Rotation.from_quat(self.q).as_matrix()
        self.R_mc = Rotation.from_quat(self.q_mc).as_matrix() # Rotation as measured by motion capture.

        self.world = world

        return

    def plot_results(self, plot_mocap, plot_estimator, plot_imu, fname=None):
        """
        Plot the results

        """

        # 3D Paths
        fig_3d = plt.figure('3D Path')
        # ax = Axes3Ds(fig)
        ax = fig_3d.add_subplot(projection='3d')
        self.world.draw(ax)
        ax.plot3D(self.x[:,0], self.x[:,1], self.x[:,2], 'b.')
        ax.plot3D(self.x_des[:,0], self.x_des[:,1], self.x_des[:,2], 'k')

        # Position and Velocity vs. Time
        (fig_posvel, axes) = plt.subplots(nrows=2, ncols=1, sharex=True, num='Pos/Vel vs Time')
        ax = axes[0]
        ax.plot(self.time, self.x_des[:,0], 'r', self.time, self.x_des[:,1], 'g', self.time, self.x_des[:,2], 'b')
        ax.plot(self.time, self.x[:,0], 'r.',    self.time, self.x[:,1], 'g.',    self.time, self.x[:,2], 'b.')
        ax.legend(('x', 'y', 'z'))
        ax.set_ylabel('position, m')
        ax.grid('major')
        ax.set_title('Position')
        ax = axes[1]
        ax.plot(self.time, self.v_des[:,0], 'r', self.time, self.v_des[:,1], 'g', self.time, self.v_des[:,2], 'b')
        ax.plot(self.time, self.v[:,0], 'r.',    self.time, self.v[:,1], 'g.',    self.time, self.v[:,2], 'b.')
        ax.legend(('x', 'y', 'z'))
        ax.set_ylabel('velocity, m/s')
        ax.set_xlabel('time, s')
        ax.grid('major')

        # Orientation and Angular Velocity vs. Time
        (fig_attrate, axes) = plt.subplots(nrows=2, ncols=1, sharex=True, num='Attitude/Rate vs Time')
        ax = axes[0]
        ax.plot(self.time, self.q_des[:,0], 'r', self.time, self.q_des[:,1], 'g', self.time, self.q_des[:,2], 'b', self.time, self.q_des[:,3], 'm')
        ax.plot(self.time, self.q[:,0], 'r.',    self.time, self.q[:,1], 'g.',    self.time, self.q[:,2], 'b.',    self.time, self.q[:,3],     'm.')
        ax.legend(('i', 'j', 'k', 'w'))
        ax.set_ylabel('quaternion')
        ax.set_xlabel('time, s')
        ax.grid('major')
        ax = axes[1]
        ax.plot(self.time, self.w[:,0], 'r.', self.time, self.w[:,1], 'g.', self.time, self.w[:,2], 'b.')
        ax.legend(('x', 'y', 'z'))
        ax.set_ylabel('angular velocity, rad/s')
        ax.set_xlabel('time, s')
        ax.grid('major')

        if plot_mocap:  # if mocap should be plotted. 
            # Motion capture position and velocity vs time
            (fig_mocapposvel, axes) = plt.subplots(nrows=2, ncols=1, sharex=True, num='Motion Capture Pos/Vel vs Time')
            ax = axes[0]
            ax.plot(self.time, self.x_mc[:,0], 'r.', self.time, self.x_mc[:,1], 'g.',    self.time, self.x_mc[:,2], 'b.')
            ax.legend(('x', 'y', 'z'))
            ax.set_ylabel('position, m')
            ax.grid('major')
            ax.set_title('MOTION CAPTURE Position/Velocity')
            ax = axes[1]
            ax.plot(self.time, self.v_mc[:,0], 'r.',    self.time, self.v_mc[:,1], 'g.',    self.time, self.v_mc[:,2], 'b.')
            ax.legend(('x', 'y', 'z'))
            ax.set_ylabel('velocity, m/s')
            ax.set_xlabel('time, s')
            ax.grid('major')
            # Motion Capture Orientation and Angular Velocity vs. Time
            (fig_mocapattrate, axes) = plt.subplots(nrows=2, ncols=1, sharex=True, num='Motion Capture Attitude/Rate vs Time')
            ax = axes[0]
            ax.plot(self.time, self.q_mc[:,0], 'r.',    self.time, self.q_mc[:,1], 'g.',    self.time, self.q_mc[:,2], 'b.',    self.time, self.q_mc[:,3],     'm.')
            ax.legend(('i', 'j', 'k', 'w'))
            ax.set_ylabel('quaternion')
            ax.set_xlabel('time, s')
            ax.grid('major')
            ax.set_title("MOTION CAPTURE Attitude/Rate")
            ax = axes[1]
            ax.plot(self.time, self.w_mc[:,0], 'r.', self.time, self.w_mc[:,1], 'g.', self.time, self.w_mc[:,2], 'b.')
            ax.legend(('x', 'y', 'z'))
            ax.set_ylabel('angular velocity, rad/s')
            ax.set_xlabel('time, s')
            ax.grid('major')

        # Commands vs. Time
        (fig_commands, axes) = plt.subplots(nrows=3, ncols=1, sharex=True, num='Commands vs Time')
        ax = axes[0]
        ax.plot(self.time, self.s_des[:,0], 'r', self.time, self.s_des[:,1], 'g', self.time, self.s_des[:,2], 'b',  self.time, self.s_des[:,3], 'k')
        ax.plot(self.time, self.s[:,0], 'r.',    self.time, self.s[:,1], 'g.',    self.time, self.s[:,2], 'b.',     self.time, self.s[:,3], 'k.')
        ax.legend(('1', '2', '3', '4'))
        ax.set_ylabel('motor speeds, rad/s')
        ax.grid('major')
        ax.set_title('Commands')
        ax = axes[1]
        ax.plot(self.time, self.M[:,0], 'r.', self.time, self.M[:,1], 'g.', self.time, self.M[:,2], 'b.')
        ax.legend(('x', 'y', 'z'))
        ax.set_ylabel('moment, N*m')
        ax.grid('major')
        ax = axes[2]
        ax.plot(self.time, self.T, 'k.')
        ax.set_ylabel('thrust, N')
        ax.set_xlabel('time, s')
        ax.grid('major')

        # Winds
        (fig_wind, axes) = plt.subplots(nrows=3, ncols=1, sharex=True, num='Winds vs Time')
        ax = axes[0]
        ax.plot(self.time, self.wind[:,0], 'r')
        ax.set_ylabel("wind X, m/s")
        ax.grid('major')
        ax.set_title('Winds')
        ax = axes[1]
        ax.plot(self.time, self.wind[:,1], 'g')
        ax.set_ylabel("wind Y, m/s")
        ax.grid('major')
        ax = axes[2]
        ax.plot(self.time, self.wind[:,2], 'b')
        ax.set_ylabel("wind Z, m/s")
        ax.set_xlabel("time, s")
        ax.grid('major')

        # IMU sensor
        if plot_imu:
            (fig_imu, axes) = plt.subplots(nrows=2, ncols=1, sharex=True, num="IMU Measurements vs Time")
            ax = axes[0]
            ax.plot(self.time, self.accel[:,0], 'r.', self.time, self.accel[:,1], 'g.', self.time, self.accel[:,2], 'b.')
            ax.plot(self.time, self.accel_gt[:,0], 'k', self.time, self.accel_gt[:,1], 'c', self.time, self.accel_gt[:,2], 'm')
            ax.set_ylabel("linear acceleration, m/s/s")
            ax.grid()
            ax = axes[1]
            ax.plot(self.time, self.gyro[:,0], 'r.', self.time, self.gyro[:,1], 'g.', self.time, self.gyro[:,2], 'b.')
            ax.set_ylabel("angular velocity, rad/s")
            ax.grid()
            ax.legend(('x','y','z'))
            ax.set_xlabel("time, s")

        if plot_estimator:
            if self.estimator_exists:
                N_filter = self.filter_state.shape[1]
                (fig_filter, axes) = plt.subplots(nrows=N_filter, ncols=1, sharex=True, num="Filter States vs Time")
                fig_filter.set_size_inches(11, 8.5)
                for i in range(N_filter):
                    ax = axes[i]
                    ax.plot(self.time, self.filter_state[:,i], 'k', )
                    ax.fill_between(self.time, self.filter_state[:,i]-self.sd[:,i], self.filter_state[:,i]+self.sd[:,i], alpha=0.3, color='k')
                    ax.set_ylabel("x"+str(i))
                ax.set_xlabel("Time, s")

                (fig_cov, axes) = plt.subplots(nrows=N_filter, ncols=1, sharex=True, num="Filter Covariance vs Time")
                fig_cov.set_size_inches(11, 8.5)
                for i in range(N_filter):
                    ax = axes[i]
                    ax.plot(self.time, self.sd[:,i]**2, 'k', )
                    ax.set_ylabel("cov(x"+str(i)+")")
                ax.set_xlabel("Time, s")

        if fname is not None:

            root_path = os.path.join(os.path.dirname(__file__),'..','data_out')
            fig_3d.savefig(os.path.join(root_path, fname+'_3dpath.png'))
            fig_posvel.savefig(os.path.join(root_path, fname+'_posvel.png'))
            fig_attrate.savefig(os.path.join(root_path, fname+'_attrate.png'))
            fig_commands.savefig(os.path.join(root_path, fname+'_commands.png'))
            fig_wind.savefig(os.path.join(root_path, fname+'_wind.png'))
            if plot_mocap:
                fig_mocapposvel.savefig(os.path.join(root_path, fname+'_mocapposvel.png'))
                fig_mocapattrate.savefig(os.path.join(root_path, fname+'_mocapattrate.png'))
            if plot_imu:
                fig_imu.savefig(os.path.join(root_path, fname+'_imu.png'))
            if plot_estimator:
                fig_filter.savefig(os.path.join(root_path, fname+'_filter.png'))
                fig_cov.savefig(os.path.join(root_path, fname+'_cov.png'))

        plt.show()

        return

    def animate_results(self, animate_wind, fname=None):
        """
        Animate the results
        
        """

        # Animation (Slow)
        # Instead of viewing the animation live, you may provide a .mp4 filename to save.
        ani = animate(self.time, self.x, self.R, self.wind, animate_wind, world=self.world, filename=fname)
        plt.show()

        return

    def unpack_results(self, result):

        # Unpack the dictionary of results
        time                = result['time']
        state               = result['state']
        control             = result['control']
        flat                = result['flat']
        imu_measurements    = result['imu_measurements']
        imu_gt              = result['imu_gt']
        mocap               = result['mocap_measurements']
        state_estimate      = result['state_estimate']

        # Unpack each result into NumPy arrays
        x = state['x']
        x_des = flat['x']
        v = state['v']
        v_des = flat['x_dot']

        q = state['q']
        q_des = control['cmd_q']
        w = state['w']

        s_des = control['cmd_motor_speeds']
        s = state['rotor_speeds']
        M = control['cmd_moment']
        T = control['cmd_thrust']

        wind = state['wind']

        accel   = imu_measurements['accel']
        gyro    = imu_measurements['gyro']

        accel_gt = imu_gt['accel']

        x_mc = mocap['x']
        v_mc = mocap['v']
        q_mc = mocap['q']
        w_mc = mocap['w']

        filter_state = state_estimate['filter_state']
        covariance = state_estimate['covariance']
        if filter_state.shape[1] > 0:
            sd = 3*np.sqrt(np.diagonal(covariance, axis1=1, axis2=2))
            self.estimator_exists = True
        else:
            sd = []
            self.estimator_exists = False

        return (time, x, x_des, v, v_des, q, q_des, w, s, s_des, M, T, wind, accel, gyro, accel_gt, x_mc, v_mc, q_mc, w_mc, filter_state, covariance, sd)

def plot_map(ax, world_data, equal_aspect=True, color=None, edgecolor=None, alpha=1, world_bounds=True, axes=True):
    """
    Plots the map in the world data in a top-down 2D view. 
    Inputs:
        ax: The axis to plot on
        world_data: The world data to plot
        equal_aspect: Determines if the aspect ratio of the plot should be equal.
        color: The color of the buildings. If None (default), it will use the color of the buildings. 
        edgecolor: The edge color of the buildings. If None (default), it will use the color of the buildings.
        alpha: The alpha value of the buildings. If None (default), it will use the color of the buildings.
        world_bounds: Whether or not to plot the world bounds as a dashed line around the 2D plot. 
        axes: Whether or not to plot the axis labels
    Outputs:
        Plots the map in the axis of interest. 
    """
    from matplotlib.patches import Rectangle

    # Add a dashed rectangle for the world bounds
    if world_bounds:
        world_patch = Rectangle((world_data['bounds']['extents'][0], world_data['bounds']['extents'][2]), 
                                world_data['bounds']['extents'][1]-world_data['bounds']['extents'][0], world_data['bounds']['extents'][3]-world_data['bounds']['extents'][2], 
                                linewidth=1, edgecolor='k', facecolor='none', linestyle='dashed')
        ax.add_patch(world_patch)

    plot_xmin = world_data['bounds']['extents'][0]
    plot_xmax = world_data['bounds']['extents'][1]
    plot_ymin = world_data['bounds']['extents'][2]
    plot_ymax = world_data['bounds']['extents'][3]

    for block in world_data['blocks']:
        xmin = block['extents'][0]
        xmax = block['extents'][1]
        ymin = block['extents'][2]
        ymax = block['extents'][3]
        if color is None:
            building_color = tuple(block['color'])
        else:
            building_color = color
        if edgecolor is None:
            building_edge_color = tuple(block['color'])
        else:
            building_edge_color = edgecolor
        block_patch = Rectangle((xmin, ymin), (xmax-xmin), (ymax-ymin), linewidth=1, edgecolor=building_edge_color, facecolor=building_color, alpha=alpha, fill=True)
        ax.add_patch(block_patch)

        if xmin < plot_xmin:
            plot_xmin = xmin
        if xmax > plot_xmax:
            plot_xmax = xmax
        if ymin < plot_ymin:
            plot_ymin = ymin
        if ymax > plot_ymax:
            plot_ymax = ymax

    ax.set_xlim([plot_xmin, plot_xmax])
    ax.set_ylim([plot_ymin, plot_ymax])

    if axes:
        ax.set_xlabel("X (m)")
        ax.set_ylabel("Y (m)")

    # Set the aspect ratio equal
    if equal_aspect:
        ax.set_aspect('equal')

    return

def plot_map(ax, world_data, equal_aspect=True, color=None, edgecolor=None, alpha=1, axes=True):
    """
    Plots the map in the world data in a top-down 2D view. 
    Inputs:
        ax: The axis to plot on
        world_data: The world data to plot
        equal_aspect: Determines if the aspect ratio of the plot should be equal.
        color: The color of the buildings. If None (default), it will use the color of the buildings. 
        edgecolor: The edge color of the buildings. If None (default), it will use the color of the buildings.
        alpha: The alpha value of the buildings. If None (default), it will use the color of the buildings.
        world_bounds: Whether or not to plot the world bounds as a dashed line around the 2D plot. 
        axes: Whether or not to plot the axis labels
    Outputs:
        Plots the map in the axis of interest. 
    """
    from matplotlib.patches import Rectangle

    plot_xmin = world_data['bounds']['extents'][0]
    plot_xmax = world_data['bounds']['extents'][1]
    plot_ymin = world_data['bounds']['extents'][2]
    plot_ymax = world_data['bounds']['extents'][3]

    for block in world_data['blocks']:
        xmin = block['extents'][0]
        xmax = block['extents'][1]
        ymin = block['extents'][2]
        ymax = block['extents'][3]
        if color is None:
            building_color = tuple(block['color'])
        else:
            building_color = color
        if edgecolor is None:
            building_edge_color = tuple(block['color'])
        else:
            building_edge_color = edgecolor
        block_patch = Rectangle((xmin, ymin), (xmax-xmin), (ymax-ymin), linewidth=1, edgecolor=building_edge_color, facecolor=building_color, alpha=alpha, fill=True)
        ax.add_patch(block_patch)

        if xmin < plot_xmin:
            plot_xmin = xmin
        if xmax > plot_xmax:
            plot_xmax = xmax
        if ymin < plot_ymin:
            plot_ymin = ymin
        if ymax > plot_ymax:
            plot_ymax = ymax

    ax.set_xlim([plot_xmin, plot_xmax])
    ax.set_ylim([plot_ymin, plot_ymax])

    if axes:
        ax.set_xlabel("X (m)")
        ax.set_ylabel("Y (m)")

    # Set the aspect ratio equal
    if equal_aspect:
        ax.set_aspect('equal')

    return

if __name__ == "__main__":

    from rotorpy.world import World

    # Get a list of the maps available under worlds. 
    available_worlds = [fname for fname in os.listdir(os.path.abspath(os.path.join(os.path.dirname(__file__),'..','worlds'))) if 'json' in fname]

    # Load a random world
    world_fname = np.random.choice(available_worlds)
    world = World.from_file(os.path.abspath(os.path.join(os.path.dirname(__file__),'..','worlds', world_fname)))
    
    # Plot the world. 
    (fig, ax) = plt.subplots(nrows=1, ncols=1, num="Top Down World View")
    plot_map(ax, world.world)
    ax.set_title(world_fname)
    plt.show()