from dataclasses import dataclass, field
from pathlib import Path
from typing import List, Literal, Mapping, Union
import matplotlib.pyplot as plt
from divergent_gravity.classes.dg import DG
from divergent_gravity.functionality.flattened_rotation_curves.galaxy import (
    GalaxyDataset,
)
import pandas as pd
import numpy as np
from icecream import ic
from tabulate import tabulate
import scipy.constants as spc
import plotly.graph_objects as go
import plotly.io as pio
import scienceplots

plt.style.use(["science", "no-latex"])

from divergent_gravity.functionality.flattened_rotation_curves.plot_saver import (
    PlotSaver,
)
from divergent_gravity.functionality.flattened_rotation_curves.sparc_validator import (
    SPARCValidator,
    generate_residual_map,
)

pio.templates.default = "plotly_dark"
from divergent_gravity.functionality.flattened_rotation_curves.sparc_handler import (
    SparcHandler,
)

ic.enable()


@dataclass
class Plotter:
    dg: DG
    """
    The core **Divergent Gravity** model. Feel free to play with the parameters 
    as you like.
    """
    cache_dir: Union[str, None] = None
    """
    Defaults to `<data directory>/divergent_gravity` for your specific file system.
    """
    spark: SparcHandler = field(init=False)
    """
    data_sets is a list of strings indicating the galaxies that should be plotted.
    Each string should match a catalog number, meaning a file name like:
    DDO161_rotmod.dat
    becomes DDO161.
    """
    data_sets: Union[List[str], Literal["All"]] = field(default_factory=lambda: "All")
    """
    A list of catalog names in the SPARC database, or "All". If "All",
    this will read in all 175 galaxies in the SPARC datbase.
    Other options include manually entering them, or using:
    ```py
    data_sets=SparcHandler.get_ai_selected_high_quality_datasets()
    ```
    or:
    ```py
    sparc = SparcHandler()
    data_sets=sparc.get_random_datasets(50)
    ```
    To get a list of all supported catalog names, use:
    ```py
    sparc = SparcHandler()
    print(sparc.all_catalog_names())
    ```
    Defaults to "All".
    """

    base_upsilon_disk: float = 0.5
    """
    This is the mass-to-light ratio applied by default. If `fine_tuning_delta != 0`, this will 
    begin a walk up and down between `fine_tuning_min` and `fine_tuning_max` by the `fine_tuning_delta` you provided.
    """
    base_upsilon_bulge: float = 0.7
    """
    This is the mass-to-light ratio applied by default. If `fine_tuning_delta != 0`, this will 
    begin a walk up and down between `fine_tuning_min` and `fine_tuning_max` by the `fine_tuning_delta` you provided.
    """

    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
    """
    `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.
    """

    df_results: pd.DataFrame = field(
        default_factory=lambda: pd.DataFrame(
            columns=[
                "Name",
                "Inclination",
                "Mass",
                "Rad (mean)",
                "Vobs (mean)",
                "Vgas (mean)",
                "Vdisk (mean)",
                "SBdisk (mean)",
                "SBbul (mean)",
                "ΛCDM MAPE",
                "Best MAPE",
                "Best MAPE ϒ (bulge)",
                "Best MAPE ϒ (disk)",
            ]  # type: ignore
        )
    )

    def __post_init__(self):
        if self.cache_dir is not None:
            self.sparc = SparcHandler(output_dir=Path(self.cache_dir))
        else:
            self.sparc = SparcHandler()
        if self.data_sets == "All":
            self.data_sets = self.sparc.all_catalog_names()

    def get_galaxy_dataset(self, catalog_name: str) -> GalaxyDataset:
        return GalaxyDataset(
            catalog_name=catalog_name,
            sparc=self.sparc,
            base_upsilon_bulge=self.base_upsilon_bulge,
            base_upsilon_disk=self.base_upsilon_disk,
            fine_tuning_delta=self.fine_tuning_delta,
            fine_tuning_max=self.fine_tuning_max,
            fine_tuning_min=self.fine_tuning_min,
            kappa=self.dg.kappa,
        )

    def __get_mape_dict(self, gds: GalaxyDataset):
        df = gds.to_dataframe()
        res = gds.calculate_galaxy_mape(
            self.dg, self.base_upsilon_disk, self.base_upsilon_bulge
        )
        if res["plotType"] == "surface":
            return (
                {
                    "Name": gds.catalog_name,
                    "Inclination": gds.inclination,
                    "Mass": gds.mass,
                    "Rad (mean)": df["Rad"].mean(),
                    "Vobs (mean)": df["Vobs"].mean(),
                    "Vgas (mean)": df["Vgas"].mean(),
                    "Vdisk (mean)": df["Vdisk"].mean(),
                    "SBdisk (mean)": df["SBdisk"].mean(),
                    "SBbul (mean)": df["SBbul"].mean(),
                    "ΛCDM MAPE": res["lamcdm_mape"],
                    "Best MAPE": res["best_mape"],
                    "Best MAPE ϒ (bulge)": res["best_mape_upsilon_bulge"],
                    "Best MAPE ϒ (disk)": res["best_mape_upsilon_disk"],
                },
                res["lam_mape_data"],
                res["dg_mape_data"],
                res["upsilon_disks"],
                res["upsilon_bulges"],
            )
        else:
            return (
                {
                    "Name": gds.catalog_name,
                    "Inclination": gds.inclination,
                    "Mass": gds.mass,
                    "Rad (mean)": df["Rad"].mean(),
                    "Vobs (mean)": df["Vobs"].mean(),
                    "Vgas (mean)": df["Vgas"].mean(),
                    "Vdisk (mean)": df["Vdisk"].mean(),
                    "SBdisk (mean)": df["SBdisk"].mean(),
                    "SBbul (mean)": df["SBbul"].mean(),
                    "ΛCDM MAPE": res["lamcdm_mape"],
                    "Best MAPE": res["best_mape"],
                    "Best MAPE ϒ (bulge)": res["best_mape_upsilon_bulge"],
                    "Best MAPE ϒ (disk)": res["best_mape_upsilon_disk"],
                },
                None,
                None,
                None,
                None,
            )

    def __calculate_mape(self, gds: GalaxyDataset):
        exists_in_column = (self.df_results["Name"] == gds.catalog_name).any()
        if exists_in_column:
            return
        res, _, _, _, _ = self.__get_mape_dict(gds)
        self.df_results.loc[len(self.df_results)] = res

    def calculate_mape(self) -> pd.DataFrame:
        for ds in self.data_sets:
            gds = self.get_galaxy_dataset(ds)
            self.__calculate_mape(gds)

        self.df_results = self.df_results.sort_values(by=["Best MAPE"])
        return self.df_results

    def plot(self, print_data: bool = True, save_props: Union[PlotSaver, None] = None):
        for ds in self.data_sets:
            gal = self.get_galaxy_dataset(ds)
            mape_dict, _, _, _, _ = self.__get_mape_dict(gal)
            if print_data:
                dict_keys = mape_dict.keys()
                table_data = list(zip(*[[mape_dict[s]] for s in dict_keys]))
                print(tabulate(table_data, headers=dict_keys, tablefmt="grid"))  # type: ignore
            gal.plot(self.dg, save_props)

    def __formt_column_name(self, col: str) -> str:
        return col.replace("Lam", r"$\Lambda$").replace("Upsilon", r"$\Upsilon$")

    def plot_database_mape(
        self, x_axis_column: str = "Mass", save_props: Union[PlotSaver, None] = None
    ):
        """
        x_axis_column is a column key in the df_results dataframe.
        x_axis_column only works when the fine_tuning_delta=0, since that
        generates a line plot. When delta != 0, a surface plot is generated.
        Defaults to "Mass"
        Valid Options:
            - Name
            - Inclination
            - Mass
            - Rad (mean)
            - Vobs (mean)
            - Vgas (mean)
            - Vdisk (mean)
            - SBdisk (mean)
            - SBbul (mean)
            - ΛCDM MAPE
            - Best MAPE
            - Best MAPE ϒ (bulge)
            - Best MAPE ϒ (disk)
        """
        self.calculate_mape()
        if self.fine_tuning_delta == 0:
            plt.figure(figsize=(10, 6))
            for col in [
                "ΛCDM MAPE",
                "Best MAPE",
                "Best MAPE ϒ (disk)",
                "Best MAPE ϒ (bulge)",
            ]:
                if col != x_axis_column:
                    plt.scatter(
                        self.df_results[x_axis_column],
                        self.df_results[col],
                        label=self.__formt_column_name(col),
                    ),
            plt.title(
                f"Mean Absolute Percentage Error across SPARC database",
                fontsize=14,
            )
            plt.xlabel(x_axis_column)
            plt.ylabel("MAPE")
            plt.legend(loc="lower right")
            if save_props is not None:
                plt.savefig(save_props.to_path("mape"))
            plt.show()
        else:
            fig = go.Figure(
                layout=dict(
                    title=dict(text="Fine-tuned MAPE across SPARC database"),
                    legend={
                        "x": 0.9,
                        "y": 0.9,
                        "xref": "container",
                        "yref": "container",
                    },
                ),
            )
            have_shown_plot = False
            for ds in self.data_sets:
                gds = self.get_galaxy_dataset(ds)
                (
                    _,
                    lam_mape_data,
                    gd_mape_data,
                    upsilon_disks,
                    upsilon_bulges,
                ) = self.__get_mape_dict(gds)
                fig.add_trace(
                    go.Surface(
                        x=upsilon_disks,
                        y=upsilon_bulges,
                        z=lam_mape_data,
                        name=f"{ds}: ΛCDM MAPE",
                        showscale=False,
                        colorscale="ice",
                        visible="legendonly" if have_shown_plot else None,
                    )
                )
                fig.add_trace(
                    go.Surface(
                        x=upsilon_disks,
                        y=upsilon_bulges,
                        z=gd_mape_data,
                        name=f"{ds}: DG MAPE",
                        colorscale="Agsunset",
                        showscale=False,
                        visible="legendonly" if have_shown_plot else None,
                    )
                )
                have_shown_plot = True
            fig.update_layout(
                scene=dict(
                    xaxis=dict(nticks=4, title=dict(text="ϒ (disk)")),
                    yaxis=dict(nticks=4, title=dict(text="ϒ (bulge)")),
                    zaxis=dict(nticks=10, title=dict(text="MAPE")),
                ),
            )
            fig.update_traces(showlegend=True)
            if save_props is not None:
                fig.to_image(save_props.to_path("mape_tuned"))
            fig.show(renderer="browser")

    def to_galaxy_dataset_list(self) -> List[GalaxyDataset]:
        items = []
        for x in self.data_sets:
            res = self.get_galaxy_dataset(x)
            res.to_dataframe()  # NOTE: Required to read from file so user's aren't left with an incomplete dataset.
            items.append(res)
        return items

    def to_data_map(self) -> Mapping[str, pd.DataFrame]:
        data = {}
        for item in self.to_galaxy_dataset_list():
            data[item.catalog_name] = item.to_dataframe()
        return data

    def plot_brightness_comparison(self):
        """
        Creates a scatter plot with the x-axis set
        to a derived surface brighness:
        surface_brightness = np.max(df["Vdisk"])**2 / df["Rad"][np.argmax(df["Vdisk"])]
        y-axis is the derived MAPE.
        """
        data = {}
        for gal in self.data_sets:
            gds = self.get_galaxy_dataset(gal)
            data[gds.catalog_name] = gds.to_dataframe()
        generate_residual_map(SPARCValidator(data), data)

    def plot_machian_normalization(
        self, ml: float = 0.5, save_props: Union[PlotSaver, None] = None
    ):
        """
        *Dotted lines*: Theoretical convergence point at `sqrt(2)` which gives rise to something
        similar to a phase shift of space itself. This is why t' = β when v = c/sqrt(2).
        *Solid Lines*: sqrt(1 + 1 / g), the predicted gain.
        *Faint Dots*: The gain observed for each galaxy in the dataset.
        """
        sparc_data = self.to_data_map()
        kappa = 1 / np.pi
        A0 = self.dg.A0
        pivot_threshold = A0 * kappa

        plt.figure(figsize=(10, 8))

        nudge_to_avoid_zero = 1e-10
        for g_id, g in sparc_data.items():
            # 1. Calculate Newtonian Accel (gn)
            # Using a fixed M/L for this visualization
            v_bar_sq = np.abs(g["Vgas"]) * g["Vgas"] + g["Vdisk"] ** 2 * ml
            v_bar = np.sqrt(np.maximum(v_bar_sq, 0))
            r_m = g["Rad"] * spc.parsec * 1000
            gn = (v_bar * 1000.0) ** 2 / (r_m + nudge_to_avoid_zero)

            # 2. Normalize by the Pivot Point
            g_star = gn / pivot_threshold

            # 3. Calculate the Observed Gain
            # Ratio of Observed Velocity to Newtonian Velocity
            v_obs_ms = g["Vobs"] * 1000.0
            v_newt_ms = v_bar * 1000.0
            gain_obs = (
                v_obs_ms / (v_newt_ms + nudge_to_avoid_zero)
            ) ** 2  # Squared to match acceleration ratio

            plt.scatter(g_star, gain_obs, alpha=0.05, s=2, color="green")

        # 4. Plot the Theoretical Divergent Identity: sqrt(1 + 1/g*)
        g_range = np.logspace(-2, 2, 500)
        theoretical_gain = np.sqrt(1 + 1 / g_range)

        plt.plot(
            g_range,
            theoretical_gain,
            "r-",
            label=r"Divergent Identity: $\sqrt{1 + 1/g^*}$",
        )
        plt.axvline(1, color="blue", linestyle="--", label=r"Pivot Point ($\sqrt{2}$)")
        plt.axhline(np.sqrt(2), color="blue", linestyle="--")

        plt.xscale("log")
        plt.yscale("log")
        plt.ylim(-10e1, 10e2)
        plt.xlabel(r"Newtonian Acceleration ($g_n / A_0\kappa$)")
        plt.ylabel(r"Machian Gain ($g_{obs} / g_n$)")
        plt.title("The Universal Pivot: Convergence of 175 Galaxies")
        plt.legend()
        plt.grid(True, which="both", alpha=0.3)
        if save_props is not None:
            plt.savefig(save_props.to_path("machian_gain"))
        plt.show()

    def describe_results(self):
        return self.calculate_mape().describe(include="all")
