"""Modified NACA 4/5-digit thickness relations."""

import numpy as np

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

from .thickness_base import (
    CLASSIC_LEADING_EDGE_INDEX_THRESHOLD,
    LEADING_EDGE_INDEX_MAX,
    LEADING_EDGE_INDEX_MIN,
    LEADING_EDGE_RADIUS_SCALE,
    MAX_THICKNESS_INDEX_LIMIT,
    MAX_THICKNESS_LOCATION_MAX,
    MAX_THICKNESS_LOCATION_MIN,
    OPEN_MODIFIED_TRAILING_EDGE_THICKNESS,
    Thickness,
)


class Naca45DigitModifiedThicknessClassic(Thickness):
    """
    Base class for the NACA modified 4-digit and 5-digit airfoil thickness.

    The thickness is parameterized on the square root of the chord location
    where the thickness is desired. This is to remove singularities that
    occur at the leading edge for the typical chord length parameterization.

    """

    def __init__(self, mti: float, lei: float, lmti: float) -> None:
        """
        Initialize a modified NACA thickness distribution.

        Parameters
        ----------
        mti : float
            Maximum-thickness index as a percent of chord.
        lei : float
            Leading-edge shape index.
        lmti : float
            Maximum-thickness-location index.

        Raises
        ------
        ValueError
            If any input index is outside the supported modified-series
            range.
        """
        # start with valid defaults for setters to work
        self._closed_te = False
        self._lei: float = 4
        self._t_m: float = 4
        self._a = np.zeros(4)
        self._d = np.zeros(4)

        # use settters to ensure valid data
        self.max_thickness_index = mti
        self.loc_max_thickness_index = lmti
        self.leading_edge_index = lei

    @property
    def max_thickness_index(self) -> float:
        """
        Return the maximum-thickness index.

        Returns
        -------
        float
            Maximum thickness as a percent-chord index.
        """
        return 100 * self._tmax

    @max_thickness_index.setter
    def max_thickness_index(self, mti: float) -> None:
        if mti < 0 or mti >= MAX_THICKNESS_INDEX_LIMIT:
            raise ValueError(
                f"Invalid NACA modified 4/5-digit maximum thickness: {mti}"
            )
        self._tmax = mti / 100.0

    @property
    def leading_edge_index(self) -> float:
        """
        Return the leading-edge index.

        Returns
        -------
        float
            Leading-edge shape index for the modified thickness relation.
        """
        return self._lei

    @leading_edge_index.setter
    def leading_edge_index(self, lei: float) -> None:
        if lei < LEADING_EDGE_INDEX_MIN or lei >= LEADING_EDGE_INDEX_MAX:
            raise ValueError(
                f"Invalid NACA modified 4/5-digit leading edge parameter: {lei}"
            )
        self._lei = lei
        self._calculate_coefficients()

    @property
    def loc_max_thickness_index(self) -> float:
        """
        Return the maximum-thickness-location index.

        Returns
        -------
        float
            Index locating where the fore and aft relations meet.
        """
        return 10 * self._t_m**2

    @loc_max_thickness_index.setter
    def loc_max_thickness_index(self, lmti: float) -> None:
        if (
            lmti < MAX_THICKNESS_LOCATION_MIN
            or lmti >= MAX_THICKNESS_LOCATION_MAX
        ):
            raise ValueError(
                "Invalid NACA modified 4/5-digit maximum thickness "
                f"location parameter: {lmti}"
            )
        self._t_m = np.sqrt(lmti / 10.0)
        self._calculate_coefficients()

    @property
    def a(self) -> FloatArray:
        """
        Return the fore-section coefficients.

        Returns
        -------
        FloatArray
            Coefficients for the fore-section thickness relation.
        """
        return self._a

    @property
    def d(self) -> FloatArray:
        """
        Return the aft-section coefficients.

        Returns
        -------
        FloatArray
            Coefficients for the aft-section thickness relation.
        """
        return self._d

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

        Returns
        -------
        List[float]
            Parametric coordinates of any discontinuities.
        """
        return [self._t_m]

    def max_thickness_parameter(self) -> float:
        """
        Return parameter coordinate of maximum thickness.

        Returns
        -------
        float
            Parameter coordinate of maximum thickness.
        """
        return self._t_m

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

        Parameters
        ----------
        t : numpy.ndarray
            Parameter location of interest. Equal to the square root of the
            desired chord location.

        Returns
        -------
        numpy.ndarray
            Thickness at specified parameter.
        """
        t = as_float_array(t)

        def fore(t: FloatArray) -> FloatArray:
            t2 = t**2
            return t * (
                self.a[0] + t * (self.a[1] + t2 * (self.a[2] + t2 * self.a[3]))
            )

        def aft(t: FloatArray) -> FloatArray:
            t2 = t**2
            term = 1 - t2
            return self.d[0] + term * (
                self.d[1] + term * (self.d[2] + term * self.d[3])
            )

        return self._tmax * np.piecewise(
            t, [t <= self._t_m, t > self._t_m], [fore, aft]
        )

    def delta_t(self, t: FloatInput) -> FloatArray:
        """
        Return first derivative of thickness at specified parameter location.

        Parameters
        ----------
        t : numpy.ndarray
            Parameter location of interest. Equal to the square root of the
            desired chord location.

        Returns
        -------
        numpy.ndarray
            First derivative of thickness at specified parameter.
        """
        t = as_float_array(t)

        def fore(t: FloatArray) -> FloatArray:
            t2 = t**2
            return self.a[0] + 2 * t * (
                self.a[1] + t2 * (2 * self.a[2] + 3 * self.a[3] * t2)
            )

        def aft(t: FloatArray) -> FloatArray:
            t2 = t**2
            return (
                -2
                * t
                * (
                    self.d[1]
                    + (1 - t2) * (2 * self.d[2] + 3 * (1 - t2) * self.d[3])
                )
            )

        return self._tmax * np.piecewise(
            t, [t <= self._t_m, t > self._t_m], [fore, aft]
        )

    def delta_tt(self, t: FloatInput) -> FloatArray:
        """
        Return second derivative of thickness at specified parameter location.

        Parameters
        ----------
        t : numpy.ndarray
            Parameter location of interest. Equal to the square root of the
            desired chord location.

        Returns
        -------
        numpy.ndarray
            Second derivative of thickness at specified parameter.
        """
        t = as_float_array(t)

        def fore(t: FloatArray) -> FloatArray:
            t2 = t**2
            return 2 * (
                self.a[1] + 3 * t2 * (2 * self.a[2] + 5 * t2 * self.a[3])
            )

        def aft(t: FloatArray) -> FloatArray:
            t2 = t**2
            return -2 * (
                self.d[1]
                + 2 * (1 - 3 * t2) * self.d[2]
                + 3 * (1 - t2) * (1 - 5 * t2) * self.d[3]
            )

        return self._tmax * np.piecewise(
            t, [t <= self._t_m, t > self._t_m], [fore, aft]
        )

    def _calculate_coefficients(self) -> None:
        """
        Recompute the fore and aft polynomial coefficients.

        Returns
        -------
        None
            This method updates the stored coefficient arrays in place.
        """
        # Pade approximation that goes through all Stack and von Doenhoff
        # (1935) values. Improves upon Riegels (1961) fit.
        p = [1.0310900853, -2.7171508529, 4.8594083156]
        q = [1.0, -1.8252487562, 1.1771499645]
        xi_m = self._t_m**2
        tau = (p[0] + p[1] * xi_m + p[2] * xi_m**2) / (
            q[0] + q[1] * xi_m + q[2] * xi_m**2
        )

        # calculate the d coefficients
        eta = 0.0 if self._closed_te else OPEN_MODIFIED_TRAILING_EDGE_THICKNESS

        self._d[0] = 0.5 * eta
        self._d[1] = tau
        self._d[2] = (2 * tau * (xi_m - 1) - 1.5 * (eta - 1)) / (xi_m - 1) ** 2
        self._d[3] = (tau * (xi_m - 1) - (eta - 1)) / (xi_m - 1) ** 3

        # calculate the a coefficients
        q_factor = LEADING_EDGE_RADIUS_SCALE
        if self._lei < CLASSIC_LEADING_EDGE_INDEX_THRESHOLD:
            q_factor /= 72
        else:
            q_factor /= 54

        k_m = (3 * (eta - 1) - 2 * tau * (xi_m - 1)) / (xi_m - 1) ** 2
        sqrt_term = np.sqrt(2 * q_factor)
        sqrt_term2 = np.sqrt(2 * q_factor * xi_m)
        self._a[0] = self._lei * sqrt_term
        self._a[1] = (
            0.5 * xi_m * (k_m + (3 - 3.75 * self._lei * sqrt_term2) / xi_m**2)
        )
        self._a[2] = -(k_m + (1.5 - 1.25 * self._lei * sqrt_term2) / xi_m**2)
        self._a[3] = (
            0.5 / xi_m * (k_m + (1 - 0.75 * self._lei * sqrt_term2) / xi_m**2)
        )


class Naca45DigitModifiedThicknessParams(Naca45DigitModifiedThicknessClassic):
    """
    Parametric NACA modified 4-digit and 5-digit thickness relation.

    The thickness is parameterized on the square root of the chord location
    where the thickness is desired. This is to remove singularities that
    occur at the leading edge for the typical chord length parameterization.

    This class extends the standard modified thickness distribution relation
    by solving for coefficients from the original constraints for
    non-integer parameter values and by allowing a closed trailing edge.
    """

    def __init__(
        self, mti: float, lei: float, lmti: float, closed_te: bool
    ) -> None:
        """
        Initialize a continuous-index modified NACA thickness distribution.

        Parameters
        ----------
        mti : float
            Maximum-thickness index as a percent of chord.
        lei : float
            Leading-edge shape index.
        lmti : float
            Maximum-thickness-location index.
        closed_te : bool
            Whether the trailing edge closes to zero thickness.
        """
        super().__init__(mti=mti, lei=lei, lmti=lmti)
        self.closed_te = closed_te

    @property
    def closed_te(self) -> bool:
        """
        Return the trailing-edge closure flag.

        Returns
        -------
        bool
            ``True`` when the trailing edge closes to zero thickness.
        """
        return self._closed_te

    @closed_te.setter
    def closed_te(self, closed_te: bool) -> None:
        if self._closed_te != closed_te:
            self._closed_te = closed_te
            self._calculate_coefficients()
