import numpy as np
import scipy
import matplotlib.pyplot as plt
from pyquaternion import Quaternion
from pysimenv.core.base import SimObject
from pysimenv.common.model import FlatEarthEnv
from pysimenv.common.orientation import quaternion_to_axis_angle
from pysimenv.multicopter.base import MulticopterDyn


class BSControl(SimObject):
    """
    Back-stepping control
    """
    def __init__(self, alpha, m, J, **kwargs):
        super(BSControl, self).__init__(**kwargs)

        self.alpha = alpha
        self.m = m
        self.J = J

        J_x, J_y, J_z = np.diag(J)[:]
        self.a_1 = (J_y - J_z)/J_x
        self.a_3 = (J_z - J_x)/J_y
        self.a_5 = (J_x - J_y)/J_z
        self.b_1 = 1./J_x
        self.b_2 = 1./J_y
        self.b_3 = 1./J_z

    # implement
    def _forward(self, dyn: MulticopterDyn, sigma_d: np.ndarray, sigma_d_dot: np.ndarray = np.zeros(4)):
        """
        :param sigma_d: [x_d, y_d, z_d, psi_d]
        :param sigma_d_dot: [x_d_dot, y_d_dot, z_d_dot, psi_d_dot]
        :return:
        """
        alp_1, alp_2, alp_3, alp_4, alp_5, alp_6, alp_7, alp_8, alp_9, alp_10, alp_11, alp_12 = self.alpha[:]

        x_1, x_3, x_5 = dyn.euler_ang[:]  # phi, theta, psi
        x_2, x_4, x_6 = dyn.ang_vel[:]  # approximation of phi_dot, theta_dot, psi_dot
        x_9, x_11, x_7 = dyn.pos[:]  # x, y, z
        x_10, x_12, x_8 = dyn.vel[:]  # x_dot, y_dot, z_dot

        # translation subsystem control
        x_9_d, x_11_d, x_7_d = sigma_d[0:3]
        x_9_d_dot, x_11_d_dot, x_7_d_dot = sigma_d_dot[0:3]

        z_7 = x_7_d - x_7
        z_7_dot = x_7_d_dot - x_8
        z_8 = z_7_dot + alp_7*z_7

        z_9 = x_9_d - x_9
        z_9_dot = x_9_d_dot - x_10
        z_10 = z_9_dot + alp_9*z_9

        z_11 = x_11_d - x_11
        z_11_dot = x_11_d_dot - x_12
        z_12 = z_11_dot + alp_11*z_11

        u_1 = -self.m/(np.cos(x_1)*np.cos(x_3))*(
                -FlatEarthEnv.grav_accel + z_7 + alp_7*(z_8 - alp_7*z_7) + alp_8*z_8)
        u_x = -self.m/u_1*(
                z_9 + alp_9*(z_10 - alp_9*z_9) + alp_10*z_10
        )
        u_y = -self.m/u_1*(
                z_11 + alp_11*(z_12 - alp_11*z_11) + alp_12*z_12
        )

        x_5_d = sigma_d[3]
        x_1_d = np.arcsin(
            np.clip(np.sin(x_5_d)*u_x - np.cos(x_5_d)*u_y, -1., 1.))
        x_3_d = np.arcsin(
            np.clip((np.cos(x_5_d)*u_x + np.sin(x_5_d)*u_y)/np.cos(x_1_d), -1., 1.))

        # rotation subsystem control
        z_1 = x_1_d - x_1
        z_1_dot = -x_2
        z_2 = z_1_dot + alp_1*z_1

        z_3 = x_3_d - x_3
        z_3_dot = -x_4
        z_4 = z_3_dot + alp_3*z_3

        z_5 = x_5_d - x_5
        z_5_dot = -x_6
        z_6 = z_5_dot + alp_5*z_5

        u_2 = 1./self.b_1*(
                -self.a_1*x_4*x_6 + z_1 + alp_1*(z_2 - alp_1*z_1) + alp_2*z_2
        )
        u_3 = 1./self.b_2*(
                -self.a_3*x_2*x_6 + z_3 + alp_3*(z_4 - alp_3*z_3) + alp_4*z_4
        )
        u_4 = 1./self.b_3*(
                -self.a_5*x_2*x_4 + z_5 + alp_5*(z_6 - alp_5*z_5) + alp_6*z_6
        )
        return {'f': u_1, 'tau': np.array([u_2, u_3, u_4])}


class SMAttControl(SimObject):
    """
    Attitude control using sliding mode
    """
    def __init__(self, m, J, c, k, **kwargs):
        super(SMAttControl, self).__init__(**kwargs)
        self.m = m
        self.J = J
        self.c = c
        self.k = k

    # implement
    def _forward(self, dyn: MulticopterDyn, eta_d: np.ndarray, omega_d: np.ndarray = np.zeros(3),
                 omega_d_dot: np.ndarray = np.zeros(3)):
        switching_fun = lambda v: 2./np.pi*np.arctan(np.linalg.norm(v))*v

        eta = dyn.euler_ang
        omega = dyn.ang_vel

        e = eta_d - eta
        e_dot = omega_d - omega

        s = e_dot + self.c*e
        tau = np.cross(omega, self.J.dot(omega)) + self.J.dot(omega_d_dot + self.c*e_dot + self.k*switching_fun(s))

        f = self.m*FlatEarthEnv.grav_accel
        return {'f': f, 'tau': tau}

class FLVelControl(SimObject):
    """
    Feedback linearization velocity control
    Reference: H. Voos, "Nonlinear Control of a Quadrotor Micro-UAV using Feedback-Linearization",
    Proceedings of the 2009 IEEE International Conference on Mechatronics, 2009.
    """
    def __init__(self, m, J,
                 k_p_att: np.ndarray, k_d_att: np.ndarray, k_p_vel: np.ndarray, **kwargs):
        super(FLVelControl, self).__init__(**kwargs)
        self.m = m
        self.J = J  # (3, 3) array
        self.k_p_att = k_p_att  # (3,) array
        self.k_d_att = k_d_att  # (3,) array
        self.k_p_vel = k_p_vel  # (3,) array

    # implement
    def _forward(self, dyn: MulticopterDyn, v_d: np.ndarray = np.zeros(3)):
        """
        :param dyn: multi-copter dynamics
        :param v_d: desired velocity
        :return:
        """
        v = dyn.vel
        eta = dyn.euler_ang
        omega = dyn.ang_vel

        u_t_1, u_t_2, u_t_3 = self.k_p_vel*(v_d - v)[:]

        f = self.m*np.sqrt(u_t_1**2 + u_t_2**2 + (u_t_3 - FlatEarthEnv.grav_accel)**2)
        phi_d = np.arcsin(self.m/f*u_t_2)
        theta_d = np.arctan(u_t_1/(u_t_3 -FlatEarthEnv.grav_accel))

        eta_d = np.array([phi_d, theta_d, 0.])
        u_star = self.k_p_att*(eta_d - eta) - self.k_d_att*omega
        tau = self.J.dot(u_star) + np.cross(omega, self.J.dot(omega))

        return {'f': f, 'tau': tau}


class QuaternionAttControl(SimObject):
    """
    Attitude control based on quaternion
    Reference: J. Carino, H. Abaunza and P. Castillo,
    "Quadrotor Quaternion Control",
    2015 International Conference on Unmanned Aircraft Systems (ICUAS), 2015.
    """
    def __init__(self, J: np.ndarray, K: np.ndarray, **kwargs):
        super(QuaternionAttControl, self).__init__(**kwargs)
        self.J = J  # (3, 3) array
        self.K = K  # (3, 6) array

    # implement
    def _forward(self, dyn: MulticopterDyn, q_d: np.ndarray, omega_d: np.ndarray = np.zeros(3)) -> np.ndarray:
        q = dyn.quaternion
        omega = dyn.ang_vel

        a, phi = quaternion_to_axis_angle(q)
        theta = a*phi
        x = np.hstack((theta, omega))

        a_d, phi_d = quaternion_to_axis_angle(q_d)
        theta_d = a_d*phi_d
        x_d = np.hstack((theta_d, omega_d))

        u = -self.K.dot(x - x_d)
        tau = self.J.dot(u) + np.cross(omega, self.J.dot(omega))
        return tau

    @classmethod
    def gain(cls, Q: np.ndarray, R: np.ndarray) -> np.ndarray:
        A = np.zeros((6, 6))
        A[0:3, 3:6] = np.identity(3)

        B = np.zeros((6, 3))
        B[3:6, :] = np.identity(3)

        P = scipy.linalg.solve_continuous_are(A, B, Q, R)
        K = np.linalg.inv(R).dot(B.transpose().dot(P))
        return K


class QuaternionPosControl(SimObject):
    e3 = np.array([0., 0., 1.])
    n = np.array([0., 0., -1.])

    def __init__(self, m: float, K: np.ndarray, **kwargs):
        super(QuaternionPosControl, self).__init__(**kwargs)
        self.m = m
        self.K = K

    # implement
    def _forward(self, dyn: MulticopterDyn, p_d: np.ndarray, v_d: np.ndarray = np.zeros(3)):
        x = np.hstack((dyn.pos, dyn.vel))
        x_d = np.hstack((p_d, v_d))
        u = -self.K.dot(x - x_d)

        u_pd = u - FlatEarthEnv.grav_accel*self.e3
        f = self.m*np.linalg.norm(u_pd)

        r_d = np.hstack((
            np.array([self.n.dot(u_pd) + np.linalg.norm(u_pd)]),
            np.cross(self.n, u_pd)
        ))

        q_d = r_d/np.linalg.norm(r_d)

        x_dot = np.hstack((dyn.vel, u))
        u_pd_dot = -self.K.dot(x_dot)
        r_d_dot = np.hstack((
            np.array([self.n.dot(u_pd_dot) + u_pd.dot(u_pd_dot)/np.linalg.norm(u_pd)]),
            np.cross(self.n, u_pd_dot)
        ))
        q_d_dot = r_d_dot/np.linalg.norm(r_d) + r_d*(
            -r_d.dot(r_d_dot)/np.linalg.norm(r_d)**3)

        omega_d = 2*Quaternion(q_d).conjugate*Quaternion(q_d_dot)
        omega_d = omega_d.elements[1:4]
        return f, q_d, omega_d

    @classmethod
    def gain(cls, Q: np.ndarray, R: np.ndarray):
        A = np.zeros((6, 6))
        A[0:3, 3:6] = np.identity(3)

        B = np.zeros((6, 3))
        B[3:6, :] = np.identity(3)

        P = scipy.linalg.solve_continuous_are(A, B, Q, R)
        K = np.linalg.inv(R).dot(B.transpose().dot(P))
        return K


class OAFControl(SimObject):
    """
    Control in one actuator failure
    Reference: Y. Wu, K. Hu, X. M. Sun, and Y. Ma, "Nonlinear Control of Quadrotor for Fault Tolerance:
    A Total Failure of One Actuator", IEEE Transactions on Systems, Man, and Cybernetics: Systems, Vol. 51, No. 5,
    2021.
    """
    def __init__(self, m, J, D_v, D_omega, G, K_p, K_d, **kwargs):
        super(OAFControl, self).__init__(**kwargs)
        self.m = m
        self.J = J
        self.d_v = np.diag(D_v)
        self.d_omega = np.diag(D_omega)
        self.G = G
        self.K_p = K_p  # P gains for [x, y, z, phi, theta]
        self.K_d = K_d  # D gains for [x, y, z, phi, theta]

    # implement
    def _forward(self, dyn: MulticopterDyn, fault_ind: int, p_d, p_d_dot = np.zeros(3), p_d_2dot = np.zeros(3)):
        p = dyn.pos
        v = dyn.vel
        eta = dyn.euler_ang
        omega = dyn.ang_vel

        phi, theta, psi = eta[:]
        if np.abs(theta) > np.pi/2. - 1e-6:
            print("")
        assert np.abs(theta) < np.pi/2. - 1e-6, "[oaf_att_control] The attitude is near the singularity."

        c_phi, s_phi = np.cos(phi), np.sin(phi)
        c_theta, s_theta = np.cos(theta), np.sin(theta)
        c_psi, s_psi = np.cos(psi), np.sin(psi)
        t_theta = s_theta/c_theta

        # Position control (altitude)
        z, v_z = p[2], v[2]
        e_1, e_1_dot = p_d[2] - z, p_d_dot[2] - v_z
        k_p_1, k_d_1 = self.K_p[2], self.K_d[2]

        F_l = -self.m/(c_phi*c_theta)*(
                -FlatEarthEnv.grav_accel + 1./self.m*self.d_v[2]*v_z + k_p_1*e_1 + k_d_1*e_1_dot + p_d_2dot[2])

        # Position control (horizontal position)
        x, y = p[0:2]
        v_x, v_y = v[0:2]
        k_p_2, k_d_2 = self.K_p[0], self.K_d[0]
        k_p_3, k_d_3 = self.K_p[1], self.K_d[1]

        e_2, e_2_dot = p_d[0] - x, p_d_dot[0] - v_x
        e_3, e_3_dot = p_d[1] - y, p_d_dot[1] - v_y
        u_x = -self.m/F_l*(1./self.m*self.d_v[0]*v_x + k_p_2*e_2 + k_d_2*e_2_dot + p_d_2dot[0])
        u_y = -self.m/F_l*(1./self.m*self.d_v[1]*v_y + k_p_3*e_3 + k_d_3*e_3_dot + p_d_2dot[1])

        phi_d = np.arcsin(
            np.clip(s_psi*u_x - c_psi*u_y, -1., 1.)
        )
        theta_d = np.arcsin(
            np.clip((c_psi*u_x + s_psi*u_y)/np.cos(phi_d), -1., 1.)
        )
        phi_d = np.clip(phi_d, np.deg2rad(-15.), np.deg2rad(15.))
        theta_d = np.clip(theta_d, np.deg2rad(-15.), np.deg2rad(15.))
        self._logger.append(t=self.time, phi_d=phi_d, theta_d=theta_d)

        # Attitude control (phi, theta)
        H_inv = np.array([
            [1., s_phi*t_theta, c_phi*t_theta],
            [0., c_phi, -s_phi],
            [0., s_phi/c_theta, c_phi/c_theta]
        ])
        eta_dot = H_inv.dot(omega)
        phi_dot, theta_dot = eta_dot[0], eta_dot[1]

        eta_r = eta[0:2]
        eta_r_dot = eta_dot[0:2]
        e_4 = np.array([phi_d, theta_d]) - eta_r
        e_4_dot = -eta_r_dot  # phi_d_dot, theta_d_dot are approximated to be 0.
        k_p_4 = self.K_p[3:5]
        k_d_4 = self.K_d[3:5]

        M_1 = np.array([
            [1., s_phi*t_theta, c_phi*t_theta],
            [0., c_phi, -s_phi]
        ])
        M_1_dot = np.array([
            [0., c_phi*t_theta*phi_dot + (s_phi/c_theta**2)*theta_dot,
             -s_phi*t_theta*phi_dot + (c_phi/c_theta**2)*theta_dot],
            [0., -s_phi*phi_dot, -c_phi*phi_dot]
        ])

        G_r, g_psi = self.reduced_thrust_model(fault_ind=fault_ind)
        n_1 = np.linalg.solve(self.J, -np.cross(omega, self.J.dot(omega)) - self.d_omega*omega
                              + np.array([0., 0., g_psi[0]])*F_l)

        M_2 = np.matmul(np.linalg.inv(self.J), np.array([[1., 0.], [0., 1.], [g_psi[1], g_psi[2]]]))

        f = M_1_dot.dot(omega) + M_1.dot(n_1)
        Q = np.matmul(M_1, M_2)

        tau_r = np.linalg.pinv(Q, rcond=1e-8).dot(-f + k_p_4*e_4 + k_d_4*e_4_dot)
        u_r = np.array([F_l, tau_r[0], tau_r[1]])
        w_r = np.linalg.solve(G_r, u_r)

        w = np.hstack((w_r[0:fault_ind], np.array([0.]), w_r[fault_ind:]))
        return w

    def reduced_thrust_model(self, fault_ind: int):
        ind_i = [0, 1, 2]  # F_l, tau_phi, tau_theta
        ind_j = []
        for j in range(4):
            if not j == fault_ind:
                ind_j.append(j)
        G_r = self.G.take(ind_i, axis=0).take(ind_j, axis=1)
        g_psi = np.dot(self.G[3].take(ind_j), np.linalg.inv(G_r))
        return G_r, g_psi

    def plot_desired_att(self):
        t = self.history('t')
        phi_d = self.history('phi_d')
        theta_d = self.history('theta_d')

        fig, ax = plt.subplots()
        ax.plot(t, np.rad2deg(phi_d), label="phi_d")
        ax.plot(t, np.rad2deg(theta_d), label="theta_d")
        ax.set_xlabel("Time (s)")
        ax.set_ylabel("Desired angle (deg)")
        ax.set_title("Desired attitude")
        ax.grid()
        ax.legend()


class OAFAttControl(SimObject):
    """
    Attitude control in one actuator failure
    """
    def __init__(self, m, J, D_v, D_omega, G, K_p, K_d, **kwargs):
        super(OAFAttControl, self).__init__(**kwargs)
        self.m = m
        self.J = J
        self.d_v = np.diag(D_v)
        self.d_omega = np.diag(D_omega)
        self.G = G
        self.K_p = K_p
        self.K_d = K_d

    # implement
    def _forward(self, dyn: MulticopterDyn, fault_ind: int, zeta_d, zeta_d_dot = np.zeros(3), zeta_d_2dot = np.zeros(3)):
        # zeta = [phi, theta, z]
        eta = dyn.euler_ang
        phi, theta = eta[0], eta[1]
        assert np.abs(theta) < np.pi/2. - 1e-6, "[oaf_att_control] The attitude is near the singularity."

        c_phi, s_phi = np.cos(phi), np.sin(phi)
        c_theta, s_theta = np.cos(theta), np.sin(theta)
        t_theta = s_theta/c_theta

        # Altitude control
        z, v_z = dyn.pos[2], dyn.vel[2]
        e_1, e_1_dot = zeta_d[2] - z, zeta_d_dot[2] - v_z
        k_p_1, k_d_1 = self.K_p[2], self.K_d[2]
        z_d_2dot = zeta_d_2dot[2]

        F_l = -self.m/(c_phi*c_theta)*(
                -FlatEarthEnv.grav_accel + 1./self.m*self.d_v[2]*v_z + k_p_1*e_1 + k_d_1*e_1_dot + z_d_2dot)

        # Attitude control
        H_inv = np.array([
            [1., s_phi*t_theta, c_phi*t_theta],
            [0., c_phi, -s_phi],
            [0., s_phi/c_theta, c_phi/c_theta]
        ])
        omega = dyn.ang_vel
        eta_dot = H_inv.dot(omega)
        phi_dot, theta_dot = eta_dot[0], eta_dot[1]

        eta_r = np.array([phi, theta])
        eta_r_dot = np.array([phi_dot, theta_dot])
        e_3 = zeta_d[0:2] - eta_r
        e_3_dot = zeta_d_dot[0:2] - eta_r_dot
        k_p_2 = self.K_p[0:2]
        k_d_2 = self.K_d[0:2]
        eta_r_d_2dot = zeta_d_2dot[0:2]

        M_1 = np.array([
            [1., s_phi*t_theta, c_phi*t_theta],
            [0., c_phi, -s_phi]
        ])
        M_1_dot = np.array([
            [0., c_phi*t_theta*phi_dot + (s_phi/c_theta**2)*theta_dot,
             -s_phi*t_theta*phi_dot + (c_phi/c_theta**2)*theta_dot],
            [0., -s_phi*phi_dot, -c_phi*phi_dot]
        ])

        G_r, g_psi = self.reduced_thrust_model(fault_ind=fault_ind)
        n_1 = np.linalg.solve(self.J, -np.cross(omega, self.J.dot(omega)) - self.d_omega*omega
                              + np.array([0., 0., g_psi[0]])*F_l)

        M_2 = np.matmul(np.linalg.inv(self.J), np.array([[1., 0.], [0., 1.], [g_psi[1], g_psi[2]]]))

        f = M_1_dot.dot(omega) + M_1.dot(n_1)
        Q = np.matmul(M_1, M_2)

        tau_r = np.linalg.pinv(Q, rcond=1e-8).dot(
            -f + k_p_2*e_3 + k_d_2*e_3_dot + eta_r_d_2dot
        )
        u_r = np.array([F_l, tau_r[0], tau_r[1]])
        w_r = np.linalg.solve(G_r, u_r)

        w = np.hstack((w_r[0:fault_ind], np.array([0.]), w_r[fault_ind:]))
        return w

    def reduced_thrust_model(self, fault_ind: int):
        ind_i = [0, 1, 2]  # F_l, tau_phi, tau_theta
        ind_j = []
        for j in range(4):
            if not j == fault_ind:
                ind_j.append(j)
        G_r = self.G.take(ind_i, axis=0).take(ind_j, axis=1)
        g_psi = np.dot(self.G[3].take(ind_j), np.linalg.inv(G_r))
        return G_r, g_psi
