"""
This file and class is exclusively used for NeuralFoil training, not inference.

It is kept here so that the training routines and the inference routines dynamically import from the same
file. This can be used to later change various training parameters (e.g., the resolution of the Kulfan
parameterization, or the spacing of x-points for boundary layer resolution), and have the inference scripts
automatically adjust to that.

End users need not worry about this file.
"""

import warnings
import aerosandbox as asb
import aerosandbox.numpy as np
from dataclasses import dataclass, field
from typing import Union, Sequence, List, Any, Dict
from scipy import interpolate


def compute_optimal_x_points(
    n_points,
):
    # theta = np.arccos(1 - 2 * x)
    # delta_Cp_thin_airfoil_theory = 4 / np.tan(theta)
    # desired_spacing = 1 / delta_Cp_thin_airfoil_theory
    # spacing_function = integral(desired_spacing, x) # rescaled to (0, 1)

    # thin_airfoil_theory_spacing_function = lambda x: (
    #                                                          -np.sqrt((1 - x) * x)
    #                                                          - 0.5 * np.arcsin(1 - 2 * x)
    #                                                  ) / (np.pi / 2) + 0.5
    #
    # uniform_spacing_function = lambda x: x
    #
    # return (
    #         thin_airfoil_theory_spacing_function(np.linspace(0, 1, n_points)) +
    #         uniform_spacing_function(np.sinspace(0, 1, n_points))
    # ) / 2
    s = np.linspace(0, 1, n_points + 1)
    return (s[1:] + s[:-1]) / 2


# def compute_optimal_x_points(
#         n_points: int = 24,
# ):
#     # theta = np.arccos(1 - 2 * x)
#     # delta_Cp_thin_airfoil_theory = 4 / np.tan(theta)
#     # desired_spacing = 1 / delta_Cp_thin_airfoil_theory
#     # spacing_function = integral(desired_spacing, x) # rescaled to (0, 1)
#
#     thin_airfoil_theory_spacing_function = lambda x: (
#                                                              -np.sqrt((1 - x) * x)
#                                                              - 0.5 * np.arcsin(1 - 2 * x)
#                                                      ) / (np.pi / 2) + 0.5
#
#     uniform_spacing_function = lambda x: x
#
#     return (
#             thin_airfoil_theory_spacing_function(np.linspace(0, 1, n_points)) +
#             uniform_spacing_function(np.sinspace(0, 1, n_points))
#     ) / 2


@dataclass
class Data:
    airfoil: asb.KulfanAirfoil
    alpha: float
    Re: float
    mach: float
    n_crit: float
    xtr_upper: float
    xtr_lower: float

    N = 32
    bl_x_points = compute_optimal_x_points(n_points=N)

    analysis_confidence: float  # Nominally 0 (no confidence) to 1 (high confidence)
    af_outputs: Dict[str, Any] = field(
        default_factory=lambda: {
            "CL": np.nan,
            "CD": np.nan,
            "CM": np.nan,
            "Top_Xtr": np.nan,
            "Bot_Xtr": np.nan,
        }
    )
    upper_bl_outputs: Dict[str, Any] = field(
        default_factory=lambda: {
            "theta": np.nan * np.ones_like(Data.bl_x_points),
            "H": np.nan * np.ones_like(Data.bl_x_points),
            "ue/vinf": np.nan * np.ones_like(Data.bl_x_points),
        }
    )  # theta, H, ue/vinf
    lower_bl_outputs: Dict[str, Any] = field(
        default_factory=lambda: {
            "theta": np.nan * np.ones_like(Data.bl_x_points),
            "H": np.nan * np.ones_like(Data.bl_x_points),
            "ue/vinf": np.nan * np.ones_like(Data.bl_x_points),
        }
    )  # theta, H, ue/vinf

    @property
    def inputs(self):
        return {
            "airfoil": self.airfoil,
            "alpha": self.alpha,
            "Re": self.Re,
            "mach": self.mach,
            "n_crit": self.n_crit,
            "xtr_upper": self.xtr_upper,
            "xtr_lower": self.xtr_lower,
        }

    @property
    def outputs(self):
        return {
            "analysis_confidence": self.analysis_confidence,
            "af_outputs": self.af_outputs,
            "upper_bl_outputs": self.upper_bl_outputs,
            "lower_bl_outputs": self.lower_bl_outputs,
        }

    @classmethod
    def from_xfoil(
        cls,
        airfoil: Union[asb.Airfoil, asb.KulfanAirfoil],
        alphas: Union[float, Sequence[float]],
        Re: float,
        mach: float,
        n_crit: float,
        xtr_upper: float,
        xtr_lower: float,
        timeout=5,
        max_iter=100,
        xfoil_command: str = "xfoil",
    ) -> List["Data"]:
        airfoil = airfoil.normalize().to_kulfan_airfoil()

        alphas = np.atleast_1d(alphas)

        xf = asb.XFoil(
            airfoil=airfoil,
            Re=Re,
            mach=mach,
            n_crit=n_crit,
            xtr_upper=xtr_upper,
            xtr_lower=xtr_lower,
            xfoil_repanel=True,
            timeout=timeout,
            max_iter=max_iter,
            xfoil_command=xfoil_command,
            include_bl_data=True,
        )
        with warnings.catch_warnings():
            warnings.filterwarnings("ignore", category=UserWarning)
            xf_outputs = xf.alpha(alphas)

        training_datas = []

        def append_empty_data():
            training_datas.append(
                cls(
                    airfoil=airfoil.to_kulfan_airfoil(),
                    alpha=alpha,
                    Re=Re,
                    mach=mach,
                    n_crit=n_crit,
                    xtr_upper=xtr_upper,
                    xtr_lower=xtr_lower,
                    analysis_confidence=0,
                )
            )

        for i, alpha in enumerate(alphas):
            ### Figure out which output corresponds to this alpha
            alpha_deviations = np.abs(xf_outputs["alpha"] - alpha)

            if (len(alpha_deviations) == 0) or (  # If the alpha is not in the output
                np.min(alpha_deviations) > 0.001
            ):
                append_empty_data()
                continue

            index = np.argmin(alpha_deviations)
            xf_output = {key: value[index] for key, value in xf_outputs.items()}
            bl_data = xf_output["bl_data"]

            # ### Trim off the wake data (rows where "Cp" first becomes NaN)
            # wake_node_indices = np.flatnonzero(bl_data["x"] > )
            # bl_data = bl_data.iloc[
            #           :np.flatnonzero(bl_data["x"] <= airfoil.x().max())[-1] + 1,
            #           :
            #           ]

            ### Split the boundary layer data into upper and lower sections
            # LE_index = bl_data["x"].argmin()
            try:
                dx = np.diff(bl_data["x"].values)
                upper_bl_data = bl_data.iloc[: np.flatnonzero(dx < 0)[-1] + 2 :].iloc[
                    ::-1, :
                ]
                lower_bl_data = bl_data.iloc[np.flatnonzero(dx > 0)[0] :, :]
            except IndexError as e:  # If the boundary layer data is too short
                print(e)
                append_empty_data()
                continue

            if (
                len(upper_bl_data) <= 4 or len(lower_bl_data) <= 4
            ):  # If the boundary layer data is too short
                print("BL data too short")
                append_empty_data()
                continue

            def interp(x: np.ndarray, y: np.ndarray) -> np.ndarray:
                return interpolate.PchipInterpolator(x, y, extrapolate=True)(
                    cls.bl_x_points
                )

            try:
                training_datas.append(
                    cls(
                        airfoil=airfoil.to_kulfan_airfoil(),
                        alpha=alpha,
                        Re=Re,
                        mach=mach,
                        n_crit=n_crit,
                        xtr_upper=xtr_upper,
                        xtr_lower=xtr_lower,
                        analysis_confidence=1,
                        af_outputs={
                            "CL": xf_output["CL"],
                            "CD": xf_output["CD"],
                            "CM": xf_output["CM"],
                            "Top_Xtr": xf_output["Top_Xtr"],
                            "Bot_Xtr": xf_output["Bot_Xtr"],
                        },
                        upper_bl_outputs={
                            "theta": interp(upper_bl_data["x"], upper_bl_data["theta"]),
                            "H": interp(upper_bl_data["x"], upper_bl_data["H"]),
                            "ue/vinf": interp(
                                upper_bl_data["x"], upper_bl_data["ue/vinf"]
                            ),
                        },
                        lower_bl_outputs={
                            "theta": interp(lower_bl_data["x"], lower_bl_data["theta"]),
                            "H": interp(lower_bl_data["x"], lower_bl_data["H"]),
                            "ue/vinf": interp(
                                lower_bl_data["x"], lower_bl_data["ue/vinf"]
                            ),
                        },
                    )
                )
            except ValueError as e:
                print(e)
                append_empty_data()
                continue

        return training_datas

    def to_vector(self) -> np.ndarray:  # dtype: float32
        items = [
            self.airfoil.upper_weights,
            self.airfoil.lower_weights,
            self.airfoil.leading_edge_weight,
            self.airfoil.TE_thickness,
            self.alpha,
            self.Re,
            self.mach,
            self.n_crit,
            self.xtr_upper,
            self.xtr_lower,
            self.analysis_confidence,
            self.af_outputs["CL"],
            self.af_outputs["CD"],
            self.af_outputs["CM"],
            self.af_outputs["Top_Xtr"],
            self.af_outputs["Bot_Xtr"],
            self.upper_bl_outputs["theta"],
            self.upper_bl_outputs["H"],
            self.upper_bl_outputs["ue/vinf"],
            self.lower_bl_outputs["theta"],
            self.lower_bl_outputs["H"],
            self.lower_bl_outputs["ue/vinf"],
        ]

        return np.concatenate(
            [np.atleast_1d(item).flatten().astype(np.float32) for item in items], axis=0
        )

    @classmethod
    def from_vector(cls, vector: np.ndarray) -> "Data":
        i = 0

        def pop(n):
            nonlocal i
            result = vector[i : i + n]
            i += n
            return result

        airfoil = asb.KulfanAirfoil(
            upper_weights=pop(8),
            lower_weights=pop(8),
            leading_edge_weight=pop(1),
            TE_thickness=pop(1),
        )
        alpha = pop(1)
        Re = pop(1)
        mach = pop(1)
        n_crit = pop(1)
        xtr_upper = pop(1)
        xtr_lower = pop(1)
        analysis_confidence = pop(1)
        af_outputs = {
            "CL": pop(1),
            "CD": pop(1),
            "CM": pop(1),
            "Top_Xtr": pop(1),
            "Bot_Xtr": pop(1),
        }
        upper_bl_outputs = {
            "theta": pop(cls.N),
            "H": pop(cls.N),
            "ue/vinf": pop(cls.N),
        }
        lower_bl_outputs = {
            "theta": pop(cls.N),
            "H": pop(cls.N),
            "ue/vinf": pop(cls.N),
        }
        if not i == len(vector):
            raise ValueError(
                f"Vector length mismatch: Got {len(vector)}, expected {i}."
            )

        return cls(
            airfoil=airfoil,
            alpha=alpha,
            Re=Re,
            mach=mach,
            n_crit=n_crit,
            xtr_upper=xtr_upper,
            xtr_lower=xtr_lower,
            analysis_confidence=analysis_confidence,
            af_outputs=af_outputs,
            upper_bl_outputs=upper_bl_outputs,
            lower_bl_outputs=lower_bl_outputs,
        )

    @classmethod
    def get_vector_input_column_names(cls):
        return [
            *[f"kulfan_upper_{i}" for i in range(8)],
            *[f"kulfan_lower_{i}" for i in range(8)],
            "kulfan_LE_weight",
            "kulfan_TE_thickness",
            "alpha",
            "Re",
            "mach",
            "n_crit",
            "xtr_upper",
            "xtr_lower",
        ]

    @classmethod
    def get_vector_output_column_names(cls):
        return [
            "analysis_confidence",
            "CL",
            "CD",
            "CM",
            "Top_Xtr",
            "Bot_Xtr",
            *[f"upper_bl_theta_{i}" for i in range(cls.N)],
            *[f"upper_bl_H_{i}" for i in range(cls.N)],
            *[f"upper_bl_ue/vinf_{i}" for i in range(cls.N)],
            *[f"lower_bl_theta_{i}" for i in range(cls.N)],
            *[f"lower_bl_H_{i}" for i in range(cls.N)],
            *[f"lower_bl_ue/vinf_{i}" for i in range(cls.N)],
        ]

    @classmethod
    def get_vector_column_names(cls):
        return (
            cls.get_vector_input_column_names() + cls.get_vector_output_column_names()
        )

    def __eq__(self, other: "Data") -> bool:
        v1 = self.to_vector()
        v2 = other.to_vector()

        return all(
            [
                np.allclose(
                    s1,
                    s2,
                    atol=0,
                    rtol=np.finfo(np.float32).eps * 10,
                )
                or np.all(np.isnan([s1, s2]))
                for s1, s2 in zip(v1, v2)
            ]
        )

        return np.allclose(
            self.to_vector(),
            other.to_vector(),
            atol=0,
            rtol=np.finfo(np.float32).eps * 10,
        )

    def validate_vector_format(self):
        assert self == Data.from_vector(self.to_vector())


if __name__ == "__main__":
    af = asb.Airfoil("ag41d-02r").normalize().to_kulfan_airfoil()

    datas = Data.from_xfoil(
        airfoil=af,
        alphas=[3, 5, 60],
        Re=1e6,
        mach=0,
        n_crit=9,
        xtr_upper=0.99,
        xtr_lower=0.99,
    )

    d = datas[0]

    # import aerosandbox as asb
    # import aerosandbox.numpy as np
    #
    # airfoil_database_path = asb._asb_root / "geometry" / "airfoil" / "airfoil_database"
    #
    # UIUC_airfoils = [
    #     asb.Airfoil(name=filename.stem).normalize().to_kulfan_airfoil()
    #     for filename in airfoil_database_path.glob("*.dat")
    # ]
    #
    # for af in UIUC_airfoils:
    #     print(af.name)
    #     datas = Data.from_xfoil(
    #         airfoil=af,
    #         alphas=[2, 60],
    #         Re=1e6,
    #         mach=0,
    #         n_crit=9,
    #         xtr_upper=0.99,
    #         xtr_lower=0.99,
    #     )
    #
    #     d = datas[0]
    #     # print(d.to_vector())
    #     # print(Data.from_vector(d.to_vector()))
    #     print(d == Data.from_vector(d.to_vector()))
