"""Models of complete paraglider systems."""

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np
import scipy.integrate
import scipy.optimize

from pfh.glidersim import foil_aerodynamics, orientation
from pfh.glidersim.util import cross3, crossmat

if TYPE_CHECKING:
    from pfh.glidersim.paraglider_harness import ParagliderHarness
    from pfh.glidersim.paraglider_wing import ParagliderWing

__all__ = [
    "ParagliderSystemDynamics6a",
    "ParagliderSystemDynamics6b",
    "ParagliderSystemDynamics6c",
    "ParagliderSystemDynamics9a",
    "ParagliderSystemDynamics9b",
    "ParagliderSystemDynamics9c",
]

def __dir__():
    return __all__

class ParagliderSystemDynamics6a:
    """
    A 6 degrees-of-freedom paraglider model; there is no relative motion
    between the wing and the harness.

    This version uses the riser connection midpoint `RM` as the reference point
    for the angular momentum, and can include the effects of apparent mass.

    Parameters
    ----------
    wing : ParagliderWing
    payload : ParagliderHarness
        The harness model includes the mass of the pilot.
    use_apparent_mass : bool, optional
        Whether to estimate the effects of apparent inertia.
    """

    def __init__(
        self,
        wing: ParagliderWing,
        payload: ParagliderHarness,
        *,
        use_apparent_mass: bool = True,
    ) -> None:
        self.wing = wing
        self.payload = payload
        self.use_apparent_mass = use_apparent_mass

    def r_CP2RM(self, delta_a: float = 0, delta_w: float = 0):
        """
        Compute the reference points for the composite Paraglider system.

        All the components of the Paraglider that experience aerodynamic forces
        need their relative wind vectors. Each component is responsible for
        creating a list of the coordinates where they need the value of the
        wind. This function then transforms them into body coordinates.

        Parameters
        ----------
        delta_a : float [percentage]
            The fraction of maximum accelerator
        delta_w : float [percentage]
            The fraction of weight shift, from -1 (left) to +1 (right)

        Returns
        -------
        ndarray of float, shape (K,3) [m]
            The position of the control points with respect to `RM`.
        """
        r_LE2RM = -self.wing.r_RM2LE(delta_a)
        wing_cps = self.wing.r_CP2LE(delta_a=delta_a) + r_LE2RM
        payload_cps = self.payload.r_CP2RM(delta_w)
        return np.vstack((wing_cps, payload_cps))

    def accelerations(
        self,
        v_RM2e,
        omega_b2e,
        g,
        delta_a: float = 0,
        delta_bl: float = 0,
        delta_br: float = 0,
        delta_w: float = 0,
        rho_air: float = 1.225,
        v_W2e=(0, 0, 0),
        reference_solution: dict | None = None,
    ):
        r"""
        Compute the translational and angular accelerations about the center of mass.

        Parameters
        ----------
        v_RM2e : array of float, shape (3,) [m/s]
            Translational velocity of `RM` in body frd coordinates, where `RM`
            is the midpoint between the two riser connection points.
        omega_b2e : array of float, shape (3,) [rad/s]
            Angular velocity of the body, in body frd coordinates.
        g : array of float, shape (3,) [m/s^s]
            The gravity vector in body frd
        delta_a : float [percentage]
            The fraction of maximum accelerator
        delta_bl : float [percentage]
            The fraction of maximum left brake
        delta_br : float [percentage]
            The fraction of maximum right brake
        delta_w : float [percentage]
            The fraction of weight shift, from -1 (left) to +1 (right)
        rho_air : float [kg/m^3], optional
            Air density
        v_W2e : ndarray of float, shape (3,) or (K,3) [m/s], optional
            The wind relative to the earth, in body frd coordinates. If it is
            a (3,) array then the wind is uniform at every control point. If
            it is a (K,3) array then it is the vectors for each control point.
        reference_solution : dictionary, optional
            FIXME: docstring. See `Phillips.__call__`

        Returns
        -------
        a_RM2e : array of float, shape (3,) [m/s^2]
            Translational acceleration of `RM` in body frd coordinates.
        alpha_b2e : array of float, shape (3,) [rad/s^2]
            Angular acceleration of the body with respect to Earth as the time
            derivative of angular velocity taken with respect to the body
            frame, expressed in body frd coordinates
            :math:`\left( ^b \dot{\omega}_{b/e}^b \right)`.
        solution : dictionary
            FIXME: docstring. See `Phillips.__call__`
        """
        r_CP2RM = self.r_CP2RM(delta_a, delta_w)
        v_W2e = np.asfarray(v_W2e)
        if v_W2e.ndim > 1 and v_W2e.shape[0] != r_CP2RM.shape[0]:
            raise ValueError(
                "Number of distinct wind vectors in v_W2e do not match the "
                "number of control points",
            )
        v_RM2e = np.asfarray(v_RM2e)
        if v_RM2e.shape != (3,):
            raise ValueError("v_RM2e.shape != (3,)")

        # -------------------------------------------------------------------
        # Compute the inertia of the body (wing + payload) with respect to the
        # riser midpoint `RM`.
        r_RM2LE = self.wing.r_RM2LE(delta_a)
        wmp = self.wing.mass_properties(rho_air, r_RM2LE)  # R = RM
        pmp = self.payload.mass_properties(delta_w, [0, 0, 0])  # R = RM
        m_b = wmp["m_b"] + pmp["m_p"]
        J_b2RM = wmp["J_b2R"] + pmp["J_p2R"]
        r_B2RM = (wmp["m_b"] * wmp["r_B2R"] + pmp["m_p"] * pmp["r_P2R"]) / m_b

        # -------------------------------------------------------------------
        # Compute the relative wind vectors for each control point.
        v_CP2e = v_RM2e + cross3(omega_b2e, r_CP2RM)
        v_W2CP = v_W2e - v_CP2e

        # FIXME: "magic" indexing established by `self.r_CP2RM`
        v_W2CP_wing = v_W2CP[:-1]
        v_W2CP_payload = v_W2CP[-1]

        # -------------------------------------------------------------------
        # Compute the forces and moments of the wing
        try:
            f_b, g_b2RM, ref = self.wing.resultant_force(
                delta_a,
                delta_bl,
                delta_br,
                v_W2CP_wing,
                rho_air,
                g,
                r_RM2LE,
                wmp,
                reference_solution,
            )
        except foil_aerodynamics.ConvergenceError:
            # Maybe it can't recover once Gamma is jacked?
            print("\nBonk! Retrying with the default reference solution")
            f_b, g_b2RM, ref = self.wing.resultant_force(
                delta_a,
                delta_bl,
                delta_br,
                v_W2CP_wing,
                rho_air,
                g,
                r_RM2LE,
                wmp,
            )

        f_p, g_p2RM = self.payload.resultant_force(
            delta_w=delta_w,
            v_W2h=v_W2CP_payload,
            rho_air=rho_air,
            g=g,
            r_R2RM=[0, 0, 0],
            mp=pmp,
        )

        # ------------------------------------------------------------------
        # Compute the accelerations \dot{v_RM2e} and \dot{omega_b2e}
        #
        # Builds a system of equations by equating derivatives of translational
        # and angular momentum to the forces and moments.

        # fmt: off

        # Real mass momentums
        v_B2e = v_RM2e + cross3(omega_b2e, r_B2RM)
        p_b2e = m_b * v_B2e  # Linear
        h_b2RM = J_b2RM @ omega_b2e + m_b * cross3(r_B2RM, v_RM2e)  # Angular

        # Build the system matrices for the real mass
        A1 = [m_b * np.eye(3), -m_b * crossmat(r_B2RM)]
        A2 = [m_b * crossmat(r_B2RM), J_b2RM]
        A = np.block([A1, A2])
        B1 = f_b + f_p - cross3(omega_b2e, p_b2e)
        B2 = (  # ref: Hughes Eq:13, p58
            g_b2RM
            + g_p2RM
            - cross3(v_RM2e, p_b2e)
            - cross3(omega_b2e, h_b2RM)
        )

        if self.use_apparent_mass:
            amp = self.wing.apparent_mass_properties(
                rho_air,
                r_RM2LE,
                v_RM2e,
                omega_b2e,
            )
            A += amp["A_a2R"]  # Incorporate the apparent inertia
            B1 += (  # Apparent inertial force (Barrows Eq:61)
                -cross3(omega_b2e, amp["p_a2e"])
            )
            B2 += (  # Apparent inertial moment (Barrows Eq:64)
                -cross3(v_RM2e, amp["p_a2e"])
                - cross3(omega_b2e, amp["h_a2R"])
                + cross3(v_RM2e, amp["M_a"] @ v_RM2e)  # Remove the steady-state term
            )

        # fmt: on

        B = np.r_[B1, B2]
        x = np.linalg.solve(A, B)  # Solve for the derivatives
        a_RM2e = x[:3]  # In frame F_b
        alpha_b2e = x[3:]  # In frames F_b and F_e
        return a_RM2e, alpha_b2e, ref

    def equilibrium_state(
        self,
        delta_a: float = 0,
        delta_b: float = 0,
        rho_air: float = 1.225,
        alpha_0: float = None,
        theta_0: float = 0,
        v_0: float = 10,
        reference_solution: dict | None = None,
    ):
        """
        Compute the equilibrium glider state for given inputs.

        Assumes that the wing is symmetric about the xz-plane.

        Parameters
        ----------
        delta_a : float, optional
            Fraction of accelerator application, where `0 <= delta_a <= 1`
        delta_b : float, optional
            Fraction of symmetric brake application, where `0 <= delta_b <= 1`
        rho_air : float [kg/m^3], optional
            Air density
        alpha_0 : float [rad], optional
            An initial proposal for the body angle of attack. If no value is
            set, the wing equilibrium alpha will be used.
        theta_0 : float [rad], optional
            An initial proposal for the body pitch angle.
        v_0 : float [m/s], optional
            An initial proposal for the body airspeed.
        reference_solution : dictionary, optional
            FIXME: docstring. See `Phillips.__call__`

        Returns
        -------
        dictionary
            alpha_b : float [radians]
                Wing angle of attack
            gamma_b : float [radians]
                Wing glide angle
            glide_ratio : float
                Units of ground distance traveled per unit of altitude lost
            Theta_b2e : array of float, shape (3,) [radians]
                Equilibrium orientation of the body relative to Earth as a set
                of Tait-Bryan yaw-pitch-role angles.
            v_RM2e : float [m/s]
                Steady-state velocity of the riser midpoint in body coordinates
            solution : dictionary
                FIXME: docstring. See `Phillips.__call__`
        """
        if alpha_0 is None:
            alpha_0 = self.wing.equilibrium_alpha(
                delta_a=delta_a,
                delta_b=delta_b,
                v_mag=v_0,
                rho_air=rho_air,
                reference_solution=reference_solution,
            )

        _state = {
            "delta_a": delta_a,
            "delta_bl": delta_b,
            "delta_br": delta_b,
            "omega_b2e": np.zeros(3),
            "rho_air": rho_air,
            "reference_solution": reference_solution,
        }

        def _helper(x, state):
            v_x, v_z, theta_b2e = x
            a_RM2e, alpha_b2e, ref = self.accelerations(
                v_RM2e=[v_x, 0, v_z],
                g=9.8 * np.array([-np.sin(theta_b2e), 0, np.cos(theta_b2e)]),
                **state,
            )
            _state["reference_solution"] = ref
            return (a_RM2e[0], a_RM2e[2], alpha_b2e[1])

        v_x_0 = v_0 * np.cos(alpha_0)
        v_z_0 = v_0 * np.sin(alpha_0)
        x0 = np.array([v_x_0, v_z_0, theta_0])
        res = scipy.optimize.root(_helper, x0, _state)
        v_RM2e = np.array([res.x[0], 0, res.x[1]])  # In body frd
        alpha_eq = np.arctan2(v_RM2e[2], v_RM2e[0])
        theta_eq = res.x[2]
        gamma_eq = alpha_eq - theta_eq

        equilibrium = {
            "alpha_b": alpha_eq,
            "gamma_b": gamma_eq,
            "glide_ratio": 1 / np.tan(gamma_eq),
            "Theta_b2e": np.array([0, theta_eq, 0]),
            "v_RM2e": v_RM2e,
            "reference_solution": _state["reference_solution"],
        }

        return equilibrium

class ParagliderSystemDynamics6b(ParagliderSystemDynamics6a):
    """
    A 6 degrees-of-freedom paraglider model; there is no relative motion
    between the wing and the harness.

    This version uses the body center of mass `B` as the reference point for
    the angular momentum. Using the center of mass produces a decoupled linear
    system, which is easier to reason about, making this model useful for
    validating other models. The system solves for `a_B2e` which is then used
    to compute `a_RM2e`.

    This model does not support apparent mass; the apparent mass model requires
    that the reference point lies in the xz-plane, which is not the case for
    `B` during weight shift control. As a result, this model is intended to
    help validate other models involving the real mass only.

    Parameters
    ----------
    wing : ParagliderWing
    payload : ParagliderHarness
        The harness model includes the mass of the pilot.
    use_apparent_mass : bool, optional
        Whether to estimate the effects of apparent inertia.
    """

    def __init__(self, wing: ParagliderWing, payload: ParagliderHarness) -> None:
        self.wing = wing
        self.payload = payload

    def accelerations(
        self,
        v_RM2e,
        omega_b2e,
        g,
        delta_a: float = 0,
        delta_bl: float = 0,
        delta_br: float = 0,
        delta_w: float = 0,
        rho_air: float = 1.225,
        v_W2e=(0, 0, 0),
        reference_solution=None,
    ):
        r"""
        Compute the translational and angular accelerations about the center of mass.

        Parameters
        ----------
        v_RM2e : array of float, shape (3,) [m/s]
            Translational velocity of `RM` in body frd coordinates, where `RM` is
            the midpoint between the two riser connection points.
        omega_b2e : array of float, shape (3,) [rad/s]
            Angular velocity of the body, in body frd coordinates.
        g : array of float, shape (3,) [m/s^s]
            The gravity vector in body frd
        delta_a : float [percentage]
            The fraction of maximum accelerator
        delta_bl : float [percentage]
            The fraction of maximum left brake
        delta_br : float [percentage]
            The fraction of maximum right brake
        delta_w : float [percentage]
            The fraction of weight shift, from -1 (left) to +1 (right)
        rho_air : float [kg/m^3], optional
            Air density
        v_W2e : ndarray of float, shape (3,) or (K,3) [m/s], optional
            The wind relative to the earth, in body frd coordinates. If it is
            a (3,) array then the wind is uniform at every control point. If
            it is a (K,3) array then it is the vectors for each control point.
        reference_solution : dictionary, optional
            FIXME: docstring. See `Phillips.__call__`

        Returns
        -------
        a_RM2e : array of float, shape (3,) [m/s^2]
            Translational acceleration of `RM` in body frd coordinates.
        alpha_b2e : array of float, shape (3,) [rad/s^2]
            Angular acceleration of the body with respect to Earth as the time
            derivative of angular velocity taken with respect to the body
            frame, expressed in body frd coordinates
            :math:`\left( ^b \dot{\omega}_{b/e}^b \right)`.
        solution : dictionary
            FIXME: docstring. See `Phillips.__call__`
        """
        r_CP2RM = self.r_CP2RM(delta_a, delta_w)
        v_W2e = np.asfarray(v_W2e)
        if v_W2e.ndim > 1 and v_W2e.shape[0] != r_CP2RM.shape[0]:
            raise ValueError(
                "Number of distinct wind vectors in v_W2e do not match the "
                "number of control points",
            )
        v_RM2e = np.asfarray(v_RM2e)
        if v_RM2e.shape != (3,):
            raise ValueError("v_RM2e.shape != (3,)")

        # -------------------------------------------------------------------
        # Compute the inertia of the body (wing + payload) with respect to its
        # center of mass `B`. Because `B` depends on the masses of both
        # components, this model must query those values and compute `B` before
        # computing properties with respect to `B`. This is wasteful, but this
        # model is mostly about checking the other models anyway.
        r_RM2LE = self.wing.r_RM2LE(delta_a)
        wmp0 = self.wing.mass_properties(rho_air, r_RM2LE)  # R = RM
        pmp0 = self.payload.mass_properties(delta_w, [0, 0, 0])  # R = RM
        m_b = wmp0["m_b"] + pmp0["m_p"]
        r_B2RM = (wmp0["m_b"] * wmp0["r_B2R"] + pmp0["m_p"] * pmp0["r_P2R"]) / m_b
        r_B2LE = r_B2RM + r_RM2LE

        # Recompute the mass properties using `B` as the reference point
        wmp = self.wing.mass_properties(rho_air, r_B2LE)  # R = B
        pmp = self.payload.mass_properties(delta_w, r_B2RM)  # R = B
        J_b2B = wmp["J_b2R"] + pmp["J_p2R"]

        # -------------------------------------------------------------------
        # Compute the relative wind vectors for each control point.
        v_CP2e = v_RM2e + cross3(omega_b2e, r_CP2RM)
        v_W2CP = v_W2e - v_CP2e

        # FIXME: "magic" indexing established by `self.r_CP2RM`
        v_W2CP_wing = v_W2CP[:-1]
        v_W2CP_payload = v_W2CP[-1]

        # -------------------------------------------------------------------
        # Compute the forces and moments of the wing
        try:
            f_b, g_b2B, ref = self.wing.resultant_force(
                delta_a,
                delta_bl,
                delta_br,
                v_W2CP_wing,
                rho_air,
                g,
                r_B2LE,
                wmp,
                reference_solution,
            )
        except foil_aerodynamics.ConvergenceError:
            # Maybe it can't recover once Gamma is jacked?
            print("\nBonk! Retrying with the default reference solution")
            f_b, g_b2B, ref = self.wing.resultant_force(
                delta_a,
                delta_bl,
                delta_br,
                v_W2CP_wing,
                rho_air,
                g,
                r_B2LE,
                wmp,
            )

        f_p, g_p2B = self.payload.resultant_force(
            delta_w=delta_w,
            v_W2h=v_W2CP_payload,
            rho_air=rho_air,
            g=g,
            r_R2RM=r_B2RM,
            mp=pmp,
        )

        # ------------------------------------------------------------------
        # Compute the accelerations \dot{v_RM2e} and \dot{omega_b2e}
        #
        # Builds a system of equations by equating derivatives of translational
        # and angular momentum to the net forces and moments.

        # Real mass momentums
        v_B2e = v_RM2e + cross3(omega_b2e, r_B2RM)
        p_b2e = m_b * v_B2e  # Linear
        h_b2B = J_b2B @ omega_b2e  # Angular

        A1 = [m_b * np.eye(3), np.zeros((3, 3))]
        A2 = [np.zeros((3, 3)), J_b2B]
        A = np.block([A1, A2])
        B1 = f_b + f_p - cross3(omega_b2e, p_b2e)
        B2 = g_b2B + g_p2B - np.cross(omega_b2e, h_b2B)  # ref: Hughes Eq:13, p58
        B = np.r_[B1, B2]
        x = np.linalg.solve(A, B)  # Solve for the derivatives
        a_B2e = x[:3]  # In frame F_b
        alpha_b2e = x[3:]  # In frames F_b and F_e
        a_RM2e = a_B2e - np.cross(alpha_b2e, r_B2RM)  # In frame F_b
        return a_RM2e, alpha_b2e, ref

class ParagliderSystemDynamics6c(ParagliderSystemDynamics6a):
    """
    A 6 degrees-of-freedom paraglider model; there is no relative motion
    between the wing and the harness.

    This version uses the body center of mass `B` as the reference point for
    the angular momentum. Similar to 6b, except it solves for v_RM2e directly.

    This model does not support apparent mass; the apparent mass model requires
    that the reference point lies in the xz-plane, which is not the case for
    `B` during weight shift control. As a result, this model is intended to
    help validate other models involving the real mass only.

    Parameters
    ----------
    wing : ParagliderWing
    payload : ParagliderHarness
        The harness model includes the mass of the pilot.
    """

    def __init__(self, wing: ParagliderWing, payload: ParagliderHarness) -> None:
        self.wing = wing
        self.payload = payload

    def accelerations(
        self,
        v_RM2e,
        omega_b2e,
        g,
        delta_a: float = 0,
        delta_bl: float = 0,
        delta_br: float = 0,
        delta_w: float = 0,
        rho_air: float = 1.225,
        v_W2e=(0, 0, 0),
        r_CP2RM=None,
        reference_solution: dict | None = None,
    ):
        r"""
        Compute the translational and angular accelerations about the center of mass.

        Parameters
        ----------
        v_RM2e : array of float, shape (3,) [m/s]
            Translational velocity of `RM` in body frd coordinates, where `RM` is
            the midpoint between the two riser connection points.
        omega_b2e : array of float, shape (3,) [rad/s]
            Angular velocity of the body, in body frd coordinates.
        g : array of float, shape (3,) [m/s^s]
            The gravity vector in body frd
        delta_a : float [percentage]
            The fraction of maximum accelerator
        delta_bl : float [percentage]
            The fraction of maximum left brake
        delta_br : float [percentage]
            The fraction of maximum right brake
        delta_w : float [percentage]
            The fraction of weight shift, from -1 (left) to +1 (right)
        rho_air : float [kg/m^3], optional
            Air density
        v_W2e : ndarray of float, shape (3,) or (K,3) [m/s], optional
            The wind relative to the earth, in body frd coordinates. If it is
            a (3,) array then the wind is uniform at every control point. If
            it is a (K,3) array then it is the vectors for each control point.
        reference_solution : dictionary, optional
            FIXME: docstring. See `Phillips.__call__`

        Returns
        -------
        a_RM2e : array of float, shape (3,) [m/s^2]
            Translational acceleration of `RM` in body frd coordinates.
        alpha_b2e : array of float, shape (3,) [rad/s^2]
            Angular acceleration of the body with respect to Earth as the time
            derivative of angular velocity taken with respect to the body
            frame, expressed in body frd coordinates
            :math:`\left( ^b \dot{\omega}_{b/e}^b \right)`.
        solution : dictionary
            FIXME: docstring. See `Phillips.__call__`
        """
        r_CP2RM = self.r_CP2RM(delta_a, delta_w)
        v_W2e = np.asfarray(v_W2e)
        if v_W2e.ndim > 1 and v_W2e.shape[0] != r_CP2RM.shape[0]:
            raise ValueError(
                "Number of distinct wind vectors in v_W2e do not match the "
                "number of control points",
            )
        v_RM2e = np.asfarray(v_RM2e)
        if v_RM2e.shape != (3,):
            raise ValueError("v_RM2e.shape != (3,)")

        # -------------------------------------------------------------------
        # Compute the inertia of the body (wing + payload) with respect to its
        # center of mass `B`. Because `B` depends on the masses of both
        # components, this model must query those values and compute `B` before
        # computing properties with respect to `B`. This is wasteful, but this
        # model is mostly about checking the other models anyway.
        r_RM2LE = self.wing.r_RM2LE(delta_a)
        wmp0 = self.wing.mass_properties(rho_air, r_RM2LE)  # R = RM
        pmp0 = self.payload.mass_properties(delta_w, [0, 0, 0])  # R = RM
        m_b = wmp0["m_b"] + pmp0["m_p"]
        r_B2RM = (wmp0["m_b"] * wmp0["r_B2R"] + pmp0["m_p"] * pmp0["r_P2R"]) / m_b
        r_B2LE = r_B2RM + r_RM2LE

        # Recompute the mass properties using `B` as the reference point
        wmp = self.wing.mass_properties(rho_air, r_B2LE)  # R = B
        pmp = self.payload.mass_properties(delta_w, r_B2RM)  # R = B
        J_b2B = wmp["J_b2R"] + pmp["J_p2R"]

        # -------------------------------------------------------------------
        # Compute the relative wind vectors for each control point.
        v_CP2e = v_RM2e + cross3(omega_b2e, r_CP2RM)
        v_W2CP = v_W2e - v_CP2e

        # FIXME: "magic" indexing established by `self.r_CP2RM`
        v_W2CP_wing = v_W2CP[:-1]
        v_W2CP_payload = v_W2CP[-1]

        # -------------------------------------------------------------------
        # Compute the forces and moments of the wing
        try:
            f_b, g_b2B, ref = self.wing.resultant_force(
                delta_a,
                delta_bl,
                delta_br,
                v_W2CP_wing,
                rho_air,
                g,
                r_B2LE,
                wmp,
                reference_solution,
            )
        except foil_aerodynamics.ConvergenceError:
            # Maybe it can't recover once Gamma is jacked?
            print("\nBonk! Retrying with the default reference solution")
            f_b, g_b2B, ref = self.wing.resultant_force(
                delta_a,
                delta_bl,
                delta_br,
                v_W2CP_wing,
                rho_air,
                g,
                r_B2LE,
                wmp,
            )

        f_p, g_p2B = self.payload.resultant_force(
            delta_w=delta_w,
            v_W2h=v_W2CP_payload,
            rho_air=rho_air,
            g=g,
            r_R2RM=r_B2RM,
            mp=pmp,
        )

        # ------------------------------------------------------------------
        # Compute the accelerations \dot{v_RM2e} and \dot{omega_b2e}
        #
        # Builds a system of equations by equating derivatives of translational
        # and angular momentum to the net forces and moments.

        # Real mass momentums
        v_B2e = v_RM2e + cross3(omega_b2e, r_B2RM)
        p_b2e = m_b * v_B2e  # Linear
        h_b2B = J_b2B @ omega_b2e  # Angular

        A1 = [m_b * np.eye(3), -m_b * crossmat(r_B2RM)]
        A2 = [np.zeros((3, 3)), J_b2B]
        A = np.block([A1, A2])
        B1 = f_b + f_p - cross3(omega_b2e, p_b2e)
        B2 = g_b2B + g_p2B - np.cross(omega_b2e, h_b2B)  # ref: Hughes Eq:13, p58
        B = np.r_[B1, B2]
        x = np.linalg.solve(A, B)  # Solve for the derivatives
        a_RM2e = x[:3]  # In frame F_b
        alpha_b2e = x[3:]  # In frames F_b and F_e
        return a_RM2e, alpha_b2e, ref

class ParagliderSystemDynamics9a:
    """
    A 9 degrees-of-freedom paraglider model, allowing rotation between the wing
    and the harness, with the connection modelled by spring-damper dynamics.

    This version uses the riser connection midpoint `RM` as the reference point
    for the angular momentum of both the body (the wing system) and the payload
    (the harness and pilot).

    Parameters
    ----------
    wing : ParagliderWing
    payload : ParagliderHarness
        The harness model includes the mass of the pilot.
    kappa_RM : array of float, shape (3,), optional
        Spring-damper coefficients for Theta_p2b (force as a linear function
        of angular displacement).
    kappa_RM_dot : array of float, shape (3,), optional
        Spring-damper coefficients for the derivative of Theta_p2b
    use_apparent_mass : bool, optional
        Whether to estimate the effects of apparent inertia. Default: True
    """

    def __init__(
        self,
        wing: ParagliderWing,
        payload: ParagliderHarness,
        kappa_RM=(0, 0, 0),
        kappa_RM_dot=(0, 0, 0),
        *,
        use_apparent_mass: bool = True,
    ) -> None:
        self.wing = wing
        self.payload = payload
        self._kappa_RM = np.asfarray(kappa_RM[:])
        self._kappa_RM_dot = np.asfarray(kappa_RM_dot[:])
        self.use_apparent_mass = use_apparent_mass

    def r_CP2RM(self, Theta_p2b, delta_a: float = 0, delta_w: float = 0):
        """
        Compute the reference points for the composite Paraglider system.

        All the components of the Paraglider that experience aerodynamic forces
        need their relative wind vectors. Each component is responsible for
        creating a list of the coordinates where they need the value of the
        wind. This function then transforms them into body coordinates.

        Parameters
        ----------
        Theta_p2b : array of float, shape (3,) [radians]
            The [phi, theta, gamma] of a yaw-pitch-roll sequence that encodes
            the relative orientation of the payload with respect to the body.
        delta_a : float [percentage]
            The fraction of maximum accelerator
        delta_w : float [percentage]
            The fraction of weight shift, from -1 (left) to +1 (right)

        Returns
        -------
        ndarray of float, shape (K,3) [m]
            The position of the control points with respect to `RM`.
        """
        r_LE2RM = -self.wing.r_RM2LE(delta_a)
        wing_cps = self.wing.r_CP2LE(delta_a=delta_a)  # In body frd
        payload_cps = self.payload.r_CP2RM(delta_w)  # In payload frd
        C_b2p = orientation.euler_to_dcm(Theta_p2b).T
        return np.vstack((wing_cps + r_LE2RM, (C_b2p @ payload_cps.T).T))

    def accelerations(
        self,
        v_RM2e,
        omega_b2e,
        omega_p2e,
        Theta_p2b,
        g,
        delta_a: float = 0,
        delta_bl: float = 0,
        delta_br: float = 0,
        delta_w: float = 0,
        rho_air: float = 1.225,
        v_W2e=(0, 0, 0),
        reference_solution: dict | None = None,
    ):
        r"""
        Compute the translational and angular accelerations about the center of mass.

        Parameters
        ----------
        v_RM2e : array of float, shape (3,) [m/s]
            Translational velocity of `RM` in body frd coordinates, where `RM` is
            the midpoint between the two riser connection points.
        omega_b2e : array of float, shape (3,) [rad/s]
            Angular velocity of the body, in body frd coordinates.
        omega_p2e : array of float, shape (3,) [rad/s]
            Angular velocity of the payload, in payload frd coordinates.
        Theta_p2b : array of float, shape (3,) [radians]
            The [phi, theta, gamma] of a yaw-pitch-roll sequence that encodes
            the relative orientation of the payload with respect to the body.
        g : array of float, shape (3,) [m/s^s]
            The gravity vector in body frd
        delta_a : float [percentage]
            The fraction of maximum accelerator
        delta_bl : float [percentage]
            The fraction of maximum left brake
        delta_br : float [percentage]
            The fraction of maximum right brake
        delta_w : float [percentage]
            The fraction of weight shift, from -1 (left) to +1 (right)
        rho_air : float [kg/m^3], optional
            Air density
        v_W2e : ndarray of float, shape (3,) or (K,3) [m/s], optional
            The wind relative to the earth, in body frd coordinates. If it is
            a (3,) array then the wind is uniform at every control point. If
            it is a (K,3) array then it is the vectors for each control point.
        reference_solution : dictionary, optional
            FIXME: docstring. See `Phillips.__call__`

        Returns
        -------
        a_RM2e : array of float, shape (3,) [m/s^2]
            Translational acceleration of `RM` in body frd coordinates.
        alpha_b2e : array of float, shape (3,) [rad/s^2]
            Angular acceleration of the body with respect to Earth as the time
            derivative of angular velocity taken with respect to the body
            frame, expressed in body frd coordinates
            :math:`\left( ^b \dot{\omega}_{b/e}^b \right)`.
        alpha_p2e : array of float, shape (3,) [rad/s^2]
            Angular acceleration of the payload with respect to Earth as the
            time derivative of angular velocity taken with respect to the
            payload frame, expressed in payload frd coordinates
            :math:`\left( ^p \dot{\omega}_{p/e}^p \right)`.
        solution : dictionary
            FIXME: docstring. See `Phillips.__call__`
        """
        r_CP2RM = self.r_CP2RM(Theta_p2b, delta_a, delta_w)
        v_W2e = np.asfarray(v_W2e)
        if v_W2e.ndim > 1 and v_W2e.shape[0] != r_CP2RM.shape[0]:
            raise ValueError(
                "Number of distinct wind vectors in v_W2e do not match the "
                "number of control points",
            )
        v_W2e = np.broadcast_to(v_W2e, r_CP2RM.shape)
        v_RM2e = np.asfarray(v_RM2e)
        if v_RM2e.shape != (3,):
            raise ValueError("v_RM2e.shape != (3,)")

        C_p2b = orientation.euler_to_dcm(Theta_p2b)
        C_b2p = C_p2b.T
        omega_p2b = C_b2p @ omega_p2e - omega_b2e  # In body frd
        omega_b2p = -omega_p2b

        # -------------------------------------------------------------------
        # Compute the inertia properties of the body and payload about `RM`.
        r_RM2LE = self.wing.r_RM2LE(delta_a)
        wmp = self.wing.mass_properties(rho_air, r_RM2LE)  # R = RM
        m_b = wmp["m_b"]
        J_b2RM = wmp["J_b2R"]
        r_B2RM = wmp["r_B2R"]

        pmp = self.payload.mass_properties(delta_w, [0, 0, 0])  # R = RM
        m_p = pmp["m_p"]
        J_p2RM = pmp["J_p2R"]
        r_P2RM = pmp["r_P2R"]

        # -------------------------------------------------------------------
        # Compute the relative wind vectors for each control point. Body
        # vectors are in body frd, payload vectors are in payload frd.
        v_B2e = v_RM2e + cross3(omega_b2e, r_B2RM)
        v_P2e = C_p2b @ v_RM2e + cross3(omega_p2e, r_P2RM)

        # FIXME: "magic" indexing established by `self.r_CP2RM`
        r_CP2RM_b = r_CP2RM[:-1]
        r_CP2RM_p = C_p2b @ r_CP2RM[-1]

        v_CP2e_b = v_B2e + cross3(omega_b2e, r_CP2RM_b - r_B2RM)
        v_CP2e_p = v_P2e + cross3(omega_p2e, r_CP2RM_p - r_P2RM)

        v_W2CP_b = v_W2e[:-1] - v_CP2e_b
        v_W2CP_p = C_p2b @ v_W2e[-1] - v_CP2e_p

        # -------------------------------------------------------------------
        # Forces and moments of the wing in body frd
        try:
            f_b, g_b2RM, ref = self.wing.resultant_force(
                delta_a,
                delta_bl,
                delta_br,
                v_W2CP_b,
                rho_air,
                g,
                r_RM2LE,
                wmp,
                reference_solution,
            )
        except foil_aerodynamics.ConvergenceError:
            # Maybe it can't recover once Gamma is jacked?
            print("\nBonk! Retrying with the default reference solution")
            f_b, g_b2RM, ref = self.wing.resultant_force(
                delta_a,
                delta_bl,
                delta_br,
                v_W2CP_b,
                rho_air,
                g,
                r_RM2LE,
                wmp,
            )

        f_p, g_p2RM = self.payload.resultant_force(  # In payload frd
            delta_w=delta_w,
            v_W2h=v_W2CP_p,
            rho_air=rho_air,
            g=C_p2b @ g,
            r_R2RM=[0, 0, 0],
            mp=pmp,
        )

        # Moment at the connection point `RM` modeled as a spring+damper system
        g_RM = Theta_p2b * self._kappa_RM + omega_p2b * self._kappa_RM_dot

        # ------------------------------------------------------------------
        # Build a system of equations by equating the time derivatives of the
        # translation and angular momentum (with respect to the Earth) of the
        # body and payload to the forces and moments on the body and payload.
        #
        # The four unknown vectors are the time derivatives of `v_RM2e`,
        # `omega_b2e` (in body frd), `omega_p2e` (in payload frd), and the
        # internal force on the risers, `F_RM` (in body frd).

        # fmt: off

        # Real mass momentums
        p_b2e = m_b * v_B2e  # Linear
        p_p2e = m_p * v_P2e
        h_b2RM = m_b * cross3(r_B2RM, v_RM2e) + J_b2RM @ omega_b2e  # Angular
        h_p2RM = m_p * cross3(r_P2RM, C_p2b @ v_RM2e) + J_p2RM @ omega_p2e

        # Build the system matrices for the real mass. A1 and A2 are the forces
        # and moments on the body, A3 and A4 are the forces and moments on the
        # payload. The unknowns are [a_RM2e^b, alpha_b2e^b, alpha_p2e^p, F_RM^b].
        I3, Z3 = np.eye(3), np.zeros((3, 3))
        A1 = [m_b * I3, -m_b * crossmat(r_B2RM), Z3, I3]
        A2 = [m_b * crossmat(r_B2RM), J_b2RM, Z3, Z3]
        A3 = [m_p * C_p2b, Z3, -m_p * crossmat(r_P2RM), -C_p2b]
        A4 = [m_p * crossmat(r_P2RM) @ C_p2b, Z3, J_p2RM, Z3]
        A = np.block([A1, A2, A3, A4])

        B1 = f_b - cross3(omega_b2e, p_b2e)
        B2 = (  # ref: Hughes Eq:13, p58
            g_b2RM
            - g_RM
            - cross3(v_RM2e, p_b2e)
            - cross3(omega_b2e, h_b2RM)
        )
        B3 = (
            f_p
            - m_p * C_p2b @ cross3(omega_b2p, v_RM2e)
            - cross3(omega_p2e, p_p2e)
        )
        B4 = (
            g_p2RM
            + C_p2b @ g_RM
            - cross3(C_p2b @ v_RM2e, p_p2e)
            - m_p * cross3(r_P2RM, C_p2b @ cross3(omega_b2p, v_RM2e))
            - cross3(omega_p2e, h_p2RM)
        )

        if self.use_apparent_mass:
            amp = self.wing.apparent_mass_properties(
                rho_air,
                r_RM2LE,
                v_RM2e,
                omega_b2e,
            )
            A[:6, :6] += amp["A_a2R"]  # Incorporate the apparent inertia
            B1 += (  # Apparent inertial force (Barrows Eq:61)
                -cross3(omega_b2e, amp["p_a2e"])
            )
            B2 += (  # Apparent inertial moment (Barrows Eq:64)
                -cross3(v_RM2e, amp["p_a2e"])
                - cross3(omega_b2e, amp["h_a2R"])
                + cross3(v_RM2e, amp["M_a"] @ v_RM2e)  # Remove the steady-state term
            )

        # fmt: on

        B = np.r_[B1, B2, B3, B4]
        x = np.linalg.solve(A, B)
        a_RM2e = x[:3]  # In frame F_b
        alpha_b2e = x[3:6]  # In frames F_b and F_e
        alpha_p2e = x[6:9]  # In frames F_p and F_e
        F_RM = x[9:]  # For debugging
        return a_RM2e, alpha_b2e, alpha_p2e, ref

    def equilibrium_state(
        self,
        delta_a: float = 0,
        delta_b: float = 0,
        rho_air: float = 1.225,
        alpha_0: float | None = None,
        theta_0: float = 0,
        v_0: float = 10,
        reference_solution=None,
    ):
        """
        Compute the equilibrium glider state for given inputs.

        Assumes that the wing is symmetric about the xz-plane.

        Parameters
        ----------
        delta_a : float, optional
            Fraction of accelerator application, where `0 <= delta_a <= 1`
        delta_b : float, optional
            Fraction of symmetric brake application, where `0 <= delta_b <= 1`
        rho_air : float [kg/m^3], optional
            Air density
        alpha_0 : float [rad], optional
            An initial proposal for the body angle of attack. If no value is
            set, the wing equilibrium alpha will be used.
        theta_0 : float [rad], optional
            An initial proposal for the body pitch angle.
        v_0 : float [m/s], optional
            An initial proposal for the body airspeed.
        reference_solution : dictionary, optional
            FIXME: docstring. See `Phillips.__call__`

        Returns
        -------
        dictionary
            alpha_b : float [radians]
                Wing angle of attack
            gamma_b : float [radians]
                Wing glide angle
            glide_ratio : float
                Units of ground distance traveled per unit of altitude lost
            Theta_b2e : array of float, shape (3,) [radians]
                Equilibrium orientation of the body relative to Earth as a set
                of Tait-Bryan yaw-pitch-role angles.
            v_RM2e : float [m/s]
                Steady-state velocity of the riser midpoint in body coordinates
            solution : dictionary
                FIXME: docstring. See `Phillips.__call__`
        """
        if alpha_0 is None:
            alpha_0 = self.wing.equilibrium_alpha(
                delta_a=delta_a,
                delta_b=delta_b,
                v_mag=v_0,
                rho_air=rho_air,
                reference_solution=reference_solution,
            )

        _state = {
            "delta_a": delta_a,
            "delta_bl": delta_b,
            "delta_br": delta_b,
            "omega_b2e": np.zeros(3),
            "omega_p2e": np.zeros(3),
            "rho_air": rho_air,
            "reference_solution": reference_solution,
        }

        def _helper(x, state):
            v_x, v_z, theta_b2e, theta_p2e = x
            theta_p2b = theta_p2e - theta_b2e
            a_RM2e, alpha_b2e, alpha_p2e, ref = self.accelerations(
                v_RM2e=[v_x, 0, v_z],
                g=9.8 * np.array([-np.sin(theta_b2e), 0, np.cos(theta_b2e)]),
                Theta_p2b=[0, theta_p2b, 0],
                **state,
            )
            _state["reference_solution"] = ref
            return (a_RM2e[0], a_RM2e[2], alpha_b2e[1], alpha_p2e[1])

        v_x_0 = v_0 * np.cos(alpha_0)
        v_z_0 = v_0 * np.sin(alpha_0)
        theta_p2e_0 = 0  # Assume the payload is hanging straight down
        x0 = np.array([v_x_0, v_z_0, theta_0, theta_p2e_0])
        res = scipy.optimize.root(_helper, x0, _state)
        v_RM2e = np.array([res.x[0], 0, res.x[1]])
        alpha_eq = np.arctan2(v_RM2e[2], v_RM2e[0])
        theta_b2e = res.x[2]
        theta_p2b = res.x[3] - res.x[2]
        gamma_eq = alpha_eq - theta_b2e

        equilibrium = {
            "alpha_b": alpha_eq,
            "gamma_b": gamma_eq,
            "glide_ratio": 1 / np.tan(gamma_eq),
            "Theta_b2e": np.array([0, theta_b2e, 0]),
            "Theta_p2b": np.array([0, theta_p2b, 0]),
            "v_RM2e": v_RM2e,
            "reference_solution": _state["reference_solution"],
        }

        return equilibrium

class ParagliderSystemDynamics9b(ParagliderSystemDynamics9a):
    """
    A 9 degrees-of-freedom paraglider model, allowing rotation between the wing
    and the harness, with the connection modelled by spring-damper dynamics.

    This model uses the body center of mass `B` as the reference point for the
    angular momentum of the body (the wing system) and the payload center of
    mass `P` for the angular momentum of the payload (the harness and pilot).

    This model does not support apparent mass. (Because it uses `B` as the
    reference point for the body dynamics, the system of equations for apparent
    mass would be in terms of `a_B2e`, which makes the equations involving the
    payload dynamics significantly messier.) Its purpose is to check for
    implementation errors in other models involving the real mass only.

    Parameters
    ----------
    wing : ParagliderWing
    payload : ParagliderHarness
        The harness model includes the mass of the pilot.
    kappa_RM : array of float, shape (3,), optional
        Spring-damper coefficients for Theta_p2b (force as a linear function
        of angular displacement).
    kappa_RM_dot : array of float, shape (3,), optional
        Spring-damper coefficients for the derivative of Theta_p2b
    """

    def __init__(
        self,
        wing: ParagliderWing,
        payload: ParagliderHarness,
        kappa_RM=(0, 0, 0),
        kappa_RM_dot=(0, 0, 0),
    ) -> None:
        self.wing = wing
        self.payload = payload
        self._kappa_RM = np.asfarray(kappa_RM[:])
        self._kappa_RM_dot = np.asfarray(kappa_RM_dot[:])

    def accelerations(
        self,
        v_RM2e,
        omega_b2e,
        omega_p2e,
        Theta_p2b,
        g,
        delta_a: float = 0,
        delta_bl: float = 0,
        delta_br: float = 0,
        delta_w: float = 0,
        rho_air: float = 1.225,
        v_W2e=(0, 0, 0),
        reference_solution: dict | None = None,
    ):
        r"""
        Compute the translational and angular accelerations about the center of mass.

        Parameters
        ----------
        v_RM2e : array of float, shape (3,) [m/s]
            Translational velocity of `RM` in body frd coordinates, where `RM` is
            the midpoint between the two riser connection points.
        omega_b2e : array of float, shape (3,) [rad/s]
            Angular velocity of the body, in body frd coordinates.
        omega_p2e : array of float, shape (3,) [rad/s]
            Angular velocity of the payload, in payload frd coordinates.
        Theta_p2b : array of float, shape (3,) [radians]
            The [phi, theta, gamma] of a yaw-pitch-roll sequence that encodes
            the relative orientation of the payload with respect to the body.
        g : array of float, shape (3,) [m/s^s]
            The gravity vector in body frd
        delta_a : float [percentage]
            The fraction of maximum accelerator
        delta_bl : float [percentage]
            The fraction of maximum left brake
        delta_br : float [percentage]
            The fraction of maximum right brake
        delta_w : float [percentage]
            The fraction of weight shift, from -1 (left) to +1 (right)
        rho_air : float [kg/m^3], optional
            Air density
        v_W2e : ndarray of float, shape (3,) or (K,3) [m/s], optional
            The wind relative to the earth, in body frd coordinates. If it is
            a (3,) array then the wind is uniform at every control point. If
            it is a (K,3) array then it is the vectors for each control point.
        reference_solution : dictionary, optional
            FIXME: docstring. See `Phillips.__call__`

        Returns
        -------
        a_RM2e : array of float, shape (3,) [m/s^2]
            Translational acceleration of `RM` in body frd coordinates.
        alpha_b2e : array of float, shape (3,) [rad/s^2]
            Angular acceleration of the body with respect to Earth as the time
            derivative of angular velocity taken with respect to the body
            frame, expressed in body frd coordinates
            :math:`\left( ^b \dot{\omega}_{b/e}^b \right)`.
        alpha_p2e : array of float, shape (3,) [rad/s^2]
            Angular acceleration of the payload with respect to Earth as the
            time derivative of angular velocity taken with respect to the
            payload frame, expressed in payload frd coordinates
            :math:`\left( ^p \dot{\omega}_{p/e}^p \right)`.
        solution : dictionary
            FIXME: docstring. See `Phillips.__call__`
        """
        r_CP2RM = self.r_CP2RM(Theta_p2b, delta_a, delta_w)
        v_W2e = np.asfarray(v_W2e)
        if v_W2e.ndim > 1 and v_W2e.shape[0] != r_CP2RM.shape[0]:
            raise ValueError(
                "Number of distinct wind vectors in v_W2e do not match the "
                "number of control points",
            )
        v_W2e = np.broadcast_to(v_W2e, r_CP2RM.shape)
        v_RM2e = np.asfarray(v_RM2e)
        if v_RM2e.shape != (3,):
            raise ValueError("v_RM2e.shape != (3,)")

        C_p2b = orientation.euler_to_dcm(Theta_p2b)
        C_b2p = C_p2b.T

        # -------------------------------------------------------------------
        # Compute the inertia properties of the body and payload with respect
        # to their centers of mass (`B` and `P`). This is wasteful since it
        # calls `mass_properties` twice just to get the centers of mass, but
        # it's simple and this model is just for validation anyway.
        r_RM2LE = self.wing.r_RM2LE(delta_a)
        wmp0 = self.wing.mass_properties(rho_air, [0, 0, 0])  # R = LE
        wmp = self.wing.mass_properties(rho_air, wmp0["r_B2R"])  # R = B
        m_b = wmp["m_b"]
        J_b2B = wmp["J_b2R"]
        r_B2RM = wmp0["r_B2R"] - r_RM2LE
        r_B2LE = r_B2RM + r_RM2LE

        pmp0 = self.payload.mass_properties(delta_w, [0, 0, 0])  # R = RM
        pmp = self.payload.mass_properties(delta_w, pmp0["r_P2R"])  # R = P
        m_p = pmp["m_p"]
        J_p2P = pmp["J_p2R"]
        r_P2RM = pmp0["r_P2R"]

        # -------------------------------------------------------------------
        # Compute the relative wind vectors for each control point. Body
        # vectors are in body frd, payload vectors are in payload frd.
        v_B2e = v_RM2e + cross3(omega_b2e, r_B2RM)
        v_P2e = C_p2b @ v_RM2e + cross3(omega_p2e, r_P2RM)

        # FIXME: "magic" indexing established by `self.r_CP2RM`
        r_CP2B_b = r_CP2RM[:-1] - r_B2RM
        r_CP2P_p = C_p2b @ r_CP2RM[-1] - r_P2RM

        v_CP2e_b = v_B2e + cross3(omega_b2e, r_CP2B_b)
        v_CP2e_p = v_P2e + cross3(omega_p2e, r_CP2P_p)

        v_W2CP_b = v_W2e[:-1] - v_CP2e_b
        v_W2CP_p = C_p2b @ v_W2e[-1] - v_CP2e_p

        # -------------------------------------------------------------------
        # Forces and moments of the wing in body frd
        try:
            f_b, g_b2B, ref = self.wing.resultant_force(
                delta_a,
                delta_bl,
                delta_br,
                v_W2CP_b,
                rho_air,
                g,
                r_B2LE,
                wmp,
                reference_solution,
            )
        except foil_aerodynamics.ConvergenceError:
            # Maybe it can't recover once Gamma is jacked?
            print("\nBonk! Retrying with the default reference solution")
            f_b, g_b2B, ref = self.wing.resultant_force(
                delta_a,
                delta_bl,
                delta_br,
                v_W2CP_b,
                rho_air,
                g,
                r_B2LE,
                wmp,
            )

        f_p, g_p2P = self.payload.resultant_force(  # In payload frd
            delta_w=delta_w,
            v_W2h=v_W2CP_p,
            rho_air=rho_air,
            g=C_p2b @ g,
            r_R2RM=pmp["r_P2RM"],
            mp=pmp,
        )

        # Moment at the connection point `RM` modeled as a spring+damper system
        omega_p2b = C_b2p @ omega_p2e - omega_b2e
        g_RM = Theta_p2b * self._kappa_RM + omega_p2b * self._kappa_RM_dot

        # ------------------------------------------------------------------
        # Build a system of equations by equating the time derivatives of the
        # translation and angular momentum (with respect to the Earth) of the
        # body and payload to the forces and moments on the body and payload.
        #
        # The four unknown vectors are the time derivatives of `v_RM2e`,
        # `omega_b2e` (in body frd), `omega_p2e` (in payload frd), and the
        # internal force on the risers, `F_RM` (in body frd).

        # fmt: off

        I3, Z3 = np.eye(3), np.zeros((3, 3))
        A1 = [m_b * I3, -m_b * crossmat(r_B2RM), Z3, I3]
        A2 = [Z3, J_b2B, Z3, -crossmat(r_B2RM)]
        A3 = [m_p * C_p2b, Z3, -m_p * crossmat(r_P2RM), -C_p2b]
        A4 = [Z3, Z3, J_p2P, crossmat(r_P2RM) @ C_p2b]
        A = np.block([A1, A2, A3, A4])
        B1 = (
            f_b
            - m_b * cross3(omega_b2e, v_RM2e)
            - m_b * cross3(omega_b2e, cross3(omega_b2e, r_B2RM))
        )
        B2 = (  # ref: Hughes Eq:13, p58
            g_b2B
            - g_RM
            - cross3(omega_b2e, J_b2B @ omega_b2e)
        )
        B3 = (
            f_p
            - m_p * C_p2b @ cross3(omega_b2e, v_RM2e)
            - m_p * cross3(omega_p2e, cross3(omega_p2e, r_P2RM))
        )
        B4 = g_p2P + C_p2b @ g_RM - cross3(omega_p2e, J_p2P @ omega_p2e)

        # fmt: on

        B = np.r_[B1, B2, B3, B4]
        x = np.linalg.solve(A, B)
        a_RM2e = x[:3]  # In frame F_b
        alpha_b2e = x[3:6]  # In frames F_b and F_e
        alpha_p2e = x[6:9]  # In frames F_p and F_e
        F_RM = x[9:]  # For debugging
        return a_RM2e, alpha_b2e, alpha_p2e, ref

class ParagliderSystemDynamics9c(ParagliderSystemDynamics9a):
    r"""
    A 9 degrees-of-freedom paraglider model, allowing rotation between the wing
    and the harness, with the connection modelled by spring-damper dynamics.

    Similar to ParagliderSystemDynamics9a, this version uses the riser midpoint
    `RM` as the reference point for both the body and the payload. Unlike
    ParagliderSystemDynamics9a, this model computes \dot{omega_p2b} instead of
    \dot{omega_p2e} and converts. Also, note that it computes everything in
    body frd and converts omega_p2e back to payload frd at the very end.

    FIXME: although 9a and 9b agree, this model produces slightly different
    answers, which might be worth looking into.

    Parameters
    ----------
    wing : ParagliderWing
    payload : ParagliderHarness
        The harness model includes the mass of the pilot.
    kappa_RM : array of float, shape (3,), optional
        Spring-damper coefficients for Theta_p2b (force as a linear function
        of angular displacement).
    kappa_RM_dot : array of float, shape (3,), optional
        Spring-damper coefficients for the derivative of Theta_p2b
    use_apparent_mass : bool, optional
        Whether to estimate the effects of apparent inertia. Default: True
    """

    def accelerations(
        self,
        v_RM2e,
        omega_b2e,
        omega_p2e,
        Theta_p2b,
        g,
        delta_a: float = 0,
        delta_bl: float = 0,
        delta_br: float = 0,
        delta_w: float = 0,
        rho_air: float = 1.225,
        v_W2e=(0, 0, 0),
        reference_solution: dict | None = None,
    ):
        r"""
        Compute the translational and angular accelerations about the center of mass.

        Parameters
        ----------
        v_RM2e : array of float, shape (3,) [m/s]
            Translational velocity of `RM` in body frd coordinates, where `RM`
            is the midpoint between the two riser connection points.
        omega_b2e : array of float, shape (3,) [rad/s]
            Angular velocity of the body, in body frd coordinates.
        omega_p2e : array of float, shape (3,) [rad/s]
            Angular velocity of the payload, in payload frd coordinates.
        Theta_p2b : array of float, shape (3,) [radians]
            The [phi, theta, gamma] of a yaw-pitch-roll sequence that encodes
            the relative orientation of the payload with respect to the body.
        g : array of float, shape (3,) [m/s^s]
            The gravity vector in body frd
        delta_a : float [percentage]
            The fraction of maximum accelerator
        delta_bl : float [percentage]
            The fraction of maximum left brake
        delta_br : float [percentage]
            The fraction of maximum right brake
        delta_w : float [percentage]
            The fraction of weight shift, from -1 (left) to +1 (right)
        rho_air : float [kg/m^3], optional
            Air density
        v_W2e : ndarray of float, shape (3,) or (K,3) [m/s], optional
            The wind relative to the earth, in body frd coordinates. If it is
            a (3,) array then the wind is uniform at every control point. If
            it is a (K,3) array then it is the vectors for each control point.
        reference_solution : dictionary, optional
            FIXME: docstring. See `Phillips.__call__`

        Returns
        -------
        a_RM2e : array of float, shape (3,) [m/s^2]
            Translational acceleration of `RM` in body frd coordinates.
        alpha_b2e : array of float, shape (3,) [rad/s^2]
            Angular acceleration of the body with respect to Earth as the time
            derivative of angular velocity taken with respect to the body
            frame, expressed in body frd coordinates
            :math:`\left( ^b \dot{\omega}_{b/e}^b \right)`.
        alpha_p2e : array of float, shape (3,) [rad/s^2]
            Angular acceleration of the payload with respect to Earth as the
            time derivative of angular velocity taken with respect to the
            payload frame, expressed in payload frd coordinates
            :math:`\left( ^p \dot{\omega}_{p/e}^p \right)`.
        solution : dictionary
            FIXME: docstring. See `Phillips.__call__`
        """
        r_CP2RM = self.r_CP2RM(Theta_p2b, delta_a, delta_w)
        v_W2e = np.asfarray(v_W2e)
        if v_W2e.ndim > 1 and v_W2e.shape[0] != r_CP2RM.shape[0]:
            raise ValueError(
                "Number of distinct wind vectors in v_W2e do not match the "
                "number of control points",
            )
        v_W2e = np.broadcast_to(v_W2e, r_CP2RM.shape)
        v_RM2e = np.asfarray(v_RM2e)
        if v_RM2e.shape != (3,):
            raise ValueError("v_RM2e.shape != (3,)")

        C_p2b = orientation.euler_to_dcm(Theta_p2b)
        C_b2p = C_p2b.T
        omega_p2e = C_b2p @ omega_p2e  # Dynamics9a uses `p`, this model uses `b`

        # -------------------------------------------------------------------
        # Compute the inertia properties of the body and payload.
        r_RM2LE = self.wing.r_RM2LE(delta_a)
        wmp = self.wing.mass_properties(rho_air, r_RM2LE)  # R = RM
        m_b = wmp["m_s"] + wmp["m_air"]
        J_b2RM = wmp["J_b2R"]
        r_B2RM = wmp["r_B2R"]

        pmp = self.payload.mass_properties(delta_w, [0, 0, 0])  # R = RM
        m_p = pmp["m_p"]
        J_p2RM = C_b2p @ pmp["J_p2R"] @ C_p2b  # In body frd
        r_P2RM = C_b2p @ pmp["r_P2R"]  # In body frd

        # -------------------------------------------------------------------
        # Compute the relative wind vectors for each control point. All vectors
        # are in body frd.
        omega_p2b = omega_p2e - omega_b2e
        omega_b2p = -omega_p2b
        v_B2e = v_RM2e + cross3(omega_b2e, r_B2RM)
        v_P2e = v_RM2e + cross3(omega_p2e, r_P2RM)

        # FIXME: "magic" indexing established by `self.r_CP2RM`
        r_CP2RM_b = r_CP2RM[:-1]
        r_CP2RM_p = r_CP2RM[-1]

        v_CP2e_b = v_B2e + cross3(omega_b2e, r_CP2RM_b - r_B2RM)
        v_CP2e_p = v_P2e + cross3(omega_p2e, r_CP2RM_p - r_P2RM)

        v_W2CP_b = v_W2e[:-1] - v_CP2e_b
        v_W2CP_p = v_W2e[-1] - v_CP2e_p

        # -------------------------------------------------------------------
        # Forces and moments of the wing in body frd
        try:
            f_b, g_b2RM, ref = self.wing.resultant_force(
                delta_a,
                delta_bl,
                delta_br,
                v_W2CP_b,
                rho_air,
                g,
                r_RM2LE,
                wmp,
                reference_solution,
            )
        except foil_aerodynamics.ConvergenceError:
            # Maybe it can't recover once Gamma is jacked?
            print("\nBonk! Retrying with the default reference solution")
            f_b, g_b2RM, ref = self.wing.resultant_force(
                delta_a,
                delta_bl,
                delta_br,
                v_W2CP_b,
                rho_air,
                g,
                r_RM2LE,
                wmp,
            )

        f_p, g_p2RM = self.payload.resultant_force(  # In payload frd
            delta_w=delta_w,
            v_W2h=C_p2b @ v_W2CP_p,
            rho_air=rho_air,
            g=C_p2b @ g,
            r_R2RM=[0, 0, 0],
            mp=pmp,
        )

        # Moment at the connection point `RM` modeled as a spring+damper system
        g_RM = Theta_p2b * self._kappa_RM + omega_p2b * self._kappa_RM_dot

        # ------------------------------------------------------------------
        # Build a system of equations by equating the time derivatives of the
        # translation and angular momentum (with respect to the Earth) of the
        # body and payload to the forces and moments on the body and payload.
        #
        # The four unknown vectors are the time derivatives of `v_RM2e^b`,
        # `omega_b2e^b`, `omega_p2e^p`, and the internal force on the risers,
        # `F_RM^b`. All derivatives are from the body frame.

        # fmt: off

        # Real mass momentums
        p_b2e = m_b * v_B2e  # Linear
        p_p2e = m_p * v_P2e
        h_b2RM = m_b * cross3(r_B2RM, v_RM2e) + J_b2RM @ omega_b2e  # Angular
        h_p2RM = m_p * cross3(r_P2RM, v_RM2e) + J_p2RM @ omega_p2e

        I3, Z3 = np.eye(3), np.zeros((3, 3))
        A1 = [m_b * I3, -m_b * crossmat(r_B2RM), Z3, I3]
        A2 = [m_b * crossmat(r_B2RM), J_b2RM, Z3, Z3]
        A3 = [m_p * I3, -m_p * crossmat(r_P2RM), -m_p * crossmat(r_P2RM), -I3]
        A4 = [m_p * crossmat(r_P2RM), J_p2RM, J_p2RM, Z3]
        A = np.block([A1, A2, A3, A4])

        B1 = f_b - cross3(omega_b2e, p_b2e)
        B2 = (  # ref: Hughes Eq:13, p58
            g_b2RM
            - g_RM
            - cross3(v_RM2e, p_b2e)
            - cross3(omega_b2e, h_b2RM)
        )
        B3 = (
            C_b2p @ f_p
            - m_p * cross3(omega_b2p, v_RM2e)
            - m_p * cross3(cross3(omega_b2p, omega_b2e), r_P2RM)
            # - m_p * cross3(cross3(omega_b2p, omega_p2e), r_P2RM)  # equivalent
            - cross3(omega_p2e, p_p2e)
        )
        B4 = (
            C_b2p @ g_p2RM
            + g_RM
            - cross3(v_RM2e, p_p2e)
            - m_p * cross3(r_P2RM, cross3(omega_b2p, v_RM2e))
            - cross3(omega_b2p, J_p2RM @ omega_p2e)
            - cross3(omega_p2e, h_p2RM)
        )

        if self.use_apparent_mass:
            amp = self.wing.apparent_mass_properties(
                rho_air,
                r_RM2LE,
                v_RM2e,
                omega_b2e,
            )
            A[:6, :6] += amp["A_a2R"]  # Incorporate the apparent inertia
            B1 += (  # Apparent inertial force (Barrows Eq:61)
                -cross3(omega_b2e, amp["p_a2e"])
            )
            B2 += (  # Apparent inertial moment (Barrows Eq:64)
                -cross3(v_RM2e, amp["p_a2e"])
                - cross3(omega_b2e, amp["h_a2R"])
                + cross3(v_RM2e, amp["M_a"] @ v_RM2e)  # Remove the steady-state term
            )

        # fmt: on

        B = np.r_[B1, B2, B3, B4]
        x = np.linalg.solve(A, B)
        a_RM2e = x[:3]  # In frame F_b
        alpha_b2e = x[3:6]  # In frames F_b and F_e
        alpha_p2b = x[6:9]  # In frames F_b and F_p
        F_RM = x[9:]  # For debugging
        alpha_p2e = alpha_p2b + alpha_b2e + cross3(omega_b2e, omega_p2b)
        alpha_p2e = C_p2b @ alpha_p2e  # In frames F_p and F_e
        return a_RM2e, alpha_b2e, alpha_p2e, ref


#----------------------------------------------------------------------------#
# Imports
#----------------------------------------------------------------------------#

from flask import Flask, render_template, request, make_response, jsonify
import numpy as np
import keras.backend as K
import keras
from keras.layers import Dense
from keras.models import Sequential
import os
from keras.models import load_model
import json
import logging
from logging import Formatter, FileHandler
from forms import *
import os

#----------------------------------------------------------------------------#
# App Config.
#----------------------------------------------------------------------------#

app = Flask(__name__)
app.config.from_object('config')
#db = SQLAlchemy(app)

# Automatically tear down SQLAlchemy.
'''
@app.teardown_request
def shutdown_session(exception=None):
    db_session.remove()
'''

# Login required decorator.
'''
def login_required(test):
    @wraps(test)
    def wrap(*args, **kwargs):
        if 'logged_in' in session:
            return test(*args, **kwargs)
        else:
            flash('You need to login first.')
            return redirect(url_for('login'))
    return wrap
'''
#----------------------------------------------------------------------------#
# Controllers.
#----------------------------------------------------------------------------#

@app.route('/')
def home():
    return render_template('pages/placeholder.home.html')

@app.route('/practise')
def practise():
    return render_template('pages/practise.html')

@app.route('/revise')
def revise():
    return render_template('pages/revise.html')

# predict
@app.route('/predict', methods=['POST'])
def predict():
    K.clear_session()
    model = load_model('space_rep_lr.h5')
    data = json.loads(request.form["data"])
    X = []
    for d in data:
        X.append(d["data"])
    preds = model.predict(np.array(X))
    resp = []
    for i in range(len(preds)):
        resp.append({
            "id": data[i]["id"],
            "pred": int(preds[i][0]*100)})
    return make_response(jsonify(resp)), 200

# Error handlers.

@app.errorhandler(500)
def internal_error(error):
    # db_session.rollback()
    return render_template('errors/500.html'), 500

@app.errorhandler(404)
def not_found_error(error):
    return render_template('errors/404.html'), 404

if not app.debug:
    file_handler = FileHandler('error.log')
    file_handler.setFormatter(
        Formatter(
            '%(asctime)s %(levelname)s: %(message)s [in %(pathname)s:%(lineno)d]')
    )
    app.logger.setLevel(logging.INFO)
    file_handler.setLevel(logging.INFO)
    app.logger.addHandler(file_handler)
    app.logger.info('errors')

#----------------------------------------------------------------------------#
# Launch.
#----------------------------------------------------------------------------#

# Default port:
if __name__ == '__main__':
    app.run()

# Or specify port manually:
'''
if __name__ == '__main__':
    port = int(os.environ.get('PORT', 5000))
    app.run(host='0.0.0.0', port=port)
'''


import numpy as np

gauss_pulse_len = 20  # ns
const_pulse_len = 100
Amp = 0.2  # Pulse Amplitude
gauss_arg = np.linspace(-3, 3, gauss_pulse_len)
gauss_wf = np.exp(-(gauss_arg ** 2) / 2)
gauss_wf = Amp * gauss_wf / np.max(gauss_wf)
readout_time = 100
# omega_10 = 4.958e9
# omega_d = omega_10 / 3

pump_len = 100
prep_len = 100
ramp_len = 20
ramp_max = 0.1

v_prep_top = 0.2
v_prep_bot = -0.2
v_readout_top = 0.2
v_readout_bot = -0.2
v_evolve_top = 0.1
v_evolve_bot = -0.1

ramp_down_prep = np.linspace(v_prep_top, v_evolve_bot, ramp_len).tolist()
ramp_up_prep = np.linspace(v_prep_bot, v_evolve_top, ramp_len).tolist()
ramp_up_readout = np.linspace(v_evolve_bot, v_readout_top, ramp_len).tolist()
ramp_down_readout = np.linspace(v_evolve_top, v_readout_bot, ramp_len).tolist()

config = {
    "version": 1,
    "controllers": {
        "con1": {
            "type": "opx1",
            "analog_outputs": {
                1: {"offset": +0.0},  # Vr
                2: {"offset": +0.0},  # Vl
                3: {"offset": +0.0},  # RF-QPC-Out
            },
            "analog_inputs": {
                1: {"offset": +0.0},  # RF-QPC-in
            },
        },
    },
    "elements": {
        "DET": {
            "singleInput": {"port": ("con1", 1)},
            "intermediate_frequency": 0e6,
            "operations": {
                "evolve": "constPulse",  # to a pulse
            },
        },
        "D1": {
            "singleInput": {"port": ("con1", 1)},
            "intermediate_frequency": 0e6,
            # 'hold_offset':{'duration': 100},
            "operations": {
                "const": "constPulse",
                "prep": "prepPosPulse",  # to a pulse
                "pump": "constPulse",
                # "prep_neg": "constNegPulse",  # to a pulse
                "evolve": "evolveD1Pulse",  # to a pulse
                "readout": "readoutPosPulse",  # to a pulse
                "ramp_down_evolve": "rampDownD1Pulse",
                "ramp_up_read": "rampUpD1Pulse",
            },
        },
        "D2": {
            "singleInput": {"port": ("con1", 2)},
            "intermediate_frequency": 0e6,
            "operations": {
                # "prep_pos": "constPosPulse",  # to a pulse
                "prep": "prepNegPulse",  # to a pulse
                "evolve": "evolveD2Pulse",  # to a pulse
                "pump": "constPulse",
                "readout": "readoutNegPulse",  # to a pulse
                "ramp_down_read": "rampDownD2Pulse",
                "ramp_up_evolve": "rampUpD2Pulse",
            },
        },
        "D1_RF": {
            "singleInput": {"port": ("con1", 1)},
            "intermediate_frequency": 0e6,
            "operations": {
                "const": "constPulse",
                # "prep": "prepPulse",  # to a pulse
                # "readout": "readoutPulse",
                # "s-pump": "pumpPulse",  # to a pulse
                # "pi_pulse": "pi_pulse_in"
            },
        },
        "D2_RF": {
            "singleInput": {"port": ("con1", 2)},
            "intermediate_frequency": 0e6,
            "operations": {
                "const": "constPulse",
                # "prep": "prepPulse",  # to a pulse
                # "readout": "readoutPulse",
                # "s-pump": "pumpPulse",  # to a pulse
                # "pi_pulse": "pi_pulse_in"
            },
        },
        "RF-QPC": {
            "singleInput": {"port": ("con1", 3)},
            "intermediate_frequency": 200e6,
            "operations": {
                "measure": "measure_pulse",
            },
            "time_of_flight": 180,  # Measurement parameters
            "smearing": 0,
            "outputs": {"out1": ("con1", 1)},
        },
    },
    "pulses": {
        "measure_pulse": {  # Readout pulse
            "operation": "measurement",
            "length": readout_time,
            "waveforms": {
                "single": "const_wf",
            },
            "integration_weights": {
                "integW1": "integW1",
                "integW2": "integW2",
            },
        },
        "constPulse": {
            "operation": "control",
            "length": const_pulse_len,  # in ns
            "waveforms": {"single": "const_wf"},
        },
        "rampDownD1Pulse": {
            "operation": "control",
            "length": ramp_len,  # in ns
            "waveforms": {"single": "ramp_down_d1_wf"},
        },
        "rampUpD1Pulse": {
            "operation": "control",
            "length": ramp_len,  # in ns
            "waveforms": {"single": "ramp_up_d1_wf"},
        },
        "rampDownD2Pulse": {
            "operation": "control",
            "length": ramp_len,  # in ns
            "waveforms": {"single": "ramp_down_d2_wf"},
        },
        "rampUpD2Pulse": {
            "operation": "control",
            "length": ramp_len,  # in ns
            "waveforms": {"single": "ramp_up_d2_wf"},
        },
        "readoutPosPulse": {
            "operation": "control",
            "length": readout_time,  # in ns
            "waveforms": {"single": "readout_pos_wf"},
        },
        "readoutNegPulse": {
            "operation": "control",
            "length": readout_time,  # in ns
            "waveforms": {"single": "readout_neg_wf"},
        },
        "evolveD2Pulse": {
            "operation": "control",
            "length": const_pulse_len,  # in ns
            "waveforms": {"single": "evolve_pos_wf"},
        },
        "evolveD1Pulse": {
            "operation": "control",
            "length": const_pulse_len,  # in ns
            "waveforms": {"single": "evolve_neg_wf"},
        },
        "prepPosPulse": {
            "operation": "control",
            "length": const_pulse_len,  # in ns
            "waveforms": {"single": "prep_pos_wf"},
        },
        "prepNegPulse": {
            "operation": "control",
            "length": const_pulse_len,  # in ns
            "waveforms": {"single": "const_neg_wf"},
        },
        "prepPulse": {
            "operation": "control",
            "length": prep_len,
            "waveforms": {"I": "const_wf", "Q": "const_wf"},
        },
        "readoutPulse": {
            "operation": "control",
            "length": readout_time,
            "waveforms": {"I": "const_wf", "Q": "const_wf"},
        },
        "pumpPulse": {
            "operation": "control",
            "length": pump_len,
            "waveforms": {"I": "const_wf", "Q": "const_wf"},
        },
        "pi_pulse_in": {  # Assumed to be calibrated
            "operation": "control",
            "length": const_pulse_len,
            "waveforms": {"I": "const_wf", "Q": "zero_wf"},
        },
        "const_pulse_in": {  # Assumed to be calibrated
            "operation": "control",
            "length": const_pulse_len,
            "waveforms": {"I": "const_wf", "Q": "zero_wf"},
        },
    },
    "waveforms": {
        "const_wf": {"type": "constant", "sample": 0.1},
        "prep_pos_wf": {"type": "constant", "sample": v_prep_top},
        "prep_neg_wf": {"type": "constant", "sample": v_prep_bot},
        "evolve_pos_wf": {"type": "constant", "sample": v_evolve_top},
        "evolve_neg_wf": {"type": "constant", "sample": v_evolve_bot},
        "readout_pos_wf": {"type": "constant", "sample": v_readout_top},
        "readout_neg_wf": {"type": "constant", "sample": v_readout_bot},
        "const_neg_wf": {"type": "constant", "sample": -0.2},
        "ev_d1_wf": {"type": "constant", "sample": -0.2},
        "ev_d2_wf": {"type": "constant", "sample": 0.1},
        "zero_wf": {"type": "constant", "sample": 0.0},
        "ramp_down_d1_wf": {"type": "arbitrary", "samples": ramp_down_prep},
        "ramp_up_d1_wf": {"type": "arbitrary", "samples": ramp_up_readout},
        "ramp_down_d2_wf": {"type": "arbitrary", "samples": ramp_down_readout},
        "ramp_up_d2_wf": {"type": "arbitrary", "samples": ramp_up_prep},
        "gauss_wf": {"type": "arbitrary", "samples": gauss_wf.tolist()},
        "exc_wf": {"type": "constant", "sample": 0.479},
        "ramp_wf": {
            "type": "arbitrary",
            "samples": np.linspace(0, ramp_max, pump_len).tolist(),
        },
    },
    "integration_weights": {  # Define integration weights for measurement demodulation
        "integW1": {"cosine": [4.0] * readout_time, "sine": [0.0] * readout_time},
        "integW2": {"cosine": [0.0] * readout_time, "sine": [4.0] * readout_time},
    },
    "mixers": {  # Potential corrections to be brought related to the IQ mixing scheme
        "mixer_res": [
            {
                "intermediate_frequency": 50e6,  # 6.15e9,
                "lo_frequency": 6.00e9,
                "correction": [1.0, 0.0, 0.0, 1.0],
            }
        ],
        "mixer_qubit": [
            {
                "intermediate_frequency": 0,  # ω_d,
                "lo_frequency": 5.10e9,
                "correction": [1.0, 0.0, 0.0, 1.0],
            }
        ],
    },
}


# Copyright 2017 Bloomberg Finance L.P.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import tensorflow as tf
import os
import json
import subprocess
from scipy.misc import imresize, imsave
import argparse
from tensorbox.train_obj_model import build_forward
from tensorbox.utils import googlenet_load
from tensorbox.utils.annolist import AnnotationLib as al
from tensorbox.utils.train_utils import add_rectangles, rescale_boxes
import cv2
import numpy as np
import pandas as pd
from sklearn import metrics
from PIL import Image
from IPython.display import Image as Image2
import time
import scatteract_logger

def main(model_dir, image_dir, true_idl, iteration, iou_threshold, conf_threshold):

    hypes_file = '{}/hypes.json'.format(model_dir)

    with open(hypes_file, 'r') as f:
        H = json.load(f)

    model_name = model_dir.split("/")[1]
    pred_idl = './output/%s_%d_val_%s.idl' % (model_name, iteration, os.path.basename(hypes_file).replace('.json', ''))
    true_annos = al.parse(true_idl)

    tf.reset_default_graph()
    googlenet = googlenet_load.init(H)
    x_in = tf.placeholder(tf.float32, name='x_in', shape=[H['image_height'], H['image_width'], 3])

    if H['use_rezoom']:
        pred_boxes, pred_logits, pred_confidences, pred_confs_deltas, pred_boxes_deltas = build_forward(H, tf.expand_dims(x_in, 0), googlenet, 'test', reuse=None)
        grid_area = H['grid_height'] * H['grid_width']
        pred_confidences = tf.reshape(tf.nn.softmax(tf.reshape(pred_confs_deltas, [grid_area * H['rnn_len'], 2])), [grid_area, H['rnn_len'], 2])
        if H['reregress']:
            pred_boxes = pred_boxes + pred_boxes_deltas
    else:
        pred_boxes, pred_logits, pred_confidences = build_forward(H, tf.expand_dims(x_in, 0), googlenet, 'test', reuse=None)

    saver = tf.train.Saver()
    with tf.Session() as sess:
        sess.run(tf.initialize_all_variables())
        #print tf.train.get_checkpoint_state(model_dir)
        saver.restore(sess, '{}/save.ckpt-{}'.format(model_dir,iteration))

        annolist = al.AnnoList()
        t = time.time()
        for i in range(len(true_annos)):
            true_anno = true_annos[i]

            img = Image.open(image_dir + "/" + true_anno.imageName)
            bg = Image.new("RGB", img.size, (255,255,255))
            bg.paste(img,img)
            img = np.array(bg)
            img_orig = np.copy(img)

            if img.shape[0] != H["image_height"] or img.shape[1] != H["image_width"]:
                true_anno = rescale_boxes(img.shape, true_anno, H["image_height"], H["image_width"])
                img = imresize(img, (H["image_height"], H["image_width"]), interp='cubic')
            feed = {x_in: img}
            (np_pred_boxes, np_pred_confidences) = sess.run([pred_boxes, pred_confidences], feed_dict=feed)
            pred_anno = al.Annotation()
            pred_anno.imageName = true_anno.imageName
            new_img, rects = add_rectangles(H, [img], np_pred_confidences, np_pred_boxes,
                                        use_stitching=True, rnn_len=H['rnn_len'], min_conf=conf_threshold)
            new_img_true = np.copy(img)
            new_img_pred = np.copy(img)

            for rect_true in true_anno.rects:
                cv2.rectangle(new_img_true,(int(rect_true.x1),int(rect_true.y1)),
                              (int(rect_true.x2),int(rect_true.y2)),
                              (0,255,0),2)

            for rect_pred in rects:
                cv2.rectangle(new_img_pred,(int(rect_pred.x1),int(rect_pred.y1)),
                              (int(rect_pred.x2),int(rect_pred.y2)),
                              (0,0,255),2)

            pred_anno.rects = rects
            pred_anno = rescale_boxes(img.shape, pred_anno, img_orig.shape[0], img_orig.shape[1])
            annolist.append(pred_anno)

            if i % 10 == 0 and i < 400:
                imsave("{}/".format(model_dir)+pred_anno.imageName.split('/')[-1][:-4]+'_pred.bmp',new_img_pred)
                imsave("{}/".format(model_dir)+pred_anno.imageName.split('/')[-1][:-4]+'_true.bmp',new_img_true)
            if (i+1) % 200 == 0 or i==0:
                mylogger.debug("Number of images analyzed: {}".format(i+1))
        avg_time = (time.time() - t) / (i + 1)
        mylogger.debug('%f images/sec' % (1. / avg_time))

    annolist.save(pred_idl)
    rpc_cmd = './tensorbox/utils/annolist/doRPC.py  --minOverlap %f %s %s' % (iou_threshold, true_idl, pred_idl)
    mylogger.debug('$ %s' % rpc_cmd)
    rpc_output = subprocess.check_output(rpc_cmd, shell=True)
    mylogger.debug(rpc_output)
    txt_file = [line for line in rpc_output.split('\n') if line.strip()][-1]
    output_png = 'output/{}_{}_results.png'.format(model_name, iteration)
    plot_cmd = './tensorbox/utils/annolist/plotSimple.py %s --output %s' % (txt_file, output_png)
    mylogger.debug('$ %s' % plot_cmd)
    plot_output = subprocess.check_output(plot_cmd, shell=True)
    Image2(filename=output_png)

    df = pd.read_csv("output/rpc-{}_{}_val_hypes_overlap{}.txt".format(model_name,iteration,iou_threshold), sep=" ", names=['precision','recall','fpii','score','accuracy'])
    auc = metrics.auc(df['recall'],df['precision'])
    mylogger.info("Average Precision: {}".format(auc))

if __name__ == "__main__":

    """
    Example of command line usage:

    python test_obj_model.py --model_dir output/lstm_rezoom_plot_labels_2017_04_11_01.14 --image_dir data/plot_test --true_idl data/plot_test/labels.idl --iteration 125000
    """

    mylogger = scatteract_logger.get_logger()
    parser = argparse.ArgumentParser()

    parser.add_argument('--model_dir', help='Model directory (string)', required=True)
    parser.add_argument('--image_dir', help='Image directory (string)', required=True)
    parser.add_argument('--true_idl', help='True idl file (string)', required=True)
    parser.add_argument('--iteration', help='Iteration number (int)', required=True)
    parser.add_argument('--iou_threshold', help='IOU threshold (float)', default=0.5, required=False)
    parser.add_argument('--conf_threshold', help='Confidence threshold (float)', default=0.3, required=False)
    args = vars(parser.parse_args())

    main(args['model_dir'], args['image_dir'], args['true_idl'], int(args['iteration']), float(args['iou_threshold']), float(args['conf_threshold']))


import cv2
import glob
import random
import sys
import os
import numpy as np

emotions = ["anger", "happy", "sadness"]
fishface = cv2.createFisherFaceRecognizer()
data = {}

def run_recognizer():
    fishface.load("trained_emoclassifier.xml")
    prediction_data1 = []
    files = glob.glob("face_cut\\try\\*")
    prediction = files[:]
    img=cv2.imread(prediction[0])
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    prediction_data1.append(gray)
    pred1, conf1 = fishface.predict(prediction_data1[0])
    return pred1

var = sys.argv[1]
sys.argv = ['face.py',var]
execfile("face.py")
correct = run_recognizer()
print emotions[correct]


# Copyright 2019 Xilinx Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# PART OF THIS FILE AT ALL TIMES.

import argparse
import pdb
import os
import shutil
import time
import math
import copy
import sys

import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim
import torch.utils.data
import torchvision.transforms as transforms
import torchvision.datasets as datasets
from torch.autograd import Variable

import numpy as np

import sys
sys.path.append('../models/')
sys.path.append('../utils/')
from load_imglist import ImageList
import face_model
import evaluate
import image_transforms
import Log
from torchsummary import summary

parser = argparse.ArgumentParser(description='PyTorch Light CNN Training')
parser.add_argument('--arch', '-a', metavar='ARCH', default='resnet20_am',
                    help='model architecture: (default: vgg19)')
parser.add_argument('--cuda', '-c', default=True)
parser.add_argument('-j', '--workers', default=4, type=int, metavar='N',
                    help='number of data loading workers (default: 16)')
parser.add_argument('--epoches', default=30, type=int, metavar='N',
                    help='number of total epoches to run')
parser.add_argument('--start-epoch', default=0, type=int, metavar='N',
                    help='manual epoch number (useful on restarts)')
parser.add_argument('-b', '--batch-size', default=1024, type=int,
                    metavar='N', help='mini-batch size (default: 128)')
parser.add_argument('--test-batch-size', default=128, type=int,
                    metavar='N', help='test mini-batch size (default: 500)')
parser.add_argument('--lr', '--learning-rate', default=0.0001, type=float,
                    metavar='LR', help='initial learning rate')
parser.add_argument('--lr-decay', '--learning-rate-decay', default=0.1, type=float,
                    metavar='LRD', help='learning rate decay rate')
parser.add_argument('--lr-list', nargs='+', default=[13,20,25], type=int)
parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
                    help='momentum')
parser.add_argument('--weight-decay', '--wd', default=1e-4, type=float,
                    metavar='W', help='weight decay (default: 1e-4)')
parser.add_argument('--print-freq', '-p', default=50, type=int,
                    metavar='N', help='print frequency (default: 100)')
parser.add_argument('--model', default='', type=str, metavar='Model',
                    help='model type: LightCNN-9, LightCNN-29, LightCNN-29v2')
parser.add_argument('--pretrained', default='', type=str, metavar='PATH',
                    help='path to pretrained checkpoint (default: none)')
parser.add_argument('--resume', default='', type=str, metavar='PATH',
                    help='path to latest checkpoint (default: none)')
parser.add_argument('--finetune', default=True, action='store_true', help='use tensorboard.')
parser.add_argument('--root_path', default='../../data/train/ms_glint',
                    type=str, metavar='PATH', help='path to root path of images (default: none)')
parser.add_argument('--train_list', default='../../data/train/ms_glint/list/msra_celebrity.txt', type=str, metavar='PATH',
                    help='path to training list (default: none)')
parser.add_argument('--testset_images', default='../../data/test/lfw/images', type=str, metavar='PATH',
                    help='path to testset (default: none)')
parser.add_argument('--testset_pairs', default='../../data/test/lfw/pairs.txt', type=str, metavar='PATH',
                    help='path to testset (default: none)')
parser.add_argument('--testset_list', default='../../data/test/lfw/lfw.txt', type=str, metavar='PATH',
                    help='path to testset (default: none)')
parser.add_argument('--save-dir', default='save', type=str, metavar='PATH',
                    help='path to save checkpoint (default: none)')
parser.add_argument('--num_classes', default=180855, type=int,
                    metavar='N', help='number of classes')
parser.add_argument('--save-prec', default=90, type=float, metavar='S',
                    help='save prec')
parser.add_argument('--test-epoch', default=0, type=int, metavar='S',
                    help='test epoch')
parser.add_argument('--fea-dim', default=512, type=int)
parser.add_argument('--prefix', default='temp', type=str, metavar='PREFIX',
                    help='checkpoint prefix (default: none)')
parser.add_argument('--summary', action='store_true', help='use tensorboard.')
parser.add_argument('--lamb', default=0.001, type=float, help='lambda for lamsoftmax')
parser.add_argument('--scale', default=80, type=float, help='scale for amsoftmax')
parser.add_argument('--margin', default=0.4, type=float, help='margin for amsoftmax or lamsoftmax')
parser.add_argument('--eval', default='lfw', type=str)
parser.add_argument('--auggamma', default='1.,1.', type=str)
parser.add_argument('--resize', default='', type=str)
parser.add_argument('--nonlinear', default='relu', type=str)
parser.add_argument('--log', default='log', type=str)

args = parser.parse_args()

Log.Log(args.log, 'w+', 1) # set log file

print("=================FLAGS==================")
for k, v in args.__dict__.items():
    print('{}: {}'.format(k, v))
print("========================================")

# tensorboard
def to_np(x):
    return x.data.cpu().numpy()

if args.summary:
    from logger import Logger
    log_dir = os.path.join('./runs/', args.prefix)
    if os.path.exists(log_dir):
        import shutil
        shutil.rmtree(log_dir)
    logger = Logger(log_dir)

start = time.time()

def main():
    global args, best_prec1, logger
    best_prec1 = 0.

    # create resnet for face recognition
    if 'resnet' in args.arch:
        truearch = args.arch.replace('_am', '').replace('_arc', '')
        if 'am' in args.arch or 'arc' in args.arch:
            ring = {'s': args.scale}
            model = face_model.__dict__[truearch](args.num_classes, wn=True, fn=True, nonlinear=args.nonlinear, fea_dim=args.fea_dim, ring=ring)
        else:
            model = face_modle.__dict__[truearch](args.num_classes, wn=False, fn=False,nonlinear=args.nonlinear, fea_dim=args.fea_dim)
    size = (3, 112, 96)
    summary(copy.deepcopy(model.cuda()), size)

    model = torch.nn.DataParallel(model).cuda()
    if args.finetune:
        checkpoint = args.pretrained
        if os.path.isfile(checkpoint):
            checkpoint = torch.load(checkpoint)
            # load part of state_dict
            pretrained_dict = checkpoint['state_dict']
            model_dict = model.state_dict()
            new_dict = {}
            for k, v in pretrained_dict.items():
                if not k in model_dict:
                    continue
                else:
                    new_dict[k] = v
            # pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
            model_dict.update(new_dict)
            model.load_state_dict(model_dict)
        else:
            print("wrong checkpoint path", checkpoint)

    # optionally resume from a checkpoint
    if args.resume:
        if os.path.isfile(args.resume):
            print("=> loading checkpoint '{}'".format(args.resume))
            checkpoint = torch.load(args.resume)
            args.start_epoch = checkpoint['epoch']
            model.load_state_dict(checkpoint['state_dict'])
            print("=> loaded checkpoint '{}' (epoch {})"
                  .format(args.resume, checkpoint['epoch']))
        else:
            print("=> no checkpoint found at '{}'".format(args.resume))

    optimizer = torch.optim.SGD(model.parameters(), args.lr,
                                momentum=args.momentum,
                                weight_decay=args.weight_decay)

    cudnn.benchmark = True

    normalize = transforms.Normalize(mean=[0.5, 0.5, 0.5],
            std=[0.5, 0.5, 0.5])
    grayscale = transforms.Grayscale(num_output_channels=1)

    auggamma = [float(x) for x in args.auggamma.split(',')]
    randomgamma = image_transforms.RandomGamma(auggamma[0], auggamma[1])
    gammabalance = image_transforms.GammaBalance()
    transform_ops = [
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            normalize,
    ]
    if args.resize != '':
        size = (int(x) for x in args.size.split(','))
        resize = transforms.Resize(size=size)
        transform_ops.insert(1, resize)
    transform=transforms.Compose(transform_ops)
    #load image
    train_loader = torch.utils.data.DataLoader(
        ImageList(root=args.root_path, fileList=args.train_list, transform=transform),
        batch_size=args.batch_size, shuffle=True,
        num_workers=args.workers, pin_memory=True)

    # define loss function and optimizer
    criterion = nn.CrossEntropyLoss()
    if args.cuda:
        criterion.cuda()

    eval_best_prec = 0.
    best_epoch = 0

    for epoch in range(args.start_epoch, args.epoches):

        if epoch in args.lr_list:
            optimizer.param_groups[0]['lr'] *= args.lr_decay
        # adjust_learning_rate(optimizer, epoch)
        # train for one epoch
        prec1 = train(train_loader, model, criterion, optimizer, epoch)

        modelname = os.path.join(args.save_dir, 'checkpoint_{}_{}.tar'.format(args.prefix, epoch+1))

        # remember best prec@1 and save checkpoint
        if epoch >= args.test_epoch:
            best_prec1 = prec1
            suffix = '_{}.bin'.format(args.arch)

            eval_prec = evaluate.eval_lfw(model, args.testset_images, args.testset_list, args.testset_pairs, batch_size=args.test_batch_size, fea_dim=args.fea_dim)
            if eval_prec > eval_best_prec:
                eval_best_prec = eval_prec
                best_epoch = epoch + 1

                save_checkpoint({
                    'epoch': best_epoch,
                    'lr': optimizer.param_groups[0]['lr'],
                    'arch': args.arch,
                    'state_dict': model.state_dict(),
                    'best_prec1': best_prec1,
                    }, filename=modelname)
                print('Eval {}\tprec: {}\tbest_prec: {}'.format(args.eval, eval_prec, eval_best_prec))

        best_prec1 = max(prec1, best_prec1)
        print('{} best prec: {}, method: {}, best epoch: {}'.format(args.eval, eval_best_prec, args.prefix, best_epoch))

def train(train_loader, model, criterion, optimizer, epoch):
    losses     = AverageMeter()
    mean_cos   = AverageMeter()
    top1       = AverageMeter()
    top5       = AverageMeter()

    model.train()

    for i, (input, target) in enumerate(train_loader):
        if args.cuda:
            input      = input.cuda()
            target     = target.cuda()
        input_var  = torch.autograd.Variable(input)
        target_var = torch.autograd.Variable(target)

        # compute output
        ringloss = torch.autograd.Variable(torch.Tensor([0]), requires_grad=True)
        feature, output = model(input_var)

        if '_am' in args.arch:
            cos = am_margin(output, target_var, args.margin)
            output = model.module.scale(output)
            mean_cos.update(cos, 1)
        loss = criterion(output, target_var)

        # measure accuracy and record loss
        prec1, prec5 = accuracy(output.data, target, topk=(1,5))
        losses.update(loss.data, input.size(0)) # loss.data[0]
        top1.update(prec1, input.size(0)) # prec1[0]
        top5.update(prec5, input.size(0)) # prec5[0]

        # compute gradient and do SGD step
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if i % args.print_freq == 0:
            lr = optimizer.param_groups[0]['lr']
            print('Epoch: [{0}][{1}/{2}] '
                    'Algo {algo}  '
                    'Time {t:.1f} '
                    'Lr {lr:.5f} '
                  'Loss {loss.val:.1f} ({loss.avg:.1f}) '
                  'Cos {mean_cos.val:.2f} '
                  'Prec@1 {top1.avg:.3f}'.format(
                          epoch, i, len(train_loader), algo=args.prefix[:-5], t=time.time()-start, lr=lr,
                   mean_cos=mean_cos, loss=losses, top1=top1))

        if args.summary:
            iters = epoch * len(train_loader) + i
            logger.scalar_summary('train_loss', float(losses.avg), iters)
            logger.scalar_summary('train_acc', float(top1.avg), iters)

    return top1.avg

def am_margin(output, target, m):
    t = to_np(copy.copy(target))
    n = output.size()[0]
    n = list(range(n))
    mean_cos = torch.mean(output[n, t]).item() # target logit
    output[n, t] = output[n, t] - m
    return mean_cos

# little different with mxnet version
# https://github.com/deepinsight/insightface/blob/1436c181056c703029d67dbadeb1408ffa14c59d/src/train_softmax.py#L228
def arc_margin(output, target, m):
    t = to_np(target)
    n = output.size()[0]
    n = list(range(n))
    cos_m = math.cos(m)
    sin_m = math.sin(m)
    mean_cos = torch.mean(output[n, t]).item() # target logit
    cos_t = 1 * output[n, t]
    sin_t = torch.sqrt(1 - cos_t * cos_t)
    output[n, t] = cos_m * cos_t - sin_m * sin_t
    return mean_cos

def save_checkpoint(state, filename):
    torch.save(state, filename)

class AverageMeter(object):
    """Computes and stores the average and current value"""
    def __init__(self):
        self.reset()

    def reset(self):
        self.val   = 0
        self.avg   = 0
        self.sum   = 0
        self.count = 0

    def update(self, val, n=1):
        self.val   = val
        self.sum   += val * n
        self.count += n
        self.avg   = self.sum / self.count

def adjust_learning_rate(optimizer, epoch):
    scale = 0.457305051927326
    step  = 10
    lr = args.lr * (scale ** (epoch // step))
    print('lr: {}'.format(lr))
    if (epoch != 0) & (epoch % step == 0):
        print('Change lr')
        for param_group in optimizer.param_groups:
            param_group['lr'] = param_group['lr'] * scale

def accuracy(output, target, topk=(1,)):
    """Computes the precision@k for the specified values of k"""
    maxk = max(topk)
    batch_size = target.size(0)

    _, pred = output.topk(maxk, 1, True, True)
    pred    = pred.t()
    correct = pred.eq(target.view(1, -1).expand_as(pred))

    res = []
    for k in topk:
        correct_k = correct[:k].view(-1).float().sum(0)
        res.append(correct_k.mul_(100.0 / batch_size))
    return res

if __name__ == '__main__':
    main()


# predict.py - predict (infer) inputs (single/batch).
import os
import sys

sys.path.append(".")  # necessary to run predict.py from command line *magic line*
import collections
import json
from argparse import ArgumentParser, Namespace
from io import BytesIO

import numpy as np
from tensorflow.keras.preprocessing.text import tokenizer_from_json

from dogapp import dogconfig, data, models, utils

sys.path.append(".")

def predict(url, data, model):
    """Predict the class for a text using
    a trained model from an experiment."""

    # url = [sample['url'] for sample in inputs]

    with open(os.path.join(os.getcwd(), "dogapp/dog_names.txt")) as f:
        dog_names = f.read().splitlines()

    # Predict
    results = []
    bottleneck_feature = utils.extract_Xception(data)
    predicted_vector = model.predict(bottleneck_feature)
    p_class = dog_names[np.argmax(predicted_vector)]
    results.append(
        {
            "input_url": url,
            "class": p_class,
        }
    )
    return results

if __name__ == "__main__":
    # Arguments
    parser = ArgumentParser()
    parser.add_argument("--url", type=str, required=True, help="image url to predict")
    args = parser.parse_args()
    inputs = [{"url": args.url}]

    # Get run components for prediction
    model = models.DogCNN()
    model.summary(input_shape=(7, 7, 2048))  # build it
    model_path = os.path.join(os.getcwd(), "embeddings/weights.best.Xception.hdf5")
    model.load_weights(model_path)
    data = utils.loadImage(inputs[0]["url"])

    # Predict
    prediction = []
    prediction.append(predict(url=inputs[0]["url"], data=data, model=model)[0])

    dogconfig.logger.info(json.dumps(prediction, indent=4, sort_keys=False))