import numpy as np

from gym_quadrotor.dynamics.coordinates import body_z, world_to_body_matrix, body_to_world_matrix, body_to_world, \
    angvel_to_euler
from gym_quadrotor.dynamics.copter import CopterParams, DynamicsState


def linear_dynamics(params, state):
    """
    Calculates the linear acceleration of a quadcopter with parameters
    `params` that is currently in the dynamics state `state`.
    :param CopterParams params: Parameters of the quadrotor.
    :param DynamicsState state: Current dynamics state.
    :return: Linear acceleration in world frame.
    """
    m = params.mass
    b = params.thrust_factor
    Kt = params.translational_drag
    O = state.rotor_speeds
    n = state.attitude
    v = state.velocity

    thrust = b/m * (O[0]**2 + O[1]**2 + O[2]**2 + O[3]**2) * body_z(n)

    Ktw = np.dot(body_to_world_matrix(n), np.dot(np.diag(Kt), world_to_body_matrix(n)))
    drag = np.dot(Ktw, v) / m

    return thrust - drag + params.gravity


def propeller_torques(params, state):
    """
    Calculates the torques that are directly generated by the propellers.
    :return:
    """
    Lb = params.arm_length * params.thrust_factor
    d = params.drag_factor
    O = state.rotor_speeds
    motor_torque = O[3] ** 2 + O[1] ** 2 - O[2] ** 2 - O[0] ** 2
    B = np.array([Lb * (O[3] ** 2 - O[1] ** 2), Lb * (O[0] ** 2 - O[2] ** 2), d * motor_torque])
    return B


def angular_momentum_body_frame(params, state):
    """
    Calculates the angular momentum of a quadcopter with parameters
    `params` that is currently in the dynamics state `state`.
    :param CopterParams params: Parameters of the quadrotor.
    :param DynamicsState state: Current dynamics state.
    :return: angular acceleration in body frame.
    """
    J = params.rotor_inertia
    I = params.frame_inertia
    O = state.rotor_speeds
    w = state.angular_velocity
    Kr = params.rotational_drag

    gyro = state.net_rotor_speed * J * np.array([w[2], -w[1], 0])
    drag = Kr * w
    Mp = propeller_torques(params, state)
    B = Mp - drag + gyro - np.cross(w, I*w)

    return B


def euler_rate(state):
    """
    Calculates the rate of change in euler angle based on the current
    attitude and angular velocity.
    :param DynamicsState state: Current dynamics state.
    :return: d/dt(attitude)
    """
    return angvel_to_euler(state.attitude, state.angular_velocity)


def simulate_quadrotor(params, state, dt):
    """
    Simulate the dynamics of the quadrotor for the timestep given
    in `dt`. First the rotor speeds are updated according to the desired
    rotor speed, and then linear and angular accelerations are calculated
    and integrated.
    :param CopterParams params: Parameters of the quadrotor.
    :param DynamicsState state: Current dynamics state.
    """

    # let rotor speed approach desired rotor speed and prevent negative rotation
    gamma = 1.0 - 0.5**(dt / params.rotor_speed_half_time)
    dw = gamma * (state.desired_rotor_speeds - state.rotor_speeds)
    state._rotorspeeds += dw
    state._rotorspeeds = np.maximum(state._rotorspeeds, 0.0)

    acceleration = linear_dynamics(params, state)
    ang_momentum = angular_momentum_body_frame(params, state)
    angular_acc = ang_momentum / params.frame_inertia

    state._position += 0.5 * dt * dt * acceleration + 0.5 * dt * state._velocity
    state._velocity += dt * acceleration

    state._angular_velocity += dt * angular_acc
    state._attitude.rotate(dt * euler_rate(state))


def rotor_rotation_for_thrust(thrust, params):
    """
    Calculate the necessary rotation speed for a rotor in order to create `thrust` Newtons of thrust.
    :param thrust: Amount of thrust to be created, in Newtons.
    :param params: Quadcopter parameters.
    :return: Necessary angular velocity (about the z axis) of the rotor.
    """
    # b w^2 = thrust
    return np.sqrt(thrust / params.thrust_factor)
