"""Training of models with hyperparameter search with cross-validation."""

from dataclasses import dataclass
from typing import Literal, cast

import cupy as cp
import cyclopts
import numpy as np
import pandas as pd
import shap.maskers
import torch
import xarray as xr
from rich import pretty, traceback
from shap import Explainer, Explanation, TreeExplainer
from sklearn import set_config
from sklearn.inspection import permutation_importance
from sklearn.metrics import confusion_matrix
from sklearn.model_selection import RandomizedSearchCV
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import (
    FunctionTransformer,
    MaxAbsScaler,
    MinMaxScaler,
    Normalizer,
    QuantileTransformer,
    StandardScaler,
)
from stopuhr import stopwatch

from entropice.ml.dataset import DatasetEnsemble, SplittedArrays, TrainingSet
from entropice.ml.inference import predict_proba
from entropice.ml.models import (
    ModelHPOConfig,
    extract_espa_feature_importance,
    extract_rf_feature_importance,
    extract_xgboost_feature_importance,
    get_model_hpo_config,
    get_splitter,
)
from entropice.utils.metrics import get_metrics, metric_functions
from entropice.utils.paths import get_training_results_dir
from entropice.utils.training import HPOCV, Training, move_data_to_device
from entropice.utils.types import HPSearch, Model, Scaler, Splitter, TargetDataset, Task

traceback.install()
pretty.install()


cli = cyclopts.App(
    "entropice-hpsearchcv",
    config=cyclopts.config.Toml("hpsearchcv-config.toml", root_keys=["tool", "entropice-hpsearchcv"]),  # ty:ignore[invalid-argument-type]
)


@cyclopts.Parameter("*")
@dataclass(frozen=True, kw_only=True)
class RunSettings:
    """Run settings for training."""

    n_iter: int = 2000
    task: Task = "binary"
    target: TargetDataset = "darts_v1"
    search: HPSearch = "random"  # TODO: Implement grid and bayesian search
    splitter: Splitter = "stratified_shuffle"
    model: Model = "espa"
    scaler: Scaler = "standard"
    normalize: bool = False

    @property
    def device(self) -> Literal["torch", "cuda", "cpu"]:
        """Get the device to use for model training.

        Note: Device Management currently is super nasty:
            - CuML (RF & kNN) expects cupy ("cuda") arrays
            - XGBoost always returns CPU arrays
            - eSPA for some reason is super slow on cupy but very fast on torch tensors
            - The sklearn permutation stuff does not work with torch tensors, only cupy and numpy arrays
            - SHAP in general only works with CPU
        """
        return "torch" if self.model == "espa" else "cuda"

    @property
    def hpo_config(self) -> ModelHPOConfig:
        """Get the hyperparameter optimization configuration for the selected model and task."""
        return get_model_hpo_config(self.model, self.task)

    def build_pipeline(self, model_hpo_config: ModelHPOConfig) -> Pipeline:
        """Build a scikit-learn Pipeline based on the settings."""
        # Add a feature scaler / normalization step if specified, but assert that it's only used for non-Tree models
        if self.model in ["rf", "xgboost"]:
            assert self.scaler == "none", f"Scaler {self.scaler} is not viable with model {self.model}"

        match self.scaler:
            case "standard":
                scaler = StandardScaler()
            case "minmax":
                scaler = MinMaxScaler()
            case "maxabs":
                scaler = MaxAbsScaler()
            case "quantile":
                scaler = QuantileTransformer(output_distribution="normal")
            case "none":
                scaler = None
            case _:
                raise ValueError(f"Unknown scaler: {self.scaler}")

        pipeline_steps = []
        if scaler is not None:
            print(f"Using {scaler.__class__.__name__} for feature scaling.")
            pipeline_steps.append(("scaler", scaler))
        if self.normalize:
            print("Using Normalizer for feature normalization.")
            pipeline_steps.append(("normalizer", Normalizer()))
        # Necessary, because scaler and normalizer are only GPU ready in 1.8.0 but autogluon requries <1.8.0
        if len(pipeline_steps) > 0 and self.device != "cpu":
            print(f"Adding steps to move data to CPU for preprocessing and back to {self.device} for model fitting.")
            to_cpu = FunctionTransformer(move_data_to_device, kw_args={"device": "cpu"})
            pipeline_steps.insert(0, ("to_cpu", to_cpu))
            to_gpu = FunctionTransformer(move_data_to_device, kw_args={"device": self.device})
            pipeline_steps.append(("to_device", to_gpu))
        pipeline_steps.append(("model", model_hpo_config.model))
        return Pipeline(pipeline_steps)

    def build_search(
        self, pipeline: Pipeline, model_hpo_config: ModelHPOConfig, metrics: list[str], refit: str
    ) -> RandomizedSearchCV:
        """Build a scikit-learn RandomizedSearchCV based on the settings."""
        if self.task in ["density", "count"]:
            assert self.splitter not in ["stratified_shuffle", "stratified_kfold"], (
                f"Splitter {self.splitter} is not viable for regression tasks"
            )

        cv = get_splitter(self.splitter, n_splits=5)
        print(f"Using {cv.__class__.__name__} for cross-validation splitting.")

        search_space = {f"model__{k}": v for k, v in model_hpo_config.search_space.items()}
        if self.search == "random":
            hp_search = RandomizedSearchCV(
                estimator=pipeline,
                param_distributions=search_space,
                n_iter=self.n_iter,
                cv=cv,
                scoring=metrics,
                refit=refit,
                random_state=42,
                verbose=2,
            )
        else:
            raise NotImplementedError(f"Search method {self.search} not implemented yet.")
        return hp_search


def _extract_cv_results(cv_results) -> pd.DataFrame:
    cv_results = pd.DataFrame(cv_results)
    # Parse the params into individual columns
    params = pd.json_normalize(cv_results["params"])  # ty:ignore[invalid-argument-type]
    cv_results = pd.concat([cv_results.drop(columns=["params"]), params], axis=1)
    return cv_results


def _compute_metrics(y: SplittedArrays, y_pred: SplittedArrays, metrics: list[str]) -> pd.DataFrame:
    m = []
    for metric in metrics:
        metric_fn = metric_functions[metric]
        for split in ["train", "test", "complete"]:
            value = metric_fn(getattr(y, split), getattr(y_pred, split))
            m.append({"metric": metric, "split": split, "score": value})
    return pd.DataFrame(m)


def _compute_confusion_matrices(
    y: SplittedArrays, y_pred: SplittedArrays, codes: np.ndarray, labels: list[str] | None
) -> xr.Dataset:
    if labels is None:
        labels = [str(code) for code in codes]
    cm = xr.Dataset(
        {
            "test": (("true_label", "predicted_label"), confusion_matrix(y.test, y_pred.test, labels=codes)),
            "train": (("true_label", "predicted_label"), confusion_matrix(y.train, y_pred.train, labels=codes)),
            "complete": (
                ("true_label", "predicted_label"),
                confusion_matrix(y.complete, y_pred.complete, labels=codes),
            ),
        },
        coords={"true_label": labels, "predicted_label": labels},
    )
    return cm


def _compute_feature_importance(model: Model, best_estimator: Pipeline, training_data: TrainingSet) -> pd.DataFrame:
    X_test = training_data.X.as_numpy().test if model in ["xgboost", "espa"] else training_data.X.test  # noqa: N806
    y_test = training_data.y.as_numpy().test if model in ["xgboost", "espa"] else training_data.y.test
    if model == "espa":
        # eSPA is super slow with cupy arrays, so we need to move the data to CPU for permutation importance
        # To to this, we manually recreate the pipeline but without the device steps
        if "scaler" in best_estimator.named_steps:
            X_test = best_estimator.named_steps["scaler"].transform(X_test)  # noqa: N806
        if "normalizer" in best_estimator.named_steps:
            X_test = best_estimator.named_steps["normalizer"].transform(X_test)  # noqa: N806
        best_estimator.named_steps["model"].to_numpy()  # inplace

    r = permutation_importance(
        best_estimator if model != "espa" else best_estimator.named_steps["model"],
        X_test,
        y_test,
        n_repeats=5,
        random_state=0,
        max_samples=min(5000, training_data.X.test.shape[0]),
    )

    if model == "espa":
        best_estimator.named_steps["model"].to_torch()  # inplace

    feature_importances = pd.DataFrame(
        {
            "importance": r["importances_mean"],
            "stddev": r["importances_std"],
        },
        index=training_data.feature_names,
    )
    match model:
        case "espa":
            model_feature_importances = extract_espa_feature_importance(
                best_estimator.named_steps["model"], training_data
            ).rename(columns={"importance": "model_feature_weights"})
        case "xgboost":
            model_feature_importances = extract_xgboost_feature_importance(
                best_estimator.named_steps["model"], training_data
            )
        case "rf":
            model_feature_importances = extract_rf_feature_importance(
                best_estimator.named_steps["model"], training_data
            )
        case _:
            model_feature_importances = None
    if model_feature_importances is not None:
        feature_importances = feature_importances.join(model_feature_importances)
    return feature_importances


def _compute_shap_explanation(model: Model, best_estimator: Pipeline, training_data: TrainingSet) -> Explanation:  # noqa: C901
    match model:
        case "espa" | "knn":
            train_transformed = training_data.X.as_numpy().train
            if "scaler" in best_estimator.named_steps:
                train_transformed = best_estimator.named_steps["scaler"].transform(train_transformed)
            if "normalizer" in best_estimator.named_steps:
                train_transformed = best_estimator.named_steps["normalizer"].transform(train_transformed)
            masker = shap.maskers.Independent(data=train_transformed)

            def _model_predict(data):
                mdl = best_estimator.named_steps["model"]
                f = mdl.predict_proba if hasattr(mdl, "predict_proba") else mdl.predict

                out = f(data)
                if isinstance(out, torch.Tensor):
                    out = out.cpu().numpy()
                elif isinstance(out, cp.ndarray):
                    out = out.get()
                return out

            explainer = Explainer(
                _model_predict,
                masker=masker,
                seed=42,
                feature_names=training_data.feature_names,
                output_names=training_data.target_labels,
            )
        case "rf":
            masker = shap.maskers.Independent(data=training_data.X.as_numpy().train)
            explainer = TreeExplainer(
                best_estimator.named_steps["model"].as_sklearn(),
                data=masker,
                feature_names=training_data.feature_names,
            )
        case "xgboost":
            explainer = TreeExplainer(best_estimator.named_steps["model"], feature_names=training_data.feature_names)
        case _:
            raise ValueError(f"Unknown model: {model}")

    samples = training_data.X.as_numpy().test
    nsamples = 2 * samples.shape[1] + 2048
    if len(samples) > nsamples:
        rng = np.random.default_rng(seed=42)
        sample_indices = rng.choice(len(samples), size=nsamples, replace=False)
        samples = samples[sample_indices]
    if "scaler" in best_estimator.named_steps:
        samples = best_estimator.named_steps["scaler"].transform(samples)
    if "normalizer" in best_estimator.named_steps:
        samples = best_estimator.named_steps["normalizer"].transform(samples)
    explanation = explainer(samples)
    return explanation


@cli.default
def hpsearch_cv(
    dataset_ensemble: DatasetEnsemble,
    settings: RunSettings = RunSettings(),
    experiment: str | None = None,
) -> Training:
    """Perform random cross-validation on the training dataset.

    Args:
        dataset_ensemble (DatasetEnsemble): The dataset ensemble configuration.
        settings (RunSettings): This runs settings.
        experiment (str | None): Optional experiment name for results directory.

    """
    results_dir = get_training_results_dir(
        experiment=experiment,
        name="random_search",
        grid=dataset_ensemble.grid,
        level=dataset_ensemble.level,
        task=settings.task,
        target=settings.target,
        model_type=settings.model,
    )

    # Since we use cuml and xgboost libraries, we can only enable array API for ESPA
    use_array_api = settings.model != "xgboost"
    set_config(array_api_dispatch=use_array_api)

    print("Creating training data...")
    training_data = dataset_ensemble.create_training_set(
        task=settings.task, target=settings.target, device=settings.device
    )

    model_hpo_config = settings.hpo_config
    print(f"Using model: {settings.model} with parameters: {model_hpo_config.hp_config}")

    metrics, refit = get_metrics(settings.task)
    print(f"Using {len(metrics)} metrics as scoring and {refit} for refitting.")

    pipeline = settings.build_pipeline(model_hpo_config)
    print(f"Pipeline steps: {pipeline.named_steps}")

    hp_search = settings.build_search(pipeline, model_hpo_config, metrics, refit)
    print(f"Starting hyperparameter search for {settings.model} with {settings.n_iter} iterations...")
    with stopwatch(f"RandomizedSearchCV fitting of {settings.model} for {settings.n_iter} candidates"):
        fit_params = {f"model__{k}": v for k, v in model_hpo_config.fit_params.items()}
        hp_search.fit(
            training_data.X.train,
            # XGBoost returns it's labels as numpy arrays instead of cupy arrays
            # Thus, for the scoring to work, we need to convert them back to numpy
            training_data.y.as_numpy().train if settings.model == "xgboost" else training_data.y.train,
            **fit_params,
        )
    print("Best parameters combination found:")
    best_estimator = cast(Pipeline, hp_search.best_estimator_)
    best_parameters = best_estimator.get_params()
    for param_name in sorted(model_hpo_config.hp_config.keys()):
        search_param_name = f"model__{param_name}"
        print(f"{param_name}: {best_parameters[search_param_name]}")

    # Compute predictions on the all sets and move them to numpy for metric computations
    y_pred = SplittedArrays(
        train=best_estimator.predict(training_data.X.train),
        test=best_estimator.predict(training_data.X.test),
    ).as_numpy()
    y = training_data.y.as_numpy()
    metrics = _compute_metrics(y, y_pred, metrics)
    if settings.task in ["binary", "count_regimes", "density_regimes"]:
        confusion_matrix = _compute_confusion_matrices(
            y, y_pred, codes=np.array(training_data.target_codes), labels=training_data.target_labels
        )
    else:
        confusion_matrix = None
    print("Metrics computed!")

    with stopwatch("Computing feature importance with permutation importance..."):
        feature_importance = _compute_feature_importance(settings.model, best_estimator, training_data)

    with stopwatch("Explaining model predictions with SHAP..."):
        explanation = _compute_shap_explanation(settings.model, best_estimator, training_data)

    print("Predicting probabilities for all cells...")
    preds = predict_proba(dataset_ensemble, model=best_estimator, task=settings.task, device=settings.device)
    print(f"Predicted probabilities DataFrame with {len(preds)} entries.")

    summary = Training(
        path=results_dir,
        dataset=dataset_ensemble,
        method=HPOCV(
            method=settings.search,
            splitter=settings.splitter,
            scaler=settings.scaler,
            normalize=settings.normalize,
            n_iter=settings.n_iter,
            hpconfig=model_hpo_config.hp_config,
        ),
        task=settings.task,
        target=settings.target,
        model=best_estimator,
        model_type=settings.model,
        metrics=metrics,
        feature_importance=feature_importance,
        shap_explanation=explanation,
        predictions=preds,
        confusion_matrix=confusion_matrix,
        cv_results=_extract_cv_results(hp_search.cv_results_),
        leaderboard=None,
    )
    summary.save()

    return summary
