"""
PID Controller

components:
    follow attitude commands
    gps commands and yaw
    waypoint following
"""
import numpy as np
from controls.frame_utils import euler2RM

DRONE_MASS_KG = 0.5
GRAVITY = -9.81
MOI = np.array([0.005, 0.005, 0.01])
MAX_THRUST = 10.0
MAX_TORQUE = 1.0


class NonlinearController(object):

    def __init__(self,
                 k_p_lateral=6.0,
                 k_d_lateral=4.0,
                 k_p_alt=6.0,
                 k_d_alt=1.5,
                 k_p_roll=8.0,
                 k_p_pitch=8.0,
                 k_p_yaw=4.5,
                 k_p_p=20,
                 k_p_q=20,
                 k_p_r=10):
        """Initialize the controller object and control gains"""
        self.k_p_alt = k_p_alt
        self.k_d_alt = k_d_alt
        self.k_p_lateral = k_p_lateral
        self.k_d_lateral = k_d_lateral
        self.k_p_roll = k_p_roll
        self.k_p_pitch = k_p_pitch
        self.k_p_yaw = k_p_yaw
        self.k_p_p = k_p_p
        self.k_p_q = k_p_q
        self.k_p_r = k_p_r
        return

    def trajectory_control(self, position_trajectory, yaw_trajectory, time_trajectory, current_time):
        """Generate a commanded position, velocity and yaw based on the trajectory
        
        Args:
            position_trajectory: list of 3-element numpy arrays, NED positions
            yaw_trajectory: list yaw commands in radians
            time_trajectory: list of times (in seconds) that correspond to the position and yaw commands
            current_time: float corresponding to the current time in seconds
            
        Returns: tuple (commanded position, commanded velocity, commanded yaw)
                
        """

        ind_min = np.argmin(np.abs(np.array(time_trajectory) - current_time))
        time_ref = time_trajectory[ind_min]

        if current_time < time_ref:
            position0 = position_trajectory[ind_min - 1]
            position1 = position_trajectory[ind_min]

            time0 = time_trajectory[ind_min - 1]
            time1 = time_trajectory[ind_min]
            yaw_cmd = yaw_trajectory[ind_min - 1]

        else:
            yaw_cmd = yaw_trajectory[ind_min]
            if ind_min >= len(position_trajectory) - 1:
                position0 = position_trajectory[ind_min]
                position1 = position_trajectory[ind_min]

                time0 = 0.0
                time1 = 1.0
            else:

                position0 = position_trajectory[ind_min]
                position1 = position_trajectory[ind_min + 1]
                time0 = time_trajectory[ind_min]
                time1 = time_trajectory[ind_min + 1]

        position_cmd = (position1 - position0) * \
                       (current_time - time0) / (time1 - time0) + position0
        velocity_cmd = (position1 - position0) / (time1 - time0)

        return (position_cmd, velocity_cmd, yaw_cmd)

    def lateral_position_control(self, local_position_cmd, local_velocity_cmd, local_position, local_velocity,
                                 acceleration_ff=np.array([0.0, 0.0])):
        """Generate horizontal acceleration commands for the vehicle in the local frame

        Args:
            local_position_cmd: desired 2D position in local frame [north, east]
            local_velocity_cmd: desired 2D velocity in local frame [north_velocity, east_velocity]
            local_position: vehicle position in the local frame [north, east]
            local_velocity: vehicle velocity in the local frame [north_velocity, east_velocity]
            acceleration_ff: feed forward acceleration command
            
        Returns: desired vehicle 2D acceleration in the local frame [north, east]
        """
        acceleration_cmd = \
            self.k_p_lateral * (local_position_cmd - local_position) + \
            self.k_d_lateral * (local_velocity_cmd - local_velocity) + \
            acceleration_ff
        return acceleration_cmd

    def altitude_control(self, altitude_cmd, vertical_velocity_cmd, altitude, vertical_velocity, attitude,
                         acceleration_ff=0.0):
        """Generate vertical acceleration (thrust) command

        Args:
            altitude_cmd: desired vertical position (+up)
            vertical_velocity_cmd: desired vertical velocity (+up)
            altitude: vehicle vertical position (+up)
            vertical_velocity: vehicle vertical velocity (+up)
            attitude: the vehicle's current attitude, 3 element numpy array (roll, pitch, yaw) in radians
            acceleration_ff: feed forward acceleration command (+up)
            
        Returns: thrust command for the vehicle (+up)
        """
        alt_acceleration_cmd = \
            self.k_p_alt * (altitude_cmd - altitude) + \
            self.k_d_alt * (vertical_velocity_cmd - vertical_velocity) + \
            acceleration_ff
        rot_mat = euler2RM(*attitude)
        thrust_cmd = DRONE_MASS_KG * alt_acceleration_cmd / rot_mat[2, 2]
        return np.clip(thrust_cmd, 0.0, MAX_THRUST)

    def roll_pitch_controller(self, acceleration_cmd, attitude, thrust_cmd):
        """ Generate the roll rate and pitch rate commands in the body frame
        
        Args:
            acceleration_cmd: 2-element numpy array (north_acceleration_cmd,east_acceleration_cmd) in m/s^2
            attitude: 3-element numpy array (roll, pitch, yaw) in radians
            thrust_cmd: vehicle thrusts command in Newton
            
        Returns: 2-element numpy array, desired roll rate (p) and pitch rate (q) commands in radians/s
        """
        if thrust_cmd == 0.0:
            return np.array([0.0, 0.0])
        b_xy_cmd = np.clip(
            acceleration_cmd / (-thrust_cmd / DRONE_MASS_KG),
            -np.pi / 4,
            np.pi / 4
        )
        rot_mat = euler2RM(*attitude)
        b_dot_x_cmd = self.k_p_roll * (b_xy_cmd[0] - rot_mat[0, 2])
        b_dot_y_cmd = self.k_p_pitch * (b_xy_cmd[1] - rot_mat[1, 2])
        p_cmd, q_cmd = \
            1. / rot_mat[2, 2] * \
            np.array([[rot_mat[1, 0], -rot_mat[0, 0]],
                      [rot_mat[1, 1], -rot_mat[0, 1]]]).dot(np.array([b_dot_x_cmd, b_dot_y_cmd]))
        return np.array([p_cmd, q_cmd])

    def body_rate_control(self, body_rate_cmd, body_rate):
        """ Generate the roll, pitch, yaw moment commands in the body frame
        
        Args:
            body_rate_cmd: 3-element numpy array (p_cmd,q_cmd,r_cmd) in radians/second^2
            body_rate: 3-element numpy array (p,q,r) in radians/second^2
            
        Returns: 3-element numpy array, desired roll moment, pitch moment, and yaw moment commands in Newtons*meters
        """
        rate = np.array([self.k_p_p, self.k_p_q, self.k_p_r]) * (body_rate_cmd - body_rate)
        torque_cmd = MOI * rate
        torque_cmd_magnitude = np.linalg.norm(torque_cmd)
        if torque_cmd_magnitude > MAX_TORQUE:
            torque_cmd *= MAX_TORQUE / torque_cmd_magnitude
        return torque_cmd

    def yaw_control(self, yaw_cmd, yaw):
        """ Generate the target yaw rate
        
        Args:
            yaw_cmd: desired vehicle yaw in radians
            yaw: vehicle yaw in radians
        
        Returns: target yaw rate in radians/sec
        """
        yaw_error = yaw_cmd % (2.0 * np.pi) - yaw
        if yaw_error > np.pi:
            yaw_error -= 2.0 * np.pi
        elif yaw_error < -np.pi:
            yaw_error += 2.0 * np.pi
        return self.k_p_yaw * yaw_error
