"""Classes associated with airfoils and their analysis."""

import numpy as np
from scipy.optimize import root_scalar

from buffalo_wings.internal.numeric import as_float_array, as_float_scalar
from buffalo_wings.type_aliases import FloatArray, FloatInput, FloatScalar

from .airfoil_schema import AirfoilDefinitionSpec
from .camber_base import Camber
from .curve import Curve
from .thickness_base import Thickness

ROOT_TOLERANCE = 1e-8
CAMBER_TOLERANCE = 1e-7
LOWER_SURFACE_BRACKET = (-1.0, 0.0)
UPPER_SURFACE_BRACKET = (0.0, 1.0)
FULL_AIRFOIL_BRACKET = (-1.0, 1.0)


class Airfoil(Curve):
    """
    Base class for airfoil specific geometries.

    The airfoil coordinates and their derivatives can be queried using two
    parameterizations.
    The standard parameterization uses ``[-1, 0]`` for the lower surface,
    running from trailing edge to leading edge, and ``[0, 1]`` for the upper
    surface, running from leading edge to trailing edge.
    This provides a smooth parameterization across the full airfoil and
    matches how many analytic airfoil descriptions are written.
    The arc-length parameterization uses surface distance measured from the
    lower trailing edge to the upper trailing edge.
    Arc-length queries are more expensive because the mapping from surface
    distance to the native airfoil parameter is not available in closed form
    for general airfoil shapes.
    """

    def __init__(self) -> None:
        """
        Initialize the cached airfoil arc-length state.

        Notes
        -----
        Subclasses should call this initializer so cached geometric values
        remain synchronized with the current airfoil shape.
        """
        super().__init__()
        self._s_max: FloatScalar | None = None

    @property
    def spec(self) -> AirfoilDefinitionSpec:
        """
        Return the schema definition used to create this airfoil.

        Returns
        -------
        AirfoilDefinitionSpec
            Serialized airfoil definition that can recreate this runtime
            object.

        Raises
        ------
        NotImplementedError
            If the concrete airfoil type does not preserve its source spec.
        """
        raise NotImplementedError(
            f"{type(self).__name__} does not provide a source spec."
        )

    def to_spec(self) -> AirfoilDefinitionSpec:
        """
        Return the schema definition needed to recreate this airfoil.

        Returns
        -------
        AirfoilDefinitionSpec
            Serialized airfoil definition that can recreate this runtime
            object.
        """
        return self.spec

    @property
    def surface_length(self) -> FloatScalar:
        """
        Return the full airfoil surface length.

        Returns
        -------
        float
            Arc length measured from the lower trailing edge to the upper
            trailing edge.
        """
        if self._s_max is None:
            self._s_max = as_float_scalar(self.arc_length(-1.0, 1.0))

        return self._s_max

    def t_from_x(self, x: FloatInput, upper: bool) -> FloatArray:
        """
        Calculate the parametric value for x-location provided.

        Parameters
        ----------
        x : FloatArray
            X-coordinate of interest.
        upper : bool
            True if want upper surface point otherwise get lower surface point.

        Returns
        -------
        numpy.ndarray
            Parameteric value for location provided.

        Raises
        ------
        ValueError
            If there is no surface point at the given x-location.
        """
        x = np.asarray(x, dtype=np.float64)
        xmin, xmax = self._surface_x_bounds(upper=upper)
        if (x < xmin - ROOT_TOLERANCE).any() or (
            x > xmax + ROOT_TOLERANCE
        ).any():
            surface_name = "upper" if upper else "lower"
            msg = (
                "Invalid x-coordinate provided. "
                f"The {surface_name} surface supports "
                f"{xmin:.6g} <= x <= {xmax:.6g}."
            )
            raise ValueError(msg)

        def fun(ti: float, x_target: float) -> float:
            x_value, _ = self.xy(ti)
            return as_float_scalar(x_value) - x_target

        it = np.nditer([x, None])
        bracket = UPPER_SURFACE_BRACKET if upper else LOWER_SURFACE_BRACKET
        with it:
            for xx, ti in it:
                xxf = float(xx)
                if xx < 0:
                    root = root_scalar(
                        fun,
                        args=(xxf,),
                        x0=0.0,
                        x1=abs(xxf),
                    )
                    ti[...] = float(root.root)
                elif np.abs(fun(bracket[0], xxf)) < ROOT_TOLERANCE:
                    ti[...] = bracket[0]
                elif np.abs(fun(bracket[1], xxf)) < ROOT_TOLERANCE:
                    ti[...] = bracket[1]
                else:
                    root = root_scalar(
                        fun,
                        args=(xxf,),
                        bracket=bracket,
                    )
                    ti[...] = float(root.root)

            return it.operands[1]

    def _surface_x_bounds(self, upper: bool) -> tuple[FloatScalar, FloatScalar]:
        """
        Return the reachable x-range for one airfoil surface.

        Parameters
        ----------
        upper : bool
            Whether to query the upper or lower surface interval.

        Returns
        -------
        float
            Minimum x-coordinate reachable on the selected surface.
        float
            Maximum x-coordinate reachable on the selected surface.
        """
        bracket = UPPER_SURFACE_BRACKET if upper else LOWER_SURFACE_BRACKET
        x0 = as_float_scalar(self.xy(bracket[0])[0])
        x1 = as_float_scalar(self.xy(bracket[1])[0])
        return min(x0, x1), max(x0, x1)

    def t_from_s(self, s: FloatInput) -> FloatArray:
        """
        Calculate the parametric value for arc-length provided.

        Parameters
        ----------
        s : numpy.ndarray
            Arc-length location of point.

        Raises
        ------
        ValueError
            When arc-length provided is larger than airfoil surface length.

        Returns
        -------
        numpy.ndarray
            Parametric value for location provided.
        """
        s_a = np.asarray(s, dtype=np.float64)
        if (s_a > self.surface_length).any() or (s_a < 0).any():
            msg = (
                "Invalid arc length provided. "
                f"Valid range is 0 <= s <= {self.surface_length:.6g}."
            )
            raise ValueError(msg)

        def fun(t: float, s_target: float) -> float:
            return as_float_scalar(self.arc_length(-1.0, t)) - s_target

        it = np.nditer([s_a, None])
        with it:
            for ss, ti in it:
                root = root_scalar(
                    fun,
                    args=(float(ss),),
                    bracket=FULL_AIRFOIL_BRACKET,
                )
                ti[...] = float(root.root)

            return it.operands[1]

    def dydx(self, t: FloatInput) -> FloatArray:
        """
        Calculate the slope at parameter location.

        Parameters
        ----------
        t : numpy.ndarray
            Parameter for desired locations.

        Returns
        -------
        numpy.ndarray
            Slope of surface at point.

        Raises
        ------
        ValueError
            If there is no surface point at the given x-location.
        """
        xt, yt = self.xy_t(t)
        return yt / xt

    def d2ydx2(self, t: FloatInput) -> FloatArray:
        """
        Calculate the second derivative at parameter location.

        Parameters
        ----------
        t : numpy.ndarray
            Parameter for desired locations.

        Returns
        -------
        numpy.ndarray
            Second derivative of surface at point.

        Raises
        ------
        ValueError
            If there is no surface point at the given x-location.
        """
        xt, yt = self.xy_t(t)
        xtt, ytt = self.xy_tt(t)
        return (xt * ytt - yt * xtt) / xt**3

    #
    # Arc-length Interface
    #
    def xy_from_s(self, s: FloatInput) -> tuple[FloatArray, FloatArray]:
        """
        Calculate the coordinates of geometry at arc-length location.

        Parameters
        ----------
        s : numpy.ndarray
            Arc-length location for point.

        Returns
        -------
        numpy.ndarray
            X-coordinate of point.
        numpy.ndarray
            Y-coordinate of point.
        """
        t = self.t_from_s(s)
        return self.xy(t)

    def xy_s(self, s: FloatInput) -> tuple[FloatArray, FloatArray]:
        """
        Calculate rates of change of the coordinates at arc-length location.

        Parameters
        ----------
        s : numpy.ndarray
            Arc-length location for point.

        Returns
        -------
        numpy.ndarray
            Arc-length rate of change of the x-coordinate of point.
        numpy.ndarray
            Arc-length rate of change of the y-coordinate of point.
        """
        t = self.t_from_s(s)
        return self.tangent(t)

    def xy_ss(self, s: FloatInput) -> tuple[FloatArray, FloatArray]:
        """
        Calculate second derivative of the coordinates at arc-length location.

        Parameters
        ----------
        s : numpy.ndarray
            Arc-length location for point.

        Returns
        -------
        numpy.ndarray
            Arc-length second derivative of the x-coordinate of point.
        numpy.ndarray
            Arc-length second derivative of the y-coordinate of point.
        """
        t = self.t_from_s(s)
        nx, ny = self.normal(t)
        k = self.k(t)
        return k * nx, k * ny

    def chord(self) -> FloatScalar:
        """
        Return the chord length of the airfoil.

        Returns
        -------
        float
            Chord length.
        """
        xle, yle = self.leading_edge()
        xte, yte = self.trailing_edge()
        return float(np.sqrt((xte - xle) ** 2 + (yte - yle) ** 2))

    def leading_edge(self) -> tuple[FloatScalar, FloatScalar]:
        """
        Return the location of the leading edge.

        Returns
        -------
        float
            X-coordinate of leading edge.
        float
            Y-coordinate of leading edge.
        """
        x_le, y_le = self.xy(0.0)
        return as_float_scalar(x_le), as_float_scalar(y_le)

    def trailing_edge(self) -> tuple[FloatScalar, FloatScalar]:
        """
        Return the location of the trailing edge.

        Notes
        -----
        Since some airfoil descriptions have gap between the upper and lower
        surface at the trailing edge (such as NACA 4-digit and 5-digit
        airfoils), the point returned is the average of the two trailing
        edge points. If the specific location of the upper or the lower
        surface of the trailing edge is desired, use :py:meth:`xy` passing in
        either -1 (lower) or +1 (upper).

        Returns
        -------
        float
            X-coordinate of trailing edge.
        float
            Y-coordinate of trailing edge.
        """
        xl, yl = self.xy(-1)
        xu, yu = self.xy(1)
        x_te = 0.5 * (as_float_scalar(xl) + as_float_scalar(xu))
        y_te = 0.5 * (as_float_scalar(yl) + as_float_scalar(yu))
        return x_te, y_te

    def _airfoil_changed(self) -> None:
        """
        Notify airfoil that shape has changed.

        This needs to be called by child classes when the airfoil geometry
        has changed so any cached values can be invalidated.

        Returns
        -------
        None
            This method updates internal caches in place.
        """
        self._s_max = None


class OrthogonalAirfoil(Airfoil):
    """
    Airfoils that can be decomposed to camber and thickness.

    This class represents airfoils that are naturally described by a camber
    line (curve) and a thickness normal to the camber line, both above and
    below the camber line. The parameterization of the camber representation
    and the thickness representation must be based on the same transformation.

    Notes
    -----
    The public airfoil parameter ``t`` is a signed surface parameter in
    ``[-1, 1]``.
    Negative values lie on the lower surface and positive values lie on the
    upper surface.
    Internally, the geometry is evaluated using the unsigned surface
    parameter ``|t|`` together with a separate surface sign.

    The camber parameter ``u`` is the chord-like parameter used by the
    camber model, with ``u = |t|**2``.
    The thickness parameter ``v`` is the square-root chord parameter used by
    the thickness model, with ``v = |t|``.
    Equivalently, ``u = v**2`` and ``v = sqrt(u)``.

    This convention is intentional.
    It removes the leading-edge square-root singularity from the thickness
    representation while preserving a smooth airfoil surface parameterization
    for coordinate and derivative evaluation.
    """

    def __init__(self, camber: Camber, thickness: Thickness) -> None:
        """
        Store the camber line and thickness for the airfoil.

        Parameters
        ----------
        camber : Camber
            Camber-line representation used by the airfoil model.
        thickness : Thickness
            Thickness distribution paired with ``camber``.
        """
        super().__init__()
        self._camber = camber
        self._thickness = thickness
        self._t_xmin: FloatScalar | None = None
        self._t_xmax: FloatScalar | None = None

    @property
    def camber(self) -> Camber:
        """
        Return the camber function for airfoil.

        Returns
        -------
        Camber
            Camber-line model associated with this airfoil.
        """
        return self._camber

    @property
    def thickness(self) -> Thickness:
        """
        Return the thickness function for airfoil.

        Returns
        -------
        Thickness
            Thickness model associated with this airfoil.
        """
        return self._thickness

    @property
    def xmin_parameter(self) -> FloatScalar:
        """
        Return the parameter of the smallest x-coordinate for the airfoil.

        Returns
        -------
        float
            Native airfoil parameter at the minimum x-location.
        """
        if self._t_xmin is None:
            max_camber_parameter = self._camber.max_camber_parameter()
            _, max_camber_value = self._camber.xy(max_camber_parameter)
            max_camber = as_float_scalar(max_camber_value)
            if abs(max_camber) < CAMBER_TOLERANCE:
                self._t_xmin = 0.0
            else:
                if max_camber > 0:
                    t0 = 1e-7
                    t1 = 0.1
                else:
                    t0 = -1e-7
                    t1 = -0.1
                root = root_scalar(
                    lambda t: as_float_scalar(self.xy_t(t)[0]),
                    bracket=(t0, t1),
                )
                self._t_xmin = float(root.root)

        return self._t_xmin

    @property
    def xmax_parameter(self) -> FloatScalar:
        """
        Return the parameter of the largest x-coordinate for the airfoil.

        Returns
        -------
        float
            Native airfoil parameter at the maximum x-location.
        """
        if self._t_xmax is None:
            if self.xy(-1)[0] >= self.xy(1)[0]:
                self._t_xmax = -1.0
            else:
                self._t_xmax = 1.0

        return self._t_xmax

    def _surface_x_bounds(self, upper: bool) -> tuple[FloatScalar, FloatScalar]:
        """
        Return the reachable x-range for one orthogonal airfoil surface.

        Notes
        -----
        Cambered orthogonal airfoils can reach their minimum x-value on one
        surface slightly ahead of the nominal leading-edge point, so the
        selected surface must include the relevant interior extremum.
        """
        bracket = UPPER_SURFACE_BRACKET if upper else LOWER_SURFACE_BRACKET
        candidate_parameters: list[FloatScalar] = [bracket[0], bracket[1]]
        if bracket[0] <= self.xmin_parameter <= bracket[1]:
            candidate_parameters.append(self.xmin_parameter)
        if bracket[0] <= self.xmax_parameter <= bracket[1]:
            candidate_parameters.append(self.xmax_parameter)

        x_values = [
            as_float_scalar(self.xy(parameter)[0])
            for parameter in candidate_parameters
        ]
        return min(x_values), max(x_values)

    def xy(self, t: FloatInput) -> tuple[FloatArray, FloatArray]:
        """
        Calculate the coordinates of geometry at parameter location.

        Notes
        -----
        Parameter goes from -1 (trailing edge lower surface) to +1 (trailing
        edge upper surface) with 0 representing the leading edge.

        Parameters
        ----------
        t : numpy.ndarray
            Parameter for desired locations.

        Returns
        -------
        numpy.ndarray
            X-coordinate of point.
        numpy.ndarray
            Y-coordinate of point.
        """
        tc, sgn = self._convert_t(t)

        u = self._convert_t_to_u(tc)[0]
        xi, eta = self._camber.xy(u)
        nx, ny = self._camber.normal(u)

        v = self._convert_t_to_v(tc)[0]
        delta = sgn * self._thickness.delta(v)

        x = xi + delta * nx
        y = eta + delta * ny
        # return self.xo + self.scale*x, self.yo + self.scale*y
        return x, y

    def xy_t(self, t: FloatInput) -> tuple[FloatArray, FloatArray]:
        """
        Calculate rates of change of the coordinates at parameter location.

        Notes
        -----
        Parameter goes from -1 (trailing edge lower surface) to +1 (trailing
        edge upper surface) with 0 representing the leading edge.

        Parameters
        ----------
        t : numpy.ndarray
            Parameter for desired locations.

        Returns
        -------
        numpy.ndarray
            Parametric rate of change of the x-coordinate of point.
        numpy.ndarray
            Parametric rate of change of the y-coordinate of point.
        """
        tc, sgn = self._convert_t(t)

        u, u_t = self._convert_t_to_u(tc)[0:2]
        xi_u, eta_u = self._camber.xy_t(u)
        k = self._camber.k(u)
        nx, ny = self._camber.normal(u)
        nx_u = -xi_u * k
        ny_u = -eta_u * k

        v, v_t = self._convert_t_to_v(tc)[0:2]
        delta = sgn * self._thickness.delta(v)
        delta_v = sgn * self._thickness.delta_t(v)

        x_t = sgn * (u_t * xi_u + v_t * delta_v * nx + delta * u_t * nx_u)
        y_t = sgn * (u_t * eta_u + v_t * delta_v * ny + delta * u_t * ny_u)

        # return self.scale*x_t, self.scale*y_t
        return x_t, y_t

    def xy_tt(self, t: FloatInput) -> tuple[FloatArray, FloatArray]:
        """
        Return second derivative of the coordinates at parameter location.

        Notes
        -----
        Parameter goes from -1 (trailing edge lower surface) to +1 (trailing
        edge upper surface) with 0 representing the leading edge.

        Parameters
        ----------
        t : numpy.ndarray
            Parameter for desired locations.

        Returns
        -------
        numpy.ndarray
            Parametric second derivative of the x-coordinate of point.
        numpy.ndarray
            Parametric second derivative of the y-coordinate of point.
        """
        tc, sgn = self._convert_t(t)

        u, u_t, u_tt = self._convert_t_to_u(tc)
        xi_u, eta_u = self._camber.xy_t(u)
        xi_uu, eta_uu = self._camber.xy_tt(u)
        k = self._camber.k(u)
        k_u = self._camber.k_t(u)
        nx, ny = self._camber.normal(u)
        nx_u = -xi_u * k
        nx_uu = -(xi_uu * k + xi_u * k_u)
        ny_u = -eta_u * k
        ny_uu = -(eta_uu * k + eta_u * k_u)

        v, v_t, v_tt = self._convert_t_to_v(tc)
        delta = sgn * self._thickness.delta(v)
        delta_v = sgn * self._thickness.delta_t(v)
        delta_vv = sgn * self._thickness.delta_tt(v)

        x_tt = (
            u_tt * xi_u
            + u_t**2 * xi_uu
            + v_t**2 * nx * delta_vv
            + u_t**2 * delta * nx_uu
            + 2 * u_t * v_t * delta_v * nx_u
            + v_tt * nx * delta_v
            + u_tt * delta * nx_u
        )
        y_tt = (
            u_tt * eta_u
            + u_t**2 * eta_uu
            + v_t**2 * ny * delta_vv
            + u_t**2 * delta * ny_uu
            + 2 * u_t * v_t * delta_v * ny_u
            + v_tt * ny * delta_v
            + u_tt * delta * ny_u
        )

        # return self.scale*x_tt, self.scale*y_tt
        return x_tt, y_tt

    def camber_location(self, t: FloatInput) -> tuple[FloatArray, FloatArray]:
        """
        Return the location of camber at specified parameter location.

        Parameters
        ----------
        t : numpy.ndarray
            Parameter location of interest.

        Returns
        -------
        numpy.ndarray
            X-coordinate of camber at specified point.
        numpy.ndarray
            Y-coordinate of camber at specified point.
        """
        return self._camber.xy(self._convert_t_to_u(t)[0])

    def thickness_value(self, t: FloatInput) -> FloatArray:
        """
        Return the amount of thickness at specified parameter location.

        Parameters
        ----------
        t : numpy.ndarray
            Parameter location of interest.

        Returns
        -------
        numpy.ndarray
            Thickness at specified point.
        """
        return self._thickness.delta(self._convert_t_to_v(t)[0])

    def joints(self) -> list[float]:
        """
        Return the locations of any joints/discontinuities in the curve.

        Returns
        -------
        List[float]
            Xi-coordinates of any discontinuities.
        """
        camber_joints = [
            as_float_scalar(self._convert_u_to_t(j))
            for j in self._camber.joints()
        ]
        thickness_joints = [
            as_float_scalar(self._convert_v_to_t(d))
            for d in self._thickness.discontinuities()
        ]
        joints = list(set(camber_joints + thickness_joints))
        joints = list(set(joints + [-x for x in joints]))
        joints.sort()
        return joints

    def _airfoil_changed(self) -> None:
        """
        Clear cached values after a geometry change.

        Returns
        -------
        None
            This method updates internal caches in place.
        """
        self._t_xmin = None
        self._t_xmax = None
        super()._airfoil_changed()

    @staticmethod
    def _convert_t_to_u(
        t: FloatInput,
    ) -> tuple[FloatArray, FloatArray, FloatArray]:
        """
        Convert the airfoil parameterization to the camber parameterization.

        This maps the unsigned surface parameter ``|t|`` to the chord-like
        camber parameter ``u = |t|**2``.

        Parameters
        ----------
        t : numpy.ndarray
            Airfoil parameterization.

        Returns
        -------
        numpy.ndarray
            Corresponding camber parameterization.
        numpy.ndarray
            Corresponding first derivative of camber parameterization.
        numpy.ndarray
            Corresponding second derivative of camber parameterization.
        """
        t_array = as_float_array(t)
        return t_array**2, 2 * t_array, 2 * np.ones_like(t_array)

    @staticmethod
    def _convert_u_to_t(u: FloatInput) -> FloatArray:
        """
        Convert the camber parameterization to the airfoil parameterization.

        This maps the chord-like camber parameter back to the unsigned
        surface parameter via ``|t| = sqrt(u)``.

        Parameters
        ----------
        u : numpy.ndarray
            Camber parameterization.

        Returns
        -------
        numpy.ndarray
            Corresponding airfoil parameterization.
        """
        u = np.asarray(u, dtype=np.float64)
        return np.sqrt(u)

    @staticmethod
    def _convert_t_to_v(
        t: FloatInput,
    ) -> tuple[FloatArray, FloatArray, FloatArray]:
        """
        Convert the airfoil parameterization to the thickness parameterization.

        This maps the unsigned surface parameter directly to the square-root
        chord thickness parameter via ``v = |t|``.

        Parameters
        ----------
        t : numpy.ndarray
            Airfoil parameterization.

        Returns
        -------
        numpy.ndarray
            Corresponding thickness parameterization.
        numpy.ndarray
            Corresponding first derivative of thickness parameterization.
        numpy.ndarray
            Corresponding second derivative of thickness parameterization.
        """
        t_array = as_float_array(t)
        return t_array, np.ones_like(t_array), np.zeros_like(t_array)

    @staticmethod
    def _convert_v_to_t(u: FloatInput) -> FloatArray:
        """
        Convert the thickness parameterization to the airfoil parameterization.

        Since ``v = |t|``, this inverse is the identity on the unsigned
        surface parameter.

        Parameters
        ----------
        u : numpy.ndarray
            Thickness parameterization.

        Returns
        -------
        numpy.ndarray
            Corresponding airfoil parameterization.
        """
        return as_float_array(u)

    @staticmethod
    def _convert_t(t: FloatInput) -> tuple[FloatArray, FloatArray]:
        """
        Split signed airfoil coordinates into magnitude and surface sign.

        Parameters
        ----------
        t : FloatArray
            Signed airfoil parameter values in ``[-1, 1]``.

        Returns
        -------
        FloatArray
            Nonnegative parameter magnitudes.
        FloatArray
            Surface sign array with ``-1`` on the lower surface and ``+1``
            on the upper surface.
        """
        tc = np.asarray(t, dtype=np.float64).copy()
        sgn = np.ones_like(tc)
        sgn[tc < 0] = -1.0
        tc[tc < 0] = -tc[tc < 0]
        return tc, sgn
