"""General Statistics shared across multiple dashboard pages.

- Dataset statistics: Feature Counts, Class Distributions, Temporal Coverage, all per grid-level-combination
"""

import pickle
from collections import defaultdict
from dataclasses import asdict, dataclass
from typing import Literal

import pandas as pd
import streamlit as st
import xarray as xr
from stopuhr import stopwatch

import entropice.spatial.grids
import entropice.utils.paths
from entropice.dashboard.utils.loaders import TrainingResult
from entropice.ml.dataset import DatasetEnsemble
from entropice.utils.types import (
    GridLevel,
    L2SourceDataset,
    TargetDataset,
    Task,
    TemporalMode,
    all_target_datasets,
    all_tasks,
    all_temporal_modes,
    grid_configs,
)


@dataclass(frozen=True)
class MemberStatistics:
    """Statistics for a specific dataset member."""

    feature_count: int  # Number of features from this member
    variable_names: list[str]  # Names of variables from this member
    dimensions: dict[str, int]  # Dimension name to size mapping
    coordinates: list[str]  # Names of coordinates from this member
    size_bytes: int  # Size of this member's data on disk in bytes

    @classmethod
    def compute(
        cls,
        e: DatasetEnsemble,
        member_datasets: dict[L2SourceDataset, xr.Dataset] | None = None,
    ) -> dict[L2SourceDataset, "MemberStatistics"]:
        """Pre-compute the statistics for a specific dataset member."""
        member_datasets = member_datasets or {}
        member_stats = {}
        for member in e.members:
            if member in member_datasets:
                ds = member_datasets[member]
            else:
                ds = e.read_member(member, lazy=True)
            size_bytes = ds.nbytes

            n_cols_member = len(ds.data_vars)
            for dim in ds.sizes:
                if dim != "cell_ids":
                    n_cols_member *= ds.sizes[dim]

            member_stats[member] = cls(
                feature_count=n_cols_member,
                variable_names=list(ds.data_vars),
                dimensions=dict(ds.sizes),
                coordinates=list(ds.coords),
                size_bytes=size_bytes,
            )
        return member_stats


@dataclass(frozen=True)
class TargetStatistics:
    """Statistics for a specific target dataset."""

    training_cells: int  # Number of cells used for training
    coverage: float  # Percentage of total cells covered by training data
    class_counts: dict[str, int] | None  # Class name to count mapping
    class_distribution: dict[str, float] | None  # Class name to percentage mapping

    @classmethod
    def compute(cls, e: DatasetEnsemble, target: TargetDataset, total_cells: int) -> dict[Task, "TargetStatistics"]:
        """Pre-compute the statistics for a specific target dataset."""
        target_stats = {}
        for task in all_tasks:
            targets = e.get_targets(target=target, task=task)

            training_cells = len(targets)
            training_coverage = len(targets) / total_cells * 100

            if task in ["count", "density"]:
                class_counts = None
                class_distribution = None
            else:
                assert targets["y"].dtype == "category", "Classification tasks must have categorical target dtype."
                counts = targets["y"].value_counts()
                distribution = counts / counts.sum() * 100
                class_counts = counts.to_dict()
                class_distribution = distribution.to_dict()
            target_stats[task] = cls(
                training_cells=training_cells,
                coverage=training_coverage,
                class_counts=class_counts,
                class_distribution=class_distribution,
            )
        return target_stats


@dataclass(frozen=True)
class DatasetStatistics:
    """Statistics for a potential dataset at a specific grid and level.

    These statistics are meant to be easily compute without loading a full dataset into memory.
    Further, it is meant to give an overview of all potential potential dataset compositions for a given grid and level.
    """

    total_features: int  # Total number of features available in the dataset
    total_cells: int  # Total number of grid cells potentially covered
    size_bytes: int  # Size of the dataset on disk in bytes
    members: dict[L2SourceDataset, MemberStatistics]  # Statistics per source dataset member
    target: dict[TargetDataset, dict[Task, TargetStatistics]]  # Statistics per target dataset and Task

    @classmethod
    def from_ensemble(
        cls,
        e: DatasetEnsemble,
        member_datasets: dict[L2SourceDataset, xr.Dataset] | None = None,
    ) -> "DatasetStatistics":
        """Compute dataset statistics from a DatasetEnsemble."""
        grid_gdf = entropice.spatial.grids.open(e.grid, e.level)  # Ensure grid is registered
        total_cells = len(grid_gdf)
        target_statistics = {}
        for target in all_target_datasets:
            if isinstance(e.temporal_mode, int) and target == "darts_mllabels":
                # darts_mllabels does not support year-based temporal modes
                continue
            target_statistics[target] = TargetStatistics.compute(e, target=target, total_cells=total_cells)
        member_statistics = MemberStatistics.compute(e, member_datasets=member_datasets)

        total_features = sum(ms.feature_count for ms in member_statistics.values())
        total_size_bytes = sum(ms.size_bytes for ms in member_statistics.values())
        return cls(
            total_features=total_features,
            total_cells=total_cells,
            size_bytes=total_size_bytes,
            members=member_statistics,
            target=target_statistics,
        )

    @staticmethod
    def get_target_sample_count_df(
        all_stats: dict[GridLevel, dict[TemporalMode, "DatasetStatistics"]],
    ) -> pd.DataFrame:
        """Convert sample count data to DataFrame."""
        rows = []
        for grid_config in grid_configs:
            mode_stats = all_stats[grid_config.id]
            for temporal_mode, stats in mode_stats.items():
                for target_name, target_stats in stats.target.items():
                    # We can assume that all tasks have equal counts, since they only affect distribution, not presence
                    target_task_stats = target_stats["binary"]
                    n_samples = target_task_stats.training_cells
                    coverage = target_task_stats.coverage
                    for task, target_task_stats in target_stats.items():
                        assert n_samples == target_task_stats.training_cells, (
                            "Inconsistent sample counts across tasks for the same target."
                        )
                        assert coverage == target_task_stats.coverage, (
                            "Inconsistent coverage across tasks for the same target."
                        )
                    rows.append(
                        {
                            "Grid": grid_config.display_name,
                            "Target": target_name.replace("_", " ").title(),
                            "Temporal Mode": str(temporal_mode).capitalize(),
                            "Samples (Coverage)": n_samples,
                            "Coverage %": coverage,
                            "Grid_Level_Sort": grid_config.sort_key,
                        }
                    )
        return pd.DataFrame(rows).sort_values("Grid_Level_Sort").drop(columns="Grid_Level_Sort")

    @staticmethod
    def get_inference_sample_count_df(
        all_stats: dict[GridLevel, dict[TemporalMode, "DatasetStatistics"]],
    ) -> pd.DataFrame:
        """Convert feature count data to DataFrame."""
        rows = []
        for grid_config in grid_configs:
            mode_stats = all_stats[grid_config.id]
            # We can assume that all temporal modes have equal counts, since they only affect features, not samples
            stats = mode_stats["feature"]
            assert len(stats.members) > 0, "Dataset must have at least one member."
            min_cells = min(member_stats.dimensions["cell_ids"] for member_stats in stats.members.values())

            for temporal_mode, stats in mode_stats.items():
                assert len(stats.members) > 0, "Dataset must have at least one member."
                assert (
                    min(member_stats.dimensions["cell_ids"] for member_stats in stats.members.values()) == min_cells
                ), "Inconsistent inference cell counts across temporal modes."

            rows.append(
                {
                    "Grid": grid_config.display_name,
                    "Inference Cells": min_cells,
                    "Grid_Level_Sort": grid_config.sort_key,
                }
            )
        return pd.DataFrame(rows).sort_values("Grid_Level_Sort").drop(columns="Grid_Level_Sort")

    @staticmethod
    def get_comparison_df(
        all_stats: dict[GridLevel, dict[TemporalMode, "DatasetStatistics"]],
    ) -> pd.DataFrame:
        """Convert comparison data to DataFrame for detailed table."""
        rows = []
        for grid_config in grid_configs:
            mode_stats = all_stats[grid_config.id]
            for temporal_mode, stats in mode_stats.items():
                min_cells = min(member_stats.dimensions["cell_ids"] for member_stats in stats.members.values())
                rows.append(
                    {
                        "Grid": grid_config.display_name,
                        "Temporal Mode": str(temporal_mode).capitalize(),
                        "Total Features": f"{stats.total_features:,}",
                        "Total Cells": f"{stats.total_cells:,}",
                        "Min Cells": f"{min_cells:,}",
                        "Size (MB)": f"{stats.size_bytes / (1024 * 1024):.2f}",
                        "Grid_Level_Sort": grid_config.sort_key,
                    }
                )
        return pd.DataFrame(rows).sort_values(["Grid_Level_Sort", "Temporal Mode"]).drop(columns="Grid_Level_Sort")

    @staticmethod
    def get_feature_breakdown_df(
        all_stats: dict[GridLevel, dict[TemporalMode, "DatasetStatistics"]],
    ) -> pd.DataFrame:
        """Convert feature breakdown data to DataFrame for stacked/donut charts."""
        rows = []
        for grid_config in grid_configs:
            mode_stats = all_stats[grid_config.id]
            added_year_temporal_mode = False
            for temporal_mode, stats in mode_stats.items():
                if added_year_temporal_mode:
                    continue
                if isinstance(temporal_mode, int):
                    # Only add one year-based temporal mode for the breakdown
                    added_year_temporal_mode = True
                    temporal_mode = "Year-Based"
                for member_name, member_stats in stats.members.items():
                    rows.append(
                        {
                            "Grid": grid_config.display_name,
                            "Temporal Mode": temporal_mode.capitalize(),
                            "Member": member_name,
                            "Number of Features": member_stats.feature_count,
                            "Grid_Level_Sort": grid_config.sort_key,
                        }
                    )
        return (
            pd.DataFrame(rows)
            .sort_values(["Grid_Level_Sort", "Temporal Mode", "Member"])
            .drop(columns="Grid_Level_Sort")
        )


@st.cache_data  # ty:ignore[invalid-argument-type]
def load_all_default_dataset_statistics() -> dict[GridLevel, dict[TemporalMode, DatasetStatistics]]:
    """Precompute dataset statistics for all grid-level combinations and temporal modes."""
    cache_file = entropice.utils.paths.get_dataset_stats_cache()
    if cache_file.exists():
        with open(cache_file, "rb") as f:
            dataset_stats = pickle.load(f)
        return dataset_stats

    dataset_stats: dict[GridLevel, dict[TemporalMode, DatasetStatistics]] = {}
    for grid_config in grid_configs:
        dataset_stats[grid_config.id] = {}
        grid_gdf = entropice.spatial.grids.open(grid_config.grid, grid_config.level)  # Ensure grid is registered
        total_cells = len(grid_gdf)
        assert total_cells > 0, "Grid must contain at least one cell."
        for temporal_mode in all_temporal_modes:
            with stopwatch(f"Loading statistics for {grid_config.grid=}, {grid_config.level=}, {temporal_mode=}"):
                e = DatasetEnsemble(grid=grid_config.grid, level=grid_config.level, temporal_mode=temporal_mode)
                target_statistics = {}
                for target in all_target_datasets:
                    if isinstance(temporal_mode, int) and target == "darts_mllabels":
                        # darts_mllabels does not support year-based temporal modes
                        continue
                    target_statistics[target] = TargetStatistics.compute(e, target=target, total_cells=total_cells)
                member_statistics = MemberStatistics.compute(e)

                total_features = sum(ms.feature_count for ms in member_statistics.values())
                total_size_bytes = sum(ms.size_bytes for ms in member_statistics.values())
                dataset_stats[grid_config.id][temporal_mode] = DatasetStatistics(
                    total_features=total_features,
                    total_cells=total_cells,
                    size_bytes=total_size_bytes,
                    members=member_statistics,
                    target=target_statistics,
                )

    with open(cache_file, "wb") as f:
        pickle.dump(dataset_stats, f)

    return dataset_stats


@dataclass(frozen=True)
class TrainingDatasetStatistics:
    """Statistics about the training dataset used for a specific training run."""

    n_samples: int  # Total number of samples in the dataset
    n_features: int  # Number of features
    feature_names: list[str]  # Names of all features
    n_training_samples: int  # Number of cells used for training
    n_test_samples: int  # Number of cells used for testing
    train_test_ratio: float  # Ratio of training to test samples

    class_labels: list[str] | None  # Ordered list of class labels
    class_intervals: list[tuple[float, float] | tuple[int, int] | tuple[None, None]] | None  # Min/max raw values
    n_classes: int | None  # Number of classes

    training_class_counts: dict[str, int] | None  # Class counts in training set
    training_class_distribution: dict[str, float] | None  # Class percentages in training set
    test_class_counts: dict[str, int] | None  # Class counts in test set
    test_class_distribution: dict[str, float] | None  # Class percentages in test set

    raw_value_min: float  # Minimum raw target value
    raw_value_max: float  # Maximum raw target value
    raw_value_mean: float  # Mean raw target value
    raw_value_median: float  # Median raw target value
    raw_value_std: float  # Standard deviation of raw target values

    imbalance_ratio: float | None  # Smallest class count / largest class count (overall)
    size_bytes: int  # Total memory usage of features in bytes

    @classmethod
    def compute(
        cls,
        ensemble: DatasetEnsemble,
        task: Task,
        target: TargetDataset,
    ) -> "TrainingDatasetStatistics":
        """Compute training dataset statistics from a DatasetEnsemble and training settings."""
        training_dataset = ensemble.create_training_set(task=task, target=target)

        # Sample counts
        n_samples = len(training_dataset)
        n_training_samples = (training_dataset.split == "train").sum()
        n_test_samples = (training_dataset.split == "test").sum()
        train_test_ratio = n_training_samples / n_test_samples if n_test_samples > 0 else 0.0

        # Feature statistics
        n_features = len(training_dataset.features.columns)
        feature_names = training_dataset.features.columns.tolist()
        size_bytes = training_dataset.features.memory_usage(deep=True).sum()

        # Class information
        class_labels = training_dataset.target_labels
        if class_labels is None:
            class_intervals = None
            n_classes = None
            training_class_counts = None
            training_class_distribution = None
            test_class_counts = None
            test_class_distribution = None
            imbalance_ratio = None
        else:
            class_intervals = training_dataset.target_intervals
            n_classes = len(class_labels)

            train_y = training_dataset.targets[training_dataset.split == "train"]["y"]
            test_y = training_dataset.targets[training_dataset.split == "test"]["y"]

            train_counts = train_y.value_counts().sort_index()
            training_class_counts = {class_labels[i]: int(train_counts.get(i, 0)) for i in range(n_classes)}
            train_total = sum(training_class_counts.values())
            training_class_distribution = {
                k: (v / train_total * 100) if train_total > 0 else 0.0 for k, v in training_class_counts.items()
            }
            test_counts = test_y.value_counts().sort_index()
            test_class_counts = {class_labels[i]: int(test_counts.get(i, 0)) for i in range(n_classes)}
            test_total = sum(test_class_counts.values())
            test_class_distribution = {
                k: (v / test_total * 100) if test_total > 0 else 0.0 for k, v in test_class_counts.items()
            }

            # Imbalance ratio (smallest class / largest class across both splits)
            all_counts = list(train_counts.values()) + list(test_class_counts.values())
            nonzero_counts = [c for c in all_counts if c > 0]
            imbalance_ratio = min(nonzero_counts) / max(nonzero_counts) if nonzero_counts else 0.0

        # Raw value statistics
        raw_values = training_dataset.targets["z"]
        raw_value_min = float(raw_values.min())
        raw_value_max = float(raw_values.max())
        raw_value_mean = float(raw_values.mean())
        raw_value_median = float(raw_values.median())
        raw_value_std = float(raw_values.std())

        return cls(
            n_samples=n_samples,
            n_features=n_features,
            feature_names=feature_names,
            n_training_samples=n_training_samples,
            n_test_samples=n_test_samples,
            train_test_ratio=train_test_ratio,
            class_labels=class_labels,
            class_intervals=class_intervals,
            n_classes=n_classes,
            training_class_counts=training_class_counts,
            training_class_distribution=training_class_distribution,
            test_class_counts=test_class_counts,
            test_class_distribution=test_class_distribution,
            raw_value_min=raw_value_min,
            raw_value_max=raw_value_max,
            raw_value_mean=raw_value_mean,
            raw_value_median=raw_value_median,
            raw_value_std=raw_value_std,
            imbalance_ratio=imbalance_ratio,
            size_bytes=size_bytes,
        )


@dataclass(frozen=True)
class CVMetricStatistics:
    """Cross-validation statistics for a specific metric."""

    best_score: float
    mean_score: float
    std_score: float
    worst_score: float
    median_score: float
    mean_cv_std: float | None

    @classmethod
    def compute(cls, result: TrainingResult, metric: str) -> "CVMetricStatistics | None":
        """Get cross-validation statistics for a metric."""
        score_col = f"mean_test_{metric}"
        std_col = f"std_test_{metric}"

        cv_results = result.run.cv_results
        if cv_results is None:
            return None

        if score_col not in cv_results.columns:
            raise ValueError(f"Metric {metric} not found in results.")

        best_score = cv_results[score_col].max()
        mean_score = cv_results[score_col].mean()
        std_score = cv_results[score_col].std()
        worst_score = cv_results[score_col].min()
        median_score = cv_results[score_col].median()

        mean_cv_std = None
        if std_col in cv_results.columns:
            mean_cv_std = cv_results[std_col].mean()

        return CVMetricStatistics(
            best_score=best_score,
            mean_score=mean_score,
            std_score=std_score,
            worst_score=worst_score,
            median_score=median_score,
            mean_cv_std=mean_cv_std,
        )


@dataclass(frozen=True)
class ParameterSpaceSummary:
    """Summary statistics for a hyperparameter in the search space."""

    parameter: str
    type: Literal["Numeric", "Categorical"]
    min: float | None
    max: float | None
    mean: float | None
    unique_values: int

    @classmethod
    def compute(cls, result: TrainingResult, param_col: str) -> "ParameterSpaceSummary | None":
        """Get cross-validation statistics for a metric."""
        if result.run.cv_results is None:
            return None
        param_name = param_col.replace("param_", "")
        param_values = result.run.cv_results[param_col].dropna()

        if pd.api.types.is_numeric_dtype(param_values):
            return ParameterSpaceSummary(
                parameter=param_name,
                type="Numeric",
                min=param_values.min(),
                max=param_values.max(),
                mean=param_values.mean(),
                unique_values=param_values.nunique(),
            )
        else:
            unique_vals = param_values.unique()
            return ParameterSpaceSummary(
                parameter=param_name,
                type="Categorical",
                min=None,
                max=None,
                mean=None,
                unique_values=len(unique_vals),
            )


@dataclass(frozen=True)
class CVResultsStatistics:
    """Cross-validation results statistics for all metrics and parameters."""

    metrics: dict[str, CVMetricStatistics]
    parameter_summary: list[ParameterSpaceSummary]

    @classmethod
    def compute(cls, result: TrainingResult) -> "CVResultsStatistics | None":
        """Get cross-validation statistics for a metric."""
        if result.run.cv_results is None:
            return None
        metrics = result.available_metrics
        metric_stats: dict[str, CVMetricStatistics] = {}
        for metric in metrics:
            stats = CVMetricStatistics.compute(result, metric)
            if stats is None:
                continue
            metric_stats[metric] = stats

        param_cols = [col for col in result.run.cv_results.columns if col.startswith("param_") and col != "params"]
        summary_data = []
        for param_col in param_cols:
            summary_data.append(ParameterSpaceSummary.compute(result, param_col))

        return CVResultsStatistics(metrics=metric_stats, parameter_summary=summary_data)

    def metrics_to_dataframe(self) -> pd.DataFrame:
        """Convert metric statistics to a DataFrame."""
        data = defaultdict(list)
        for metric, stats in self.metrics.items():
            data["Metric"].append(metric)
            data["Best Score"].append(stats.best_score)
            data["Mean Score"].append(stats.mean_score)
            data["Std Dev"].append(stats.std_score)
            data["Worst Score"].append(stats.worst_score)
            data["Median Score"].append(stats.median_score)
            data["Mean CV Std Dev"].append(stats.mean_cv_std)

        return pd.DataFrame(data)

    def parameters_to_dataframe(self) -> pd.DataFrame:
        """Convert parameter summary to a DataFrame."""
        return pd.DataFrame([asdict(p) for p in self.parameter_summary])
