"""Canonical wing-geometry evaluation from the documented schema."""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np

from buffalo_wings.airfoil import Airfoil, AirfoilFactory, DatAirfoilSpec
from buffalo_wings.internal.numeric import as_float_array, as_float_scalar
from buffalo_wings.type_aliases import FloatArray, FloatScalar

from .wing_schema import (
    DistributionSpec,
    PanelSpec,
    PiecewiseLinearDistribution,
    ReferenceAxis,
    SingleAirfoilRef,
    WingSpec,
)

_MIN_STATION_COUNT = 2


@dataclass(slots=True)
class EvaluatedStation:
    """
    Evaluated panel properties at a span station.

    Attributes
    ----------
    panel_id : str
        Identifier of the panel evaluated at this station.
    eta : FloatScalar
        Normalized spanwise coordinate used for evaluation.
    y : FloatScalar
        Spanwise position in the wing reference frame.
    x_ref : FloatScalar
        Reference-axis x-position at the span station.
    z_ref : FloatScalar
        Reference-axis z-position at the span station.
    chord : FloatScalar
        Local chord length at the span station.
    twist_rad : FloatScalar
        Local twist angle in radians.
    """

    panel_id: str
    eta: FloatScalar
    y: FloatScalar
    x_ref: FloatScalar
    z_ref: FloatScalar
    chord: FloatScalar
    twist_rad: FloatScalar


@dataclass(slots=True)
class SectionCurves:
    """
    Sampled upper and lower 3D section curves at one span station.

    Attributes
    ----------
    station : EvaluatedStation
        Evaluated panel properties used to place the section in 3D space.
    chordwise_samples : FloatArray
        One-dimensional normalized chordwise sample locations in ``[0, 1]``.
    upper_curve : FloatArray
        Array of shape ``(n_samples, 3)`` containing upper-surface points.
    lower_curve : FloatArray
        Array of shape ``(n_samples, 3)`` containing lower-surface points.
    """

    station: EvaluatedStation
    chordwise_samples: FloatArray
    upper_curve: FloatArray
    lower_curve: FloatArray


class WingCanonical:
    """
    Canonical evaluator for wing-section placement and sampling.

    Parameters
    ----------
    spec : WingSpec
        Parsed wing schema used for panel evaluation and section sampling.

    Raises
    ------
    ValueError
        If the configured reference-axis and twist-axis combination is not
        supported by the current implementation.
    """

    def __init__(self, spec: WingSpec) -> None:
        """
        Store the parsed wing specification and validate axis choices.

        Parameters
        ----------
        spec : WingSpec
            Parsed wing schema used for panel evaluation and section
            sampling.

        Raises
        ------
        ValueError
            If the configured reference-axis and twist-axis combination is
            not supported by the current implementation.
        """
        self._spec = spec
        self._panel_map = {panel.id: panel for panel in spec.wing.panels}
        self._validate_axes()

    @classmethod
    def from_spec(cls, spec: WingSpec) -> WingCanonical:
        """
        Build a canonical evaluator from a parsed ``WingSpec``.

        Parameters
        ----------
        spec : WingSpec
            Parsed wing schema used for panel evaluation and section
            sampling.

        Returns
        -------
        WingCanonical
            Canonical evaluator backed by ``spec``.

        Raises
        ------
        ValueError
            If the configured reference-axis and twist-axis combination is
            not supported by the current implementation.
        """
        return cls(spec)

    @property
    def spec(self) -> WingSpec:
        """
        Return the underlying wing specification.

        Returns
        -------
        WingSpec
            Wing schema used by the evaluator.
        """
        return self._spec

    def evaluate_distribution(
        self,
        distribution: DistributionSpec,
        eta: FloatScalar,
    ) -> FloatScalar:
        """
        Evaluate a scalar distribution at ``eta``.

        Parameters
        ----------
        distribution : DistributionSpec
            Distribution definition to evaluate.
        eta : FloatScalar
            Spanwise parameter at which to evaluate the distribution.

        Returns
        -------
        FloatScalar
            Interpolated scalar value at ``eta``.

        Raises
        ------
        ValueError
            If ``eta`` lies outside the distribution support or the
            distribution data are invalid.
        """
        return self._evaluate_piecewise_linear(distribution, eta)

    def sample_span_stations(
        self,
        panel_id: str,
        count: int,
        *,
        spacing: str = "uniform",
    ) -> FloatArray:
        """
        Generate span stations over a panel ``eta_range``.

        Parameters
        ----------
        panel_id : str
            Identifier of the panel to sample.
        count : int
            Number of stations to generate.
        spacing : str, default="uniform"
            Spacing rule used between the panel endpoints.
            Supported values are ``"uniform"`` and ``"cosine"``.

        Returns
        -------
        FloatArray
            One-dimensional array of sampled ``eta`` values.

        Raises
        ------
        KeyError
            If ``panel_id`` does not match a known panel.
        ValueError
            If ``count < 2`` or ``spacing`` is unsupported.
        """
        if count < _MIN_STATION_COUNT:
            msg = "Span station count must be at least 2."
            raise ValueError(msg)

        panel = self._get_panel(panel_id)
        eta0, eta1 = panel.eta_range
        if spacing == "uniform":
            return np.linspace(eta0, eta1, count, dtype=np.float64)
        if spacing == "cosine":
            beta = np.linspace(0.0, np.pi, count, dtype=np.float64)
            blend = 0.5 * (1.0 - np.cos(beta))
            return eta0 + (eta1 - eta0) * blend
        msg = f"Unsupported spacing method: {spacing}"
        raise ValueError(msg)

    def evaluate_panel(
        self, panel_id: str, eta: FloatScalar
    ) -> EvaluatedStation:
        """
        Evaluate the spanwise panel properties at ``eta``.

        Parameters
        ----------
        panel_id : str
            Identifier of the panel to evaluate.
        eta : FloatScalar
            Spanwise parameter within the panel ``eta_range``.

        Returns
        -------
        EvaluatedStation
            Evaluated reference-line position, chord, and twist at ``eta``.

        Raises
        ------
        KeyError
            If ``panel_id`` does not match a known panel.
        ValueError
            If ``eta`` lies outside the panel ``eta_range``.
        """
        panel = self._get_panel(panel_id)
        self._validate_eta(panel, eta)
        return EvaluatedStation(
            panel_id=panel_id,
            eta=eta,
            y=self._spec.wing.half_span * eta,
            x_ref=self.evaluate_distribution(panel.ref_line.x_ref, eta),
            z_ref=self.evaluate_distribution(panel.ref_line.z_ref, eta),
            chord=self.evaluate_distribution(panel.chord, eta),
            twist_rad=self._angle_to_radians(
                self.evaluate_distribution(panel.twist, eta)
            ),
        )

    def section_curves(
        self,
        panel_id: str,
        *,
        eta: FloatScalar,
        chordwise_samples: FloatArray,
    ) -> SectionCurves:
        """
        Generate sampled upper and lower 3D curves for one section.

        Parameters
        ----------
        panel_id : str
            Identifier of the panel to sample.
        eta : float
            Spanwise parameter within the panel ``eta_range``.
        chordwise_samples : FloatArray
            One-dimensional normalized chordwise coordinates in ``[0, 1]``.

        Returns
        -------
        SectionCurves
            Sampled upper and lower section curves in the canonical wing
            frame.

        Raises
        ------
        KeyError
            If ``panel_id`` or the referenced airfoil name is unknown.
        ValueError
            If ``eta`` lies outside the panel range or
            ``chordwise_samples`` is not a one-dimensional array in
            ``[0, 1]``.
        NotImplementedError
            If the panel uses an unsupported airfoil reference or airfoil
            type.
        """
        station = self.evaluate_panel(panel_id, eta)
        samples = as_float_array(chordwise_samples)
        if samples.ndim != 1:
            msg = "Chordwise samples must be a 1D array."
            raise ValueError(msg)
        if np.any(samples < 0.0) or np.any(samples > 1.0):
            msg = "Chordwise samples must lie in [0, 1]."
            raise ValueError(msg)

        airfoil = self._get_panel_airfoil(self._get_panel(panel_id))
        upper_curve = self._section_surface_curve(
            airfoil,
            station,
            samples,
            upper=True,
        )
        lower_curve = self._section_surface_curve(
            airfoil,
            station,
            samples,
            upper=False,
        )
        return SectionCurves(
            station=station,
            chordwise_samples=samples,
            upper_curve=upper_curve,
            lower_curve=lower_curve,
        )

    def _section_surface_curve(
        self,
        airfoil: Airfoil,
        station: EvaluatedStation,
        chordwise_samples: FloatArray,
        *,
        upper: bool,
    ) -> FloatArray:
        """
        Build one 3D section curve in the canonical wing frame.

        Parameters
        ----------
        airfoil : Airfoil
            Runtime airfoil object used to evaluate the section shape.
        station : EvaluatedStation
            Evaluated panel properties at the sampled span station.
        chordwise_samples : FloatArray
            One-dimensional normalized chordwise coordinates in ``[0, 1]``.
        upper : bool
            If ``True``, build the upper-surface curve.

        Returns
        -------
        FloatArray
            Array of shape ``(n_samples, 3)`` containing section points.
        """
        t_surface = airfoil.t_from_x(chordwise_samples, upper=upper)
        x_airfoil, z_airfoil = airfoil.xy(t_surface)
        axis_fraction = self._reference_axis_fraction(
            self._spec.wing.reference_axis
        )
        local_x = station.chord * (x_airfoil - axis_fraction)
        local_z = station.chord * z_airfoil
        cos_theta = as_float_scalar(np.cos(station.twist_rad))
        sin_theta = as_float_scalar(np.sin(station.twist_rad))
        x_rotated = cos_theta * local_x + sin_theta * local_z
        z_rotated = -sin_theta * local_x + cos_theta * local_z
        curve = np.empty((chordwise_samples.size, 3), dtype=np.float64)
        curve[:, 0] = station.x_ref + x_rotated
        curve[:, 1] = station.y
        curve[:, 2] = station.z_ref + z_rotated
        return curve

    def _get_panel_airfoil(self, panel: PanelSpec) -> Airfoil:
        """
        Resolve the airfoil definition associated with ``panel``.

        Parameters
        ----------
        panel : PanelSpec
            Wing panel whose referenced airfoil should be resolved.

        Returns
        -------
        Airfoil
            Runtime airfoil object for the panel.

        Raises
        ------
        KeyError
            If the panel references an unknown named airfoil.
        NotImplementedError
            If the panel uses an unsupported airfoil reference or airfoil
            type.
        """
        if not isinstance(panel.airfoil, SingleAirfoilRef):
            msg = (
                "Only single-airfoil panel definitions are supported by "
                "WingCanonical."
            )
            raise NotImplementedError(msg)

        try:
            airfoil_spec = self._spec.airfoils[panel.airfoil.name]
        except KeyError as exc:
            msg = f"Unknown airfoil reference: {panel.airfoil.name}"
            raise KeyError(msg) from exc

        if isinstance(airfoil_spec, DatAirfoilSpec):
            msg = "DAT airfoils are not yet supported by WingCanonical."
            raise NotImplementedError(msg)

        try:
            return AirfoilFactory.from_spec(airfoil_spec)
        except NotImplementedError as exc:
            msg = (
                "Only analytic NACA airfoils supported by AirfoilFactory are "
                "currently supported by WingCanonical."
            )
            raise NotImplementedError(msg) from exc

    def _validate_axes(self) -> None:
        """
        Reject unsupported reference-axis and twist-axis combinations.

        Returns
        -------
        None
            This method validates the stored wing specification in place.

        Raises
        ------
        ValueError
            If the configured axes are unsupported by the canonical
            evaluator.
        """
        reference_axis = self._spec.wing.reference_axis
        twist_axis = self._spec.wing.twist_axis
        if reference_axis == "elastic_axis" or twist_axis == "elastic_axis":
            msg = (
                "elastic_axis is not yet supported because the schema does "
                "not define its chordwise location."
            )
            raise ValueError(msg)
        if reference_axis != twist_axis:
            msg = (
                "WingCanonical currently requires reference_axis and "
                "twist_axis to match."
            )
            raise ValueError(msg)

    def _get_panel(self, panel_id: str) -> PanelSpec:
        """
        Return the panel definition for ``panel_id``.

        Parameters
        ----------
        panel_id : str
            Identifier of the panel to retrieve.

        Returns
        -------
        PanelSpec
            Panel definition matching ``panel_id``.

        Raises
        ------
        KeyError
            If ``panel_id`` does not match a known panel.
        """
        try:
            return self._panel_map[panel_id]
        except KeyError as exc:
            msg = f"Unknown panel id: {panel_id}"
            raise KeyError(msg) from exc

    @staticmethod
    def _validate_eta(panel: PanelSpec, eta: FloatScalar) -> None:
        """
        Ensure ``eta`` lies within the panel spanwise interval.

        Parameters
        ----------
        panel : PanelSpec
            Panel whose spanwise interval should contain ``eta``.
        eta : FloatScalar
            Spanwise parameter to validate.

        Returns
        -------
        None
            This method validates the inputs in place.

        Raises
        ------
        ValueError
            If ``eta`` lies outside ``panel.eta_range``.
        """
        eta0, eta1 = panel.eta_range
        if eta < eta0 or eta > eta1:
            msg = (
                f"eta={eta} lies outside panel {panel.id!r} range "
                f"[{eta0}, {eta1}]."
            )
            raise ValueError(msg)

    @staticmethod
    def _evaluate_piecewise_linear(
        distribution: PiecewiseLinearDistribution,
        eta: FloatScalar,
    ) -> FloatScalar:
        """
        Interpolate a piecewise-linear scalar distribution at ``eta``.

        Parameters
        ----------
        distribution : PiecewiseLinearDistribution
            Piecewise-linear distribution definition to evaluate.
        eta : float
            Spanwise parameter at which to evaluate the distribution.

        Returns
        -------
        float
            Interpolated distribution value at ``eta``.

        Raises
        ------
        ValueError
            If the distribution data are empty, non-monotonic, or do not
            bracket ``eta``.
        """
        if not distribution.data:
            msg = (
                "Piecewise-linear distributions must contain at least one "
                "point."
            )
            raise ValueError(msg)

        eta_values = np.array(
            [point[0] for point in distribution.data],
            dtype=np.float64,
        )
        data_values = np.array(
            [point[1] for point in distribution.data],
            dtype=np.float64,
        )
        deltas = np.diff(eta_values)
        if np.any(deltas <= 0.0):
            msg = "Piecewise-linear distribution eta values must be increasing."
            raise ValueError(msg)
        eta_min = as_float_scalar(eta_values[0])
        eta_max = as_float_scalar(eta_values[-1])
        if eta < eta_min or eta > eta_max:
            msg = (
                f"eta={eta} lies outside distribution range "
                f"[{eta_min}, {eta_max}]."
            )
            raise ValueError(msg)
        return as_float_scalar(np.interp(eta, eta_values, data_values))

    def _angle_to_radians(self, angle_value: FloatScalar) -> FloatScalar:
        """
        Convert an input angle to radians using the schema units.

        Parameters
        ----------
        angle_value : FloatScalar
            Angle value expressed in the schema-declared units.

        Returns
        -------
        FloatScalar
            Angle converted to radians.
        """
        if self._spec.units.angle == "rad":
            return angle_value
        return as_float_scalar(np.deg2rad(angle_value))

    @staticmethod
    def _reference_axis_fraction(axis: ReferenceAxis) -> float:
        """
        Return the chordwise fraction associated with a reference axis.

        Parameters
        ----------
        axis : ReferenceAxis
            Reference-axis label from the wing schema.

        Returns
        -------
        float
            Chordwise fraction associated with ``axis``.

        Raises
        ------
        ValueError
            If ``axis`` is unsupported by the canonical evaluator.
        """
        if axis == "leading_edge":
            return 0.0
        if axis == "quarter_chord":
            return 0.25
        msg = f"Unsupported reference axis: {axis}"
        raise ValueError(msg)


__all__ = [
    "EvaluatedStation",
    "SectionCurves",
    "WingCanonical",
]
