"""
Purpose-built QX-250 6-DOF quadcopter model + cascade flight controller.
Correct 125 mm geometry, real mass/inertia, motor lag, and a transparent
rate -> attitude -> position/heading cascade. Design & validation reference
for the Amesim realization. Pure numpy/scipy; no Amesim dependency.
"""
import numpy as np
from scipy.integrate import solve_ivp

# ----------------- QX-250 physical parameters -----------------
M   = 0.50            # all-up mass [kg]
G   = 9.80665         # gravity [m/s^2]
IXX = 0.0025          # roll inertia [kg m^2]
IYY = 0.0025          # pitch inertia
IZZ = 0.0045          # yaw inertia
ARM = 0.125           # centre->motor [m] (250 mm wheelbase / 2)
D   = ARM*np.sin(np.pi/4)   # roll/pitch moment arm for X layout [m]
CT  = 0.016           # prop torque/thrust ratio [m]  (yaw torque = CT * thrust diff)
TMAX = 6.1            # max thrust per motor [N]  (~5:1 thrust:weight)
TAU_M = 0.03          # motor+ESC first-order lag [s]
J = np.diag([IXX, IYY, IZZ])
Jinv = np.linalg.inv(J)

# ----------------- Cascade controller gains (tuned, well-damped) -----------------
# position (outer) -> desired velocity
KP_POS = np.array([1.6, 1.6, 3.0])       # x,y,z
KP_VEL = np.array([3.2, 3.2, 6.0])       # velocity P -> desired accel
KD_VEL = np.array([0.6, 0.6, 1.0])
# attitude -> desired body rate
KP_ATT = np.array([9.0, 9.0, 4.0])       # roll,pitch,yaw
# body rate -> moment (PD)
KP_RATE = np.array([0.055, 0.055, 0.06])
KD_RATE = np.array([0.0016, 0.0016, 0.0018])
TILT_MAX = np.radians(35)                # limit commanded tilt

def Rzyx(phi, th, psi):
    # body->world rotation, ZYX (yaw-pitch-roll)
    cph,sph=np.cos(phi),np.sin(phi); cth,sth=np.cos(th),np.sin(th); cps,sps=np.cos(psi),np.sin(psi)
    return np.array([
        [cth*cps, sph*sth*cps-cph*sps, cph*sth*cps+sph*sps],
        [cth*sps, sph*sth*sps+cph*cps, cph*sth*sps-sph*cps],
        [-sth,    sph*cth,             cph*cth]])

def euler_rates_matrix(phi, th):
    # body rates (p,q,r) -> euler angle rates
    cph,sph=np.cos(phi),np.sin(phi); cth=np.cos(th); tth=np.tan(th)
    return np.array([[1, sph*tth, cph*tth],
                     [0, cph, -sph],
                     [0, sph/cth, cph/cth]])

# X-quad motor layout (orthogonal mixer). Motors 1 FR, 2 FL, 3 RL, 4 RR.
# Spin: 1 CW, 2 CCW, 3 CW, 4 CCW.  Sign patterns are mutually orthogonal.
MIX = np.array([
    [ 1,  1,  1,  1],   # T
    [-D,  D,  D, -D],   # Mx roll  (right motors 1,4 down)
    [ D,  D, -D, -D],   # My pitch (front motors 1,2)
    [CT,-CT, CT,-CT]])  # Mz yaw   (CW 1,3 vs CCW 2,4)

def mixer_inverse(T, Mx, My, Mz):
    t = np.linalg.solve(MIX, np.array([T, Mx, My, Mz]))
    return np.clip(t, 0.0, TMAX)

def forces_moments(t_motors, phi, th, psi):
    fm = MIX.dot(t_motors)          # [T, Mx, My, Mz]
    return fm[0], fm[1:4]

def controller(state, ref):
    pos=state[0:3]; vel=state[3:6]; eul=state[6:9]; omega=state[9:12]
    phi,th,psi = eul
    pos_ref, yaw_ref = ref
    # outer position -> desired velocity -> desired accel
    vel_des = KP_POS*(pos_ref - pos)
    acc_des = KP_VEL*(vel_des - vel) - KD_VEL*vel
    # desired thrust (world z) and tilt
    az = acc_des[2] + G
    T_des = M*az / max(np.cos(phi)*np.cos(th), 0.5)
    T_des = np.clip(T_des, 0.0, 4*TMAX)
    # desired roll/pitch from horizontal accel (yaw-rotated)
    ax_b =  acc_des[0]*np.cos(psi) + acc_des[1]*np.sin(psi)
    ay_b = -acc_des[0]*np.sin(psi) + acc_des[1]*np.cos(psi)
    th_des  = np.clip( ax_b/G, -TILT_MAX, TILT_MAX)
    phi_des = np.clip(-ay_b/G, -TILT_MAX, TILT_MAX)
    # attitude -> desired body rate
    att_err = np.array([phi_des-phi, th_des-th, np.arctan2(np.sin(yaw_ref-psi),np.cos(yaw_ref-psi))])
    rate_des = KP_ATT*att_err
    # rate -> moments (PD)
    M_cmd = J.dot(KP_RATE*(rate_des-omega) - KD_RATE*omega) if False else \
            (KP_RATE*(rate_des-omega) - KD_RATE*omega)
    return T_des, M_cmd

def dynamics(t, s, tm_state, ref, dvec):
    pos=s[0:3]; vel=s[3:6]; eul=s[6:9]; omega=s[9:12]
    phi,th,psi=eul
    t_motors = tm_state
    T, Mvec = forces_moments(t_motors, phi, th, psi)
    # translational: gravity + body-thrust rotated to world + external disturbance force (ZOH)
    Rwb = Rzyx(phi,th,psi)
    thrust_world = Rwb.dot(np.array([0,0,T]))
    acc = np.array([0,0,-G]) + thrust_world/M + np.asarray(dvec)/M
    # rotational
    domega = Jinv.dot(Mvec - np.cross(omega, J.dot(omega)))
    deul = euler_rates_matrix(phi,th).dot(omega)
    return np.concatenate([vel, acc, deul, domega])

def simulate(pos0=(0,0,0), eul0=(0,0,0), ref=((0,0,0),0.0),
             disturb=lambda t: np.zeros(3), T_end=8.0, dt=0.005):
    s = np.zeros(12); s[0:3]=pos0; s[6:9]=eul0
    tm = np.full(4, M*G/4.0)          # motors start at hover thrust
    ts=[0.0]; S=[s.copy()]; TM=[tm.copy()]
    n=int(T_end/dt)
    for k in range(n):
        T_des, M_cmd = controller(s, ref)
        t_cmd = mixer_inverse(T_des, M_cmd[0], M_cmd[1], M_cmd[2])
        # motor first-order lag
        tm = tm + (t_cmd - tm)*(dt/TAU_M)
        tm = np.clip(tm, 0, TMAX)
        dvec = np.asarray(disturb(ts[-1]))     # ZOH disturbance at current global time
        sol = solve_ivp(dynamics, (0,dt), s, args=(tm, ref, dvec),
                        method='RK45', rtol=1e-7, atol=1e-9, max_step=dt)
        s = sol.y[:,-1]
        ts.append(ts[-1]+dt); S.append(s.copy()); TM.append(tm.copy())
    return np.array(ts), np.array(S), np.array(TM)

if __name__=="__main__":
    # quick self-check: disturbance recovery from 15 deg roll
    t,S,TM = simulate(eul0=(np.radians(15),0,0), ref=((0,0,0),0.0), T_end=6)
    roll=np.degrees(S[:,6]);
    print("roll0=%.1f  min=%.2f  final=%.3f  |final|<1deg=%s"%(roll[0], roll.min(), roll[-1], abs(roll[-1])<1))
    # station-keep: 1 m x step
    t2,S2,_=simulate(ref=((1,0,0),0.0), T_end=8)
    x=S2[:,0]; print("x step: final x=%.3f (target 1.0), overshoot max=%.3f"%(x[-1], x.max()))