"""Training utilities for Entropice."""

import pickle
from dataclasses import asdict, dataclass
from functools import cached_property
from pathlib import Path
from typing import Any, Literal

import cupy as cp
import geopandas as gpd
import numpy as np
import pandas as pd
import toml
import torch
import xarray as xr
from shap import Explanation

from entropice.ml.dataset import DatasetEnsemble, TrainingSet
from entropice.ml.models import (
    HPConfig,
    extract_espa_state,
    get_search_space,
)
from entropice.utils.types import HPSearch, Model, Scaler, Splitter, TargetDataset, Task

type ndarray = np.ndarray | torch.Tensor | cp.ndarray


def move_data_to_device(data: ndarray, device: Literal["torch", "cuda", "cpu"]) -> ndarray:
    """Move the given data to the specified device (CPU, CUDA, or PyTorch tensor)."""
    match device:
        case "torch":
            torch_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
            return torch.asarray(data, device=torch_device)
        case "cuda":
            if isinstance(data, cp.ndarray):
                return data
            with cp.cuda.Device(0):
                return cp.asarray(data)
        case "cpu":
            if isinstance(data, np.ndarray):
                return data
            elif isinstance(data, torch.Tensor):
                return data.cpu().numpy()
            elif isinstance(data, cp.ndarray):
                return data.get()
        case _:
            raise ValueError(f"Unknown device: {device}")


@dataclass
class HPOCV:  # noqa: D101
    method: HPSearch
    splitter: Splitter
    scaler: Scaler
    normalize: bool
    n_iter: int
    hpconfig: HPConfig

    @cached_property
    def search_space(self):  # noqa: D102
        return get_search_space(self.hpconfig)


@dataclass
class AutoML:  # noqa: D101
    time_budget: int
    preset: str
    hpo: bool


@dataclass
class Training:
    """Configuration and results of a training run.

    A training run involves the complete approach of training a machine learning model, currently:

    1. A hyperparameter optimization using cross-validation
    2. Use of AutoML techniques

    Thus, a training run is always defined as:

    f(dataset, method) -> (model, metrics)

    Metrics refer to a simple dataframe in long format, with the columns: "metric", "split", "score".
    Split is either "train", "test" or "complete".
    """

    path: Path
    dataset: DatasetEnsemble
    method: HPOCV | AutoML
    task: Task
    target: TargetDataset
    model: Any
    model_type: Model
    metrics: pd.DataFrame
    feature_importance: pd.DataFrame
    shap_explanation: Explanation | None
    predictions: gpd.GeoDataFrame
    confusion_matrix: xr.Dataset | None  # only for classification tasks
    cv_results: pd.DataFrame | None  # only for HPOCV
    leaderboard: pd.DataFrame | None  # only for AutoGluon

    def __repr__(self) -> str:  # noqa: D105
        return (
            f"Training("
            f"path={self.path}, "
            f"dataset={self.dataset.grid}-{self.dataset.level}{self.dataset.members}, "
            f"method={type(self.method).__name__}, "
            f"task={self.task}, "
            f"target={self.target})"
        )

    @cached_property
    def metric_names(self) -> list[str]:
        """Get the list of metric names from the metrics DataFrame."""
        return self.metrics["metric"].unique().tolist()

    @cached_property
    def training_set(self) -> TrainingSet:
        """Get the training set for this training run."""
        return self.dataset.create_training_set(self.task, self.target, device="cpu")

    @property
    def method_type(self) -> Literal["HPOCV", "AutoML"]:
        """Get the type of method used in this training run."""
        if isinstance(self.method, HPOCV):
            return "HPOCV"
        elif isinstance(self.method, AutoML):
            return "AutoML"
        else:
            raise ValueError(f"Unknown method type: {type(self.method)}")

    @property
    def n_trials(self) -> int | None:
        """Get the number of trials in the hyperparameter search, if applicable."""
        if self.method_type == "HPOCV" and self.cv_results is not None:
            return len(self.cv_results)
        elif self.method_type == "AutoML" and self.leaderboard is not None:
            return len(self.leaderboard)
        else:
            return None

    @property
    def get_state(self) -> xr.Dataset | pd.DataFrame | None:
        """Get the inner state of the trained model, if available."""
        if self.model_type == "espa":
            return extract_espa_state(self.model, self.training_set)
        else:
            return None

    def get_metrics_from_split(self, split: Literal["train", "test", "complete"]) -> dict[str, float]:
        """Get a dictionary of metric names and values for the specified split."""
        return self.metrics[self.metrics["split"] == split].set_index("metric")["score"].to_dict()  # ty:ignore[invalid-return-type]

    def save(self):
        """Save the training results to the specified path."""
        self.path.mkdir(parents=True, exist_ok=True)
        config_file = self.path / "training_config.toml"
        model_file = self.path / "model.pkl"
        metrics_file = self.path / "metrics.parquet"
        feature_importance_file = self.path / "feature_importance.parquet"
        predictions_file = self.path / "predictions.parquet"
        # Save config
        with open(config_file, "w") as f:
            toml.dump(
                {
                    "dataset": asdict(self.dataset),
                    "method": asdict(self.method),
                    "method_type": type(self.method).__name__,
                    "task": self.task,
                    "target": self.target,
                    "model_type": self.model_type,
                },
                f,
            )

        # Save model, metrics, explanations and predictions
        model_file.write_bytes(pickle.dumps(self.model))
        self.metrics.to_parquet(metrics_file)
        self.feature_importance.to_parquet(feature_importance_file)
        self.predictions.to_parquet(predictions_file)

        # Save SHAP explanation if it exists
        if self.shap_explanation is not None:
            explanations_file = self.path / "shap_explanation.pkl"
            explanations_file.write_bytes(pickle.dumps(self.shap_explanation))

        # Save the confusion matrix if it exists
        if self.confusion_matrix is not None:
            cm_file = self.path / "confusion_matrix.nc"
            self.confusion_matrix.to_netcdf(cm_file, engine="h5netcdf")

        if self.cv_results is not None:
            results_file = self.path / "search_results.parquet"
            self.cv_results.to_parquet(results_file)

        # Save the leaderboard if it exists
        if self.leaderboard is not None:
            leaderboard_file = self.path / "leaderboard.parquet"
            self.leaderboard.to_parquet(leaderboard_file)

    @classmethod
    def load(cls, path: Path) -> "Training":
        """Load a training run from the specified path."""
        config_file = path / "training_config.toml"
        model_file = path / "model.pkl"
        metrics_file = path / "metrics.parquet"
        feature_importance_file = path / "feature_importance.parquet"
        predictions_file = path / "predictions.parquet"
        explanations_file = path / "shap_explanation.pkl"
        cm_file = path / "confusion_matrix.nc"
        cv_results_file = path / "search_results.parquet"
        leaderboard_file = path / "leaderboard.parquet"

        # Load config
        with open(config_file) as f:
            config = toml.load(f)

        task = config["task"]
        target = config["target"]
        model_type = config["model_type"]

        dataset = DatasetEnsemble(**config["dataset"])

        method_type = config["method_type"]
        if method_type == "HPOCV":
            method = HPOCV(**config["method"])
        elif method_type == "AutoML":
            method = AutoML(**config["method"])
        else:
            raise ValueError(f"Unknown method type: {method_type}")

        # Load model, metrics, explanations and predictions
        model = pickle.loads(model_file.read_bytes())
        metrics = pd.read_parquet(metrics_file)
        feature_importance = pd.read_parquet(feature_importance_file)
        predictions = gpd.read_parquet(predictions_file)

        # Load SHAP explanation if it exists
        shap_explanation = None
        if explanations_file.exists():
            shap_explanation = pickle.loads(explanations_file.read_bytes())

        # Load confusion matrix if it exists
        confusion_matrix = None
        if cm_file.exists():
            confusion_matrix = xr.load_dataset(cm_file, engine="h5netcdf")

        cv_results = None
        if cv_results_file.exists():
            cv_results = pd.read_parquet(cv_results_file)

        # Load leaderboard if it exists
        leaderboard = None
        if leaderboard_file.exists():
            leaderboard = pd.read_parquet(leaderboard_file)

        return cls(
            path=path,
            dataset=dataset,
            method=method,
            task=task,
            target=target,
            model=model,
            model_type=model_type,
            metrics=metrics,
            feature_importance=feature_importance,
            shap_explanation=shap_explanation,
            predictions=predictions,
            confusion_matrix=confusion_matrix,
            cv_results=cv_results,
            leaderboard=leaderboard,
        )
