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 dg.classes.dg import DG
from dg.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 dg.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,
        )

    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")
