"""Fixed-wing 6-DoF flight dynamics, batched torch: the same equations serve the plant and the oracle world model.

Internally the textbook frames of Beard & McLain (ch. 3-4): world NED, body FRD, state
    s = [p_ned(3), v_body uvw(3), q wxyz body->NED(4), w_frd pqr(3)]                         (B, 13)
and the action a = [throttle 0..1, aileron, elevator, rudder (rad)]                         (B, 4).
At the boundary everything is converted to dronepilot's shared 18-dim layout (world ENU z up, body FLU):
    x = [p, v world, R columns (body -> world), w body]                                    (B, 18)
so a plane is read by the same Pose / tasks / aura as every quad.

Forces: lift / drag with a flat-plate blend past stall (alpha0, M), side force, gravity, a simple electric
propeller (Beard & McLain 1st ed.). Moments: roll, pitch, yaw from the stability and control derivatives.
Integration: RK4 at a fixed substep, quaternion renormalized every substep.
"""
from __future__ import annotations

import math
from dataclasses import dataclass, field

import torch

T_NE = torch.tensor([[0.0, 1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, -1.0]])   # NED <-> ENU (its own inverse)
F_FL = torch.diag(torch.tensor([1.0, -1.0, -1.0]))                         # FRD <-> FLU


@dataclass
class Airframe:
    mass: float
    J: torch.Tensor                     # (3, 3) FRD
    S: float
    b: float
    c: float
    rho: float
    g: float
    aero: dict
    prop: dict
    surface_max: float
    extra: dict = field(default_factory=dict)

    @classmethod
    def from_cfg(cls, cfg: dict, dtype=torch.float32) -> "Airframe":
        I = cfg["inertia"]
        J = torch.tensor([[I["Jx"], 0.0, -I["Jxz"]], [0.0, I["Jy"], 0.0], [-I["Jxz"], 0.0, I["Jz"]]], dtype=dtype)
        g = cfg["geometry"]
        return cls(float(cfg["mass"]), J, g["S"], g["b"], g["c"], cfg["rho"], cfg["gravity"], dict(cfg["aero"]),
                   dict(cfg["propulsion"]), cfg["limits"]["surface_rad"])

    @property
    def AR(self) -> float:
        return self.b**2 / self.S


# ---------------------------------------------------------------- quaternions
def quat_to_R(q: torch.Tensor) -> torch.Tensor:
    w, x, y, z = (q / q.norm(dim=-1, keepdim=True)).unbind(-1)
    return torch.stack([
        torch.stack([1 - 2 * (y * y + z * z), 2 * (x * y - w * z), 2 * (x * z + w * y)], -1),
        torch.stack([2 * (x * y + w * z), 1 - 2 * (x * x + z * z), 2 * (y * z - w * x)], -1),
        torch.stack([2 * (x * z - w * y), 2 * (y * z + w * x), 1 - 2 * (x * x + y * y)], -1)], -2)


def R_to_quat(R: torch.Tensor) -> torch.Tensor:
    """(..., 3, 3) -> (..., 4) wxyz, robust branch on the largest diagonal term."""
    m = R
    tr = m[..., 0, 0] + m[..., 1, 1] + m[..., 2, 2]
    cands = torch.stack([tr, m[..., 0, 0], m[..., 1, 1], m[..., 2, 2]], -1)
    k = cands.argmax(-1)
    out = torch.empty(R.shape[:-2] + (4,), dtype=R.dtype)
    s0 = torch.sqrt((1 + tr).clamp(min=1e-12)) * 2
    q0 = torch.stack([0.25 * s0, (m[..., 2, 1] - m[..., 1, 2]) / s0, (m[..., 0, 2] - m[..., 2, 0]) / s0,
                      (m[..., 1, 0] - m[..., 0, 1]) / s0], -1)
    s1 = torch.sqrt((1 + m[..., 0, 0] - m[..., 1, 1] - m[..., 2, 2]).clamp(min=1e-12)) * 2
    q1 = torch.stack([(m[..., 2, 1] - m[..., 1, 2]) / s1, 0.25 * s1, (m[..., 0, 1] + m[..., 1, 0]) / s1,
                      (m[..., 0, 2] + m[..., 2, 0]) / s1], -1)
    s2 = torch.sqrt((1 + m[..., 1, 1] - m[..., 0, 0] - m[..., 2, 2]).clamp(min=1e-12)) * 2
    q2 = torch.stack([(m[..., 0, 2] - m[..., 2, 0]) / s2, (m[..., 0, 1] + m[..., 1, 0]) / s2, 0.25 * s2,
                      (m[..., 1, 2] + m[..., 2, 1]) / s2], -1)
    s3 = torch.sqrt((1 + m[..., 2, 2] - m[..., 0, 0] - m[..., 1, 1]).clamp(min=1e-12)) * 2
    q3 = torch.stack([(m[..., 1, 0] - m[..., 0, 1]) / s3, (m[..., 0, 2] + m[..., 2, 0]) / s3,
                      (m[..., 1, 2] + m[..., 2, 1]) / s3, 0.25 * s3], -1)
    for i, qi in enumerate((q0, q1, q2, q3)):
        sel = k == i
        out[sel] = qi[sel]
    out = out * torch.where(out[..., :1] < 0, -1.0, 1.0)
    return out / out.norm(dim=-1, keepdim=True)


# ---------------------------------------------------------------- frames
def to_shared(s: torch.Tensor) -> torch.Tensor:
    """Internal (B, 13) NED/FRD -> shared (B, 18) ENU/FLU."""
    T, F = T_NE.to(s.dtype), F_FL.to(s.dtype)
    R_nf = quat_to_R(s[..., 6:10])                                   # body FRD -> NED
    R_ef = T @ R_nf @ F                                              # body FLU -> ENU
    p = s[..., 0:3] @ T.T
    v = (R_nf @ s[..., 3:6, None])[..., 0] @ T.T
    w = s[..., 10:13] @ F.T
    return torch.cat([p, v, R_ef[..., :, 0], R_ef[..., :, 1], R_ef[..., :, 2], w], -1)


def from_shared(x: torch.Tensor) -> torch.Tensor:
    """Shared (B, 18) ENU/FLU -> internal (B, 13) NED/FRD."""
    T, F = T_NE.to(x.dtype), F_FL.to(x.dtype)
    R_ef = torch.stack([x[..., 6:9], x[..., 9:12], x[..., 12:15]], -1)
    R_nf = T @ R_ef @ F
    p = x[..., 0:3] @ T.T
    v_ned = x[..., 3:6] @ T.T
    v_body = (R_nf.transpose(-1, -2) @ v_ned[..., None])[..., 0]
    w = x[..., 15:18] @ F.T
    return torch.cat([p, v_body, R_to_quat(R_nf), w], -1)


# ---------------------------------------------------------------- forces and moments
def air_data(s: torch.Tensor, wind_ned: torch.Tensor | None):
    R = quat_to_R(s[..., 6:10])
    v_r = s[..., 3:6]
    if wind_ned is not None:
        v_r = v_r - (R.transpose(-1, -2) @ wind_ned[..., None])[..., 0]
    u, v, w = v_r.unbind(-1)
    Va = v_r.norm(dim=-1).clamp(min=0.5)
    alpha = torch.atan2(w, u)
    beta = torch.asin((v / Va).clamp(-1.0, 1.0))
    return R, Va, alpha, beta


def forces_moments(af: Airframe, s: torch.Tensor, a: torch.Tensor, wind_ned=None, mass: float | None = None):
    """Body FRD force (B, 3) and moment (B, 3). mass: the true mass if a payload changed it (gravity)."""
    A = af.aero
    R, Va, alpha, beta = air_data(s, wind_ned)
    p, q, r = s[..., 10:13].unbind(-1)
    dt, da, de, dr = a.unbind(-1)
    qbar = 0.5 * af.rho * Va**2 * af.S
    # lift with a flat-plate blend past stall
    M, a0 = A["M"], A["alpha0"]
    e1, e2 = torch.exp(-M * (alpha - a0)), torch.exp(M * (alpha + a0))
    sigma = (1 + e1 + e2) / ((1 + e1) * (1 + e2))
    CL_lin = A["C_L_0"] + A["C_L_alpha"] * alpha
    CL = (1 - sigma) * CL_lin + sigma * 2 * torch.sign(alpha) * torch.sin(alpha) ** 2 * torch.cos(alpha)
    CD = A["C_D_0"] + CL_lin**2 / (math.pi * A["e"] * af.AR)
    c2V, b2V = af.c / (2 * Va), af.b / (2 * Va)
    F_lift = qbar * (CL + A["C_L_q"] * c2V * q + A["C_L_delta_e"] * de)
    F_drag = qbar * (CD + A["C_D_q"] * c2V * q + A["C_D_delta_e"] * de.abs())
    ca, sa = torch.cos(alpha), torch.sin(alpha)
    fx = -ca * F_drag + sa * F_lift
    fz = -sa * F_drag - ca * F_lift
    fy = qbar * (A["C_Y_0"] + A["C_Y_beta"] * beta + A["C_Y_p"] * b2V * p + A["C_Y_r"] * b2V * r
                 + A["C_Y_delta_a"] * da + A["C_Y_delta_r"] * dr)
    # propeller along body x
    P = af.prop
    thrust = (0.5 * af.rho * P["S_prop"] * P["C_prop"] * ((P["k_motor"] * dt) ** 2 - Va**2)).clamp(min=0.0)
    fx = fx + thrust
    # gravity
    m = af.mass if mass is None else mass
    g_body = (R.transpose(-1, -2) @ torch.tensor([0.0, 0.0, m * af.g], dtype=s.dtype)[:, None])[..., 0]
    F = torch.stack([fx, fy, fz], -1) + g_body
    l = qbar * af.b * (A["C_ell_0"] + A["C_ell_beta"] * beta + A["C_ell_p"] * b2V * p + A["C_ell_r"] * b2V * r
                       + A["C_ell_delta_a"] * da + A["C_ell_delta_r"] * dr)
    pitch = qbar * af.c * (A["C_m_0"] + A["C_m_alpha"] * alpha + A["C_m_q"] * c2V * q + A["C_m_delta_e"] * de)
    n = qbar * af.b * (A["C_n_0"] + A["C_n_beta"] * beta + A["C_n_p"] * b2V * p + A["C_n_r"] * b2V * r
                       + A["C_n_delta_a"] * da + A["C_n_delta_r"] * dr)
    return F, torch.stack([l, pitch, n], -1)


def derivative(af: Airframe, s: torch.Tensor, a: torch.Tensor, wind_ned=None,
               mass: float | None = None, res: torch.Tensor | None = None,
               payload: tuple | None = None) -> torch.Tensor:
    """res (B, 6): a learned residual [linear acc (body FRD, m/s^2), angular acc (body FRD, rad/s^2)] added to the
    physics (fixed_wing/learned.py). None = pure physics.
    payload (m_p, r_frd): a point mass m_p (kg, already counted in `mass`) at r_frd (m, body FRD, from the CG): its
    weight's moment about the CG. Its inertia is the caller's (Airframe.J with the parallel-axis term)."""
    m = af.mass if mass is None else mass
    F, Mo = forces_moments(af, s, a, wind_ned, mass)
    q = s[..., 6:10]
    v = s[..., 3:6]
    w = s[..., 10:13]
    R = quat_to_R(q)
    p_dot = (R @ v[..., None])[..., 0]
    v_dot = -torch.cross(w, v, dim=-1) + F / m
    qw, qx, qy, qz = q.unbind(-1)
    p, qq, r = w.unbind(-1)
    q_dot = 0.5 * torch.stack([-qx * p - qy * qq - qz * r, qw * p + qy * r - qz * qq,
                               qw * qq + qz * p - qx * r, qw * r + qx * qq - qy * p], -1)
    J = af.J.to(s.dtype)
    Jw = w @ J.T
    w_dot = torch.linalg.solve(J, (Mo - torch.cross(w, Jw, dim=-1)).unsqueeze(-1)).squeeze(-1)
    if payload is not None:
        m_p, r = payload
        f_p = (R.transpose(-1, -2) @ torch.tensor([0.0, 0.0, m_p * af.g], dtype=s.dtype)[:, None])[..., 0]
        tau = torch.cross(torch.as_tensor(r, dtype=s.dtype).expand_as(f_p), f_p, dim=-1)
        w_dot = w_dot + torch.linalg.solve(J, tau.unsqueeze(-1)).squeeze(-1)
    if res is not None:
        v_dot = v_dot + res[..., 0:3]
        w_dot = w_dot + res[..., 3:6]
    return torch.cat([p_dot, v_dot, q_dot, w_dot], -1)


def rk4(af: Airframe, s: torch.Tensor, a: torch.Tensor, h: float, wind_ned=None, mass=None, res=None,
        payload=None) -> torch.Tensor:
    k1 = derivative(af, s, a, wind_ned, mass, res, payload)
    k2 = derivative(af, s + 0.5 * h * k1, a, wind_ned, mass, res, payload)
    k3 = derivative(af, s + 0.5 * h * k2, a, wind_ned, mass, res, payload)
    k4 = derivative(af, s + h * k3, a, wind_ned, mass, res, payload)
    out = s + (h / 6.0) * (k1 + 2 * k2 + 2 * k3 + k4)
    q = out[..., 6:10]
    return torch.cat([out[..., :6], q / q.norm(dim=-1, keepdim=True), out[..., 10:]], -1)