from dataclasses import dataclass, field
from os.path import dirname
from typing import TYPE_CHECKING, Union
import io
import matplotlib.pyplot as plt
import astropy.units as u
from numpy.typing import NDArray
import pandas as pd
import numpy as np
import json
from pathlib import Path

from scipy.constants import pi
from divergent_gravity.functionality.flattened_rotation_curves.sparc_handler import (
    SparcHandler,
)
from divergent_gravity.classes.dg import DG


KPC_TO_M = 3.086e19
KM_S_TO_M_S = 1000.0


@dataclass
class GalaxyDataset:
    catalog_name: str
    sparc: SparcHandler
    mass: float = field(init=False)
    inclination: float = field(init=False)

    base_upsilon_disk: float = 0.6
    base_upsilon_bulge: float = 0.6

    """
    fine_tuning_max and fine_tuning_min are used to walk Upsilon up
    and down by the fine_tuning_delta looking for a best fit. By default,
    fine_tuning_delta = 0, making this value constant at 0.6.
    """
    fine_tuning_max: float = 1

    """
    fine_tuning_max and fine_tuning_min are used to walk Upsilon up
    and down by the fine_tuning_delta looking for a best fit. By default,
    fine_tuning_delta = 0, making this value constant at 0.6.
    """
    fine_tuning_min: float = 0.3
    """
    fine_tuning_max and fine_tuning_min are used to walk Upsilon up
    and down by the fine_tuning_delta looking for a best fit. By default,
    fine_tuning_delta = 0, making this value constant at 0.6.
    """
    fine_tuning_delta: float = 0
    df_cache: Union[pd.DataFrame, None] = None
    kappa: float = field(default_factory=lambda: 1 / pi)
    use_dynamic_kappa: bool = False

    def __post_init__(self):
        if not self.is_valid():
            print("Invalid catalog name. Cannot continue.")
            exit(1)
        p = Path(__file__).parent.parent.parent / "data" / "sparc_masses.json"
        masses = json.loads(p.read_text())
        self.mass = masses[self.catalog_name]["M_bar"]
        self.inclination = masses[self.catalog_name]["inc"]

    def is_valid(self) -> bool:
        zf = self.sparc.initialize_sparc_data()
        if not zf:
            return False
        try:
            with zf.open(f"{self.catalog_name}_rotmod.dat") as f:
                return True
        except FileNotFoundError:
            return False

    def to_dataframe(self) -> pd.DataFrame:
        if self.df_cache is not None:
            return self.df_cache
        zf = self.sparc.initialize_sparc_data()
        if not zf:
            print("Cannot load data. Cannot continue.")
            exit(1)
        try:
            with zf.open(f"{self.catalog_name}_rotmod.dat") as f:
                lines = []
                initial_lines = f.readlines()
                for l in initial_lines:
                    _l = l.decode("utf-8")
                    if not _l.startswith("#"):
                        lines.append(_l)
                file_object = io.StringIO("\n".join(lines))
                df = pd.read_table(
                    file_object,
                    sep=r"\s+",
                    names=[
                        "Rad",
                        "Vobs",
                        "errV",
                        "Vgas",
                        "Vdisk",
                        "Vbul",
                        "SBdisk",
                        "SBbul",
                    ],
                )
                self.df_cache = df
                return df
        except KeyError:
            print(f"Galaxy {self.catalog_name} not found in the dataset.")
            exit(1)

    def calculate_lam_cdm_mape(self, dg: DG):
        """
        Computes predicted velocity based on Divergent Gravity using raw SI
        conversion to ensure unit-safety and 0.5% precision.
        """
        df = self.to_dataframe()
        # 1. Convert to raw SI floats (meters and meters/second)
        v_gas_ms = (df["Vgas"].values * u.km / u.s).to(u.m / u.s).value
        v_disk_ms = (df["Vdisk"].values * u.km / u.s).to(u.m / u.s).value
        v_bul_ms = (df["Vbul"].values * u.km / u.s).to(u.m / u.s).value

        # 2. Compute Baryonic Potential (v_bar^2) in SI
        v_bar_sq = v_gas_ms**2 + v_disk_ms**2 + v_bul_ms**2

        # 4. Return as Astropy Quantities for standard plotting
        v_bar_final = (np.sqrt(v_bar_sq) * u.m / u.s).to(u.km / u.s)

        return v_bar_final

    def get_dg_velocities(
        self,
        dg: DG,
        upsilon_disk: Union[float, None] = None,
        upsilon_bulge: Union[float, None] = None,
    ):
        """
        Calculates the predicted velocities for the DG model.
        Returns:
            v_dg_obs: The velocity projected to the observed plane (for plotting vs Vobs).
            v_dg_true: The actual circular velocity (the physical model).
            v_newton_obs: The baryonic-only velocity projected (for baseline comparison).
        """
        if upsilon_disk is None:
            upsilon_disk = self.base_upsilon_disk
        if upsilon_bulge is None:
            upsilon_bulge = self.base_upsilon_bulge
        df = self.to_dataframe()
        kpc_to_m = 3.086e19
        km_s_to_m_s = 1000.0

        r_meters = df["Rad"].values * kpc_to_m
        sin_i = np.sin(np.radians(self.inclination))
        kappa = (
            self.kappa
            if not self.use_dynamic_kappa
            else self.calculate_dynamic_kappa(upsilon_disk, upsilon_bulge)
        )

        # 1. Newtonian Baryonic sum with Upsilon scaling
        v_bar_sq_si = (
            df["Vgas"] ** 2
            + upsilon_disk * df["Vdisk"] ** 2
            + upsilon_bulge * df["Vbul"] ** 2
        ).values * (km_s_to_m_s**2)

        # 2. DG Acceleration Identity: g_dg = sqrt(gn^2 + gn * a0)
        g_n = v_bar_sq_si / r_meters
        g_dg = np.sqrt(g_n**2 + g_n * dg.A0 * kappa)

        # 3. Convert accelerations back to velocities (km/s)
        v_dg_true_km_s = np.sqrt(g_dg * r_meters) / km_s_to_m_s
        v_newton_true_km_s = np.sqrt(v_bar_sq_si) / km_s_to_m_s

        # 4. Project to the observed plane (Multiply by sin(i))
        # This is critical so the plot matches the Vobs data points
        v_dg_obs = v_dg_true_km_s * sin_i
        v_newton_obs = v_newton_true_km_s * sin_i

        return v_dg_obs, v_dg_true_km_s, v_newton_obs

    def calculate_mape_from_velocities(self, v_predicted_obs: np.ndarray):
        """
        Calculates the Mean Absolute Percentage Error against SPARC observations.
        """
        df = self.to_dataframe()
        v_obs = df["Vobs"].values

        # Mask out zero or invalid observations to prevent division by zero
        mask = (v_obs > 0) & (~np.isnan(v_obs))

        actual = v_obs[mask]
        predicted = v_predicted_obs[mask]

        mape = np.mean(np.abs((actual - predicted) / actual)) * 100
        return mape

    def get_mape_dict(
        self,
        dg: DG,
        upsilon_disk: Union[float, None] = None,
        upsilon_bulge: Union[float, None] = None,
    ):
        if upsilon_disk is None:
            upsilon_disk = self.base_upsilon_disk
        if upsilon_bulge is None:
            upsilon_bulge = self.base_upsilon_bulge
        df = self.to_dataframe()
        kpc_to_m = 3.086e19
        km_s_to_m_s = 1000.0
        kappa = (
            self.kappa
            if not self.use_dynamic_kappa
            else self.calculate_dynamic_kappa(upsilon_disk, upsilon_bulge)
        )

        r_meters = df["Rad"].values * kpc_to_m
        v_obs_km_s = df["Vobs"].values

        # 1. Baryonic sum with Upsilon scaling
        v_bar_sq_si = (
            df["Vgas"] ** 2
            + upsilon_disk * df["Vdisk"] ** 2
            + upsilon_bulge * df["Vbul"] ** 2
        ).values * (km_s_to_m_s**2)

        # 2. Core DG Identity: Acceleration based
        g_n = v_bar_sq_si / r_meters
        g_dg = np.sqrt(g_n**2 + g_n * dg.A0 * kappa)

        # 3. Convert back to velocity and PROJECT to observed plane
        v_dg_true_km_s = np.sqrt(g_dg * r_meters) / km_s_to_m_s

        # FIX: Account for inclination
        sin_i = np.sin(np.radians(self.inclination))
        v_dg_predicted_observable = v_dg_true_km_s * sin_i

        # 4. MAPE Calculation
        mask = v_obs_km_s > 0
        mape = (
            np.mean(
                np.abs(
                    (v_obs_km_s[mask] - v_dg_predicted_observable[mask])
                    / v_obs_km_s[mask]
                )
            )
            * 100
        )

        # Newtonian comparison (Baryons only)
        v_newton_obs = (np.sqrt(v_bar_sq_si) / km_s_to_m_s) * sin_i
        newton_mape = (
            np.mean(np.abs((v_obs_km_s[mask] - v_newton_obs[mask]) / v_obs_km_s[mask]))
            * 100
        )

        return mape, newton_mape

    def calculate_galaxy_mape(
        self,
        dg: DG,
        upsilon_disk: Union[float, None] = None,
        upsilon_bulge: Union[float, None] = None,
        best_mape: Union[float, None] = None,
        best_mape_upsilon_disk=-1.0,
        best_mape_upsilon_bulge=-1.0,
    ) -> dict:
        if upsilon_disk is None:
            upsilon_disk = self.base_upsilon_disk
        if upsilon_bulge is None:
            upsilon_bulge = self.base_upsilon_bulge

        mape, lamcdm_mape = self.get_mape_dict(dg, upsilon_disk, upsilon_bulge)

        _best_mape = mape if (best_mape is None or mape < best_mape) else best_mape
        _best_mape_upsilon_disk = (
            upsilon_disk
            if (best_mape is None or mape < best_mape)
            else best_mape_upsilon_disk
        )
        _best_mape_upsilon_bulge = (
            upsilon_bulge
            if (best_mape is None or mape < best_mape)
            else best_mape_upsilon_bulge
        )

        if self.fine_tuning_delta == 0:
            # There is no delta, so don't walk at all.
            return {
                "best_mape": _best_mape,
                "best_mape_upsilon_disk": _best_mape_upsilon_disk,
                "best_mape_upsilon_bulge": _best_mape_upsilon_bulge,
                "lamcdm_mape": lamcdm_mape,
                "plotType": "line",
            }

        dg_mapes = []
        lam_mapes = []
        upsilon_disks = np.arange(
            self.fine_tuning_min, self.fine_tuning_max, self.fine_tuning_delta
        )
        upsilon_bulges = upsilon_disks

        for disk in upsilon_disks:
            mape_row = []
            lam_mape_row = []
            for bulge in upsilon_bulges:
                mape, lamcdm_mape = self.get_mape_dict(dg, disk, bulge)
                mape_row.append(mape)
                lam_mape_row.append(lamcdm_mape)
            dg_mapes.append(mape_row)
            lam_mapes.append(lam_mape_row)

        return {
            "best_mape": _best_mape,
            "best_mape_upsilon_disk": _best_mape_upsilon_disk,
            "best_mape_upsilon_bulge": _best_mape_upsilon_bulge,
            "lamcdm_mape": lamcdm_mape,
            "lam_mape_data": lam_mapes,
            "dg_mape_data": dg_mapes,
            "upsilon_disks": upsilon_disks,
            "upsilon_bulges": upsilon_bulges,
            "plotType": "surface",
        }

    def plot(self, dg: DG, output_dir: Union[str, Path, None] = None):
        df = self.to_dataframe()
        v_dg_obs, v_dg_true, v_newton_obs = self.get_dg_velocities(
            dg, self.base_upsilon_disk, self.base_upsilon_bulge
        )
        v_dg_projected = v_dg_true * np.sin(np.radians(self.inclination))

        # Calculate the error for the title
        current_mape = self.calculate_mape_from_velocities(v_dg_obs)

        plt.figure(figsize=(10, 6))
        plt.errorbar(
            df["Rad"],
            df["Vobs"],
            yerr=df["errV"],
            fmt="ko",
            markersize=4,
            label="SPARC Observed Data",
            alpha=0.7,
        )

        # 2. Newtonian Baseline (Projected)
        plt.plot(df["Rad"], v_newton_obs, "g--", label="Baryonic (Newtonian Baseline)")

        # 3. DG Prediction (Projected)
        plt.plot(
            df["Rad"],
            v_dg_obs,
            "b-",
            linewidth=2,
            label=f"Divergent Gravity (MAPE: {current_mape:.2f}%)",
        )

        plt.title(f"Divergent Gravity Rotation Curve: {self.catalog_name}", fontsize=14)
        plt.xlabel("Radius (kpc)")
        plt.ylabel("Velocity (km/s)")
        plt.legend(loc="lower right")
        plt.grid(True, linestyle=":", alpha=0.5)
        if output_dir is not None:
            plt.savefig(Path(output_dir) / f"{self.catalog_name}.png")
        plt.show()
