import numpy as np
import matplotlib.pyplot as plt
from typing import Optional, Tuple
from pysimenv.core.system import MultipleSystem
from pysimenv.missile.model import PlanarMissile2dof, PlanarManVehicle2dof
from pysimenv.missile.guidance import PurePNG2dim
from pysimenv.missile.util import RelKin2dim, CloseDistCond, miss_distance


class Engagement2dim(MultipleSystem):
    def __init__(self, missile: PlanarMissile2dof, target: PlanarManVehicle2dof):
        super(Engagement2dim, self).__init__()
        self.missile = missile
        self.target = target
        self.rel_kin = RelKin2dim()
        self.close_dist_cond = CloseDistCond(r_threshold=10.0)

        self.attach_sim_objects([self.missile, self.target])

    # override
    def reset(self):
        super(Engagement2dim, self).reset()
        self.close_dist_cond.reset()

    # implement
    def initialize(self):
        x_M = self.missile.state
        x_T = self.target.state
        self.rel_kin.evaluate(x_M, x_T)

    # implement
    def forward(self):
        x_M = self.missile.state
        x_T = self.target.state
        self.rel_kin.evaluate(x_M, x_T)
        self.close_dist_cond.evaluate(self.rel_kin.r)
        self.target.forward()

    # implement
    def check_stop_condition(self) -> Tuple[bool, int]:
        to_stop = False

        missile_stop, _ = self.missile.check_stop_condition()
        if self.intercepted():  # probable interception
            to_stop = True
            self.flag = 1

        if missile_stop:  # stop due to the missile
            to_stop = True
            self.flag = 2

        if self.is_out_of_view():  # out of field-of-view limit
            to_stop = True
            self.flag = 3

        return to_stop, self.flag

    def intercepted(self) -> bool:
        return self.close_dist_cond.check()

    def is_out_of_view(self) -> bool:
        sigma = self.rel_kin.sigma
        return self.missile.is_out_of_view(sigma)

    def miss_distance(self) -> float:
        x_M = self.missile.history('x')
        x_T = self.target.history('x')

        p_M = x_M[:, 0:2]
        p_T = x_T[:, 0:2]
        d_miss = miss_distance(p_M, p_T)
        return d_miss

    def report(self):
        self.missile.report()
        if self.flag == 1:
            print("[engagement] The target has been intercepted!")
        else:
            print("[engagement] The target has been missed!")
        print("[engagement] Miss distance: {:.6f} (m)".format(self.miss_distance()))

    def plot(self):
        fig_axs = dict()

        fig_ax = self.missile.plot_path()
        self.target.plot_path(fig_ax)

        fig1, ax1 = fig_ax['fig'], fig_ax['ax']
        fig1.suptitle("2-dim flight path")
        fig_axs['path'] = {'fig': fig1, 'ax': ax1}

        # plots for rel. kin.
        time_list = self.missile.history('t')
        x_M_list = self.missile.history('x')
        x_T_list = self.target.history('x')

        relKin = RelKin2dim()
        r_list = []
        sigma_list = []
        lam_list = []
        omega_list = []

        zem_list = []

        for i in range(x_M_list.shape[0]):
            x_M = x_M_list[i, :]
            x_T = x_T_list[i, :]

            relKin.evaluate(x_M, x_T)

            r_list.append(relKin.r)
            sigma_list.append(relKin.sigma)
            lam_list.append(relKin.lam)
            omega_list.append(relKin.omega)

            zem_list.append(relKin.zem)

        r_list = np.array(r_list)
        sigma_list = np.array(sigma_list)
        lam_list = np.array(lam_list)
        omega_list = np.array(omega_list)

        zem_list = np.array(zem_list)

        fig2, ax2 = plt.subplots(2, 2)
        fig2.tight_layout()
        ax2_1 = ax2[0, 0]
        ax2_1.set_title("Rel. dist")
        ax2_1.plot(time_list[:-1], r_list[:-1], label="Rel. dist")
        ax2_1.set_xlabel("Time (s)")
        ax2_1.set_ylabel("Rel. dist (m)")
        ax2_1.grid()
        ax2_1.legend()

        ax2_2 = ax2[0, 1]
        ax2_2.set_title("Look angle")
        ax2_2.plot(time_list[:-1], np.rad2deg(sigma_list[:-1]), label="look angle")
        ax2_2.set_xlabel("Time (s)")
        ax2_2.set_ylabel("Look angle (deg)")
        ax2_2.grid()
        ax2_2.legend()

        ax2_3 = ax2[1, 0]
        ax2_3.set_title("LOS angle")
        ax2_3.plot(time_list[:-1], np.rad2deg(lam_list[:-1]), label="LOS angle")
        ax2_3.set_xlabel("Time (s)")
        ax2_3.set_ylabel("LOS angle (deg)")
        ax2_3.grid()
        ax2_3.legend()

        ax2_4 = ax2[1, 1]
        ax2_4.set_title("LOS rate")
        ax2_4.plot(time_list[:-1], np.rad2deg(omega_list[:-1]), label="LOS rate")
        ax2_4.set_xlabel("Time (s)")
        ax2_4.set_ylabel("LOS rate (deg/s)")
        ax2_4.grid()
        ax2_4.legend()

        fig_axs['Rel. Kin.'] = {'fig': fig2, 'ax': ax2}

        fig3, ax3 = plt.subplots()
        ax3.set_title("ZEM")
        ax3.plot(time_list[:-1], zem_list[:-1], label="ZEM")
        ax3.set_xlabel("Time (s)")
        ax3.set_ylabel("ZEM (m)")
        ax3.grid()
        ax3.legend()

        fig_axs['Rel. Kin. add.'] = {'fig': fig3, 'ax': ax3}

        return fig_axs


class PurePNG2dimEngagement(Engagement2dim):
    def __init__(self, missile: PlanarMissile2dof, target: PlanarManVehicle2dof):
        super(PurePNG2dimEngagement, self).__init__(missile, target)

        self.pure_png = PurePNG2dim(N=3.0)
        self.attach_sim_objects([self.pure_png])

    # implement
    def forward(self):
        super(PurePNG2dimEngagement, self).forward()
        V_M = self.missile.V
        omega = self.rel_kin.omega

        a_M = self.pure_png.evaluate(V_M, omega)
        self.missile.forward(np.array([0, a_M]))
