from dataclasses import dataclass, field
from pathlib import Path
from typing import List, Literal, Union
import io
import numpy as np
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
from icecream import ic
from tabulate import tabulate
import plotly.graph_objects as go
import plotly.io as pio

pio.templates.default = "plotly_dark"

from divergent_gravity.functionality.flattened_rotation_curves.sparc_handler import (
    SparcHandler,
)

ic.enable()


@dataclass
class RotationCurvePlotter:
    dg: DG
    cache_dir: Union[str, None] = None
    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")

    base_upsilon_disk: float = 0.5
    base_upsilon_bulge: float = 0.7

    """
    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_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)",
            ]
        )
    )

    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 if self.dg.kappa != "dynamic" else None,
        )

    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, output_dir: Union[str, Path, 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"))
            gal.plot(self.dg, output_dir)

    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"):
        """
        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")
            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",
                        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",
                        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")),
                ),
                coloraxis_colorbar=dict(
                    title=dict(text="MAPE", side="top"),
                    tickmode="array",
                    ticks="outside",
                    nticks=5,
                ),
            )
            fig.update_traces(showlegend=True)
            fig.show(renderer="browser")

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