from dataclasses import dataclass
from typing import Any, Dict, Tuple

import numpy as np
from scipy.integrate import solve_ivp
from scipy.optimize import OptimizeResult

from flying_discs.frispy.scipy_backend.eom import EOM
from flying_discs.frispy.scipy_backend.model import Model


@dataclass
class Disc:
    # pylint: disable=too-many-instance-attributes
    """Flying spinning disc object.

    One can specify initial coordinates, initial kinematic variables, the two
    non-zero elements of the moments of inertia tensor, the face area, mass,
    air density, and gravitational acceleration. Note that the default area,
    mass, and inertial moments are for Discraft Ultrastars (0.175 kg).

    MKS units used throughout. Angular units use radians (rad), and
    angular velocities are in rad/s.

    Args:
        x: horizontal position
        y: horizontal position
        z: vertical position
        vx: horizontal speed
        vy: horizontal speed
        vz: vertical speed
        phi: first Euler angle. Pitch angle when velocity is horizontal
        theta: second Euler angle. Roll angle when velocity is horizontal
        gamma: third Euler angle. Spin angle
        dphi: angular velocity
        dtheta: angular velocity
        dgamma: spin angular velocity
        area: face area
        I_xx: pitch and roll moments of inertia
        I_zz: spin moment of inertia
        air_density: default is sea level
        g: gravitational acceleration
        model: defines how accelerations depend on velocities
        eom: equations of motion
    """

    def __init__(
        self,
        area: float = 0.058556,
        I_xx: float = 0.001219,
        I_zz: float = 0.002352,
        mass: float = 0.175,
        air_density: float = 1.225,
        g: float = 9.81,
    ) -> None:
        self.area = area
        self.I_xx = I_xx
        self.I_zz = I_zz
        self.mass = mass
        self.air_density = air_density
        self.g = g
        self.model = Model()
        self.eom = EOM(
            area=self.area,
            I_xx=self.I_xx,
            I_zz=self.I_zz,
            mass=self.mass,
            air_density=self.air_density,
            g=self.g,
        )

    def compute_trajectory(
        self,
        x: float = 0.0,
        y: float = 0.0,
        z: float = 1.0,
        vx: float = 10.0,
        vy: float = 0.0,
        vz: float = 0.0,
        phi: float = 0.0,
        theta: float = 0.0,
        gamma: float = 0.0,
        dphi: float = 0.0,
        dtheta: float = 0.0,
        dgamma: float = 62.0,
        flight_time: float = 3.0,
        n_times: int = 100,
        **solver_kwargs: Any,
    ) -> Tuple[Dict[str, np.ndarray], OptimizeResult]:
        # pylint: disable=too-many-positional-arguments, too-many-arguments, too-many-locals
        """Call the differential equation solver to compute the trajectory.

        The kinematic variables and timesteps are saved
        as the `current_trajectory` attribute, which is a dictionary,
        which is also returned by this function.

        See `the scipy docs
        <https://docs.scipy.org/doc/scipy/reference/generated/scipy.integrate.solve_ivp.html#scipy.integrate.solve_ivp>`_
        for more information on the solver.

        .. warning::

           `flight_time` is ignored if `t_span` is a key in
           `solver_kwargs`.

        Args:
            flight_time: time in seconds that the simulation
                will run over
            n_times: Number of samples in time for the trajectory.
                Samples are spaced evenly in time from ``(0, flight_time)``.
            solver_kwargs: extra keyword arguments to pass
                to the :meth:`scipy.integrate.solve_ivp`

        Returns:
            first element is a dict of all kinematic variables and the
            second is the ``solve_ivp`` results object
        """
        # Pop out these kwargs and take defaults based on our API
        t_span: tuple[float, float] = solver_kwargs.pop("t_span", (0, flight_time))
        t_eval = solver_kwargs.pop("t_eval", np.linspace(t_span[0], t_span[1], n_times))

        # Call the solver
        result = solve_ivp(
            fun=self.eom.compute_derivatives,
            t_span=t_span,
            y0=[
                x,
                y,
                z,
                vx,
                vy,
                vz,
                phi,
                theta,
                gamma,
                dphi,
                dtheta,
                dgamma,
            ],
            t_eval=t_eval,
            **solver_kwargs,
        )

        return (
            {
                "times": result.t,
                "x": result.y[0],
                "y": result.y[1],
                "z": result.y[2],
                "vx": result.y[3],
                "vy": result.y[4],
                "vz": result.y[5],
                "phi": result.y[6],
                "theta": result.y[7],
                "gamma": result.y[8],
                "dphi": result.y[9],
                "dtheta": result.y[10],
                "dgamma": result.y[11],
            },
            result,
        )
