"""Training of models with AutoGluon."""

from dataclasses import dataclass
from typing import cast

import cyclopts
import pandas as pd
import shap.maskers
import xarray as xr
from autogluon.tabular import TabularDataset, TabularPredictor
from rich import pretty, traceback
from shap import Explainer, Explanation
from sklearn import set_config
from stopuhr import stopwatch

from entropice.ml.dataset import DatasetEnsemble
from entropice.ml.inference import predict_proba
from entropice.utils.paths import get_training_results_dir
from entropice.utils.training import AutoML, Training
from entropice.utils.types import TargetDataset, Task

traceback.install()
pretty.install()


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


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

    time_limit: int = 3600  # Time limit in seconds (1 hour default)
    presets: str = "extreme"
    task: Task = "binary"
    target: TargetDataset = "darts_v1"
    verbosity: int = 2  # Verbosity level (0-4)


def _compute_metrics_and_confusion_matrix(  # noqa: C901
    predictor: TabularPredictor, train_data: pd.DataFrame, test_data: pd.DataFrame, complete_data: pd.DataFrame
) -> tuple[pd.DataFrame, xr.Dataset | None]:
    train_scores = predictor.evaluate(train_data, display=True, detailed_report=True)
    test_scores = predictor.evaluate(test_data, display=True, detailed_report=True)
    complete_scores = predictor.evaluate(complete_data, display=True, detailed_report=True)
    m = []
    cm = {}
    for split, scores in zip(["train", "test", "complete"], [train_scores, test_scores, complete_scores]):
        for metric, score in scores.items():
            if metric == "confusion_matrix":
                score = cast(pd.DataFrame, score)
                confusion_matrix = xr.DataArray(
                    score.to_numpy(),
                    dims=("y_true", "y_pred"),
                    coords={"y_true": score.index.tolist(), "y_pred": score.columns.tolist()},
                )
                cm[split] = confusion_matrix
            elif metric == "classification_report":
                score = cast(dict[str, dict[str, float]], score)
                score.pop("accuracy")  # Accuracy is already included as a separate metric
                macro_avg = score.pop("macro avg")
                for macro_avg_metric, macro_avg_score in macro_avg.items():
                    metric_name = f"macro_avg_{macro_avg_metric}"
                    m.append({"split": split, "metric": metric_name, "score": macro_avg_score})
                weighted_avg = score.pop("weighted avg")
                for weighted_avg_metric, weighted_avg_score in weighted_avg.items():
                    metric_name = f"weighted_avg_{weighted_avg_metric}"
                    m.append({"split": split, "metric": metric_name, "score": weighted_avg_score})
                for class_name, class_scores in score.items():
                    class_name = class_name.replace(" ", "-")
                    for class_metric, class_score in class_scores.items():
                        m.append({"split": split, "metric": f"{class_name}_{class_metric}", "score": class_score})
            else:  # Scalar metric
                m.append({"split": split, "metric": metric, "score": score})
    if len(cm) == 0:
        return pd.DataFrame(m), None
    elif len(cm) == 3:
        return pd.DataFrame(m), xr.Dataset(cm)
    else:
        raise RuntimeError("Confusion matrix should be computed for all datasets or none.")


def _compute_shap_explanation(
    predictor: TabularPredictor,
    train_data: pd.DataFrame,
    test_data: pd.DataFrame,
    feature_names: list[str],
    target_labels: list[str] | None,
) -> Explanation:
    masker = shap.maskers.Independent(data=train_data.drop(columns=["label"]))
    explainer = Explainer(
        predictor.predict_proba if predictor.problem_type in ["binary", "multiclass"] else predictor.predict,
        masker=masker,
        seed=42,
        feature_names=feature_names,
        output_names=target_labels,
    )
    samples = test_data.drop(columns=["label"])
    if len(samples) > 100:
        samples = samples.sample(n=100, random_state=42)
    explanation = explainer(samples)
    return explanation


@cli.default
def train(
    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.

    """
    set_config(array_api_dispatch=False)
    results_dir = get_training_results_dir(
        experiment=experiment,
        name="autogluon",
        grid=dataset_ensemble.grid,
        level=dataset_ensemble.level,
        task=settings.task,
        target=settings.target,
    )
    print(f"\n💾 Results directory: {results_dir}")

    print("Creating training data...")
    training_data = dataset_ensemble.create_training_set(task=settings.task, target=settings.target)
    # Convert to AutoGluon TabularDataset
    train_data: pd.DataFrame = TabularDataset(training_data.to_dataframe("train"))  # ty:ignore[invalid-assignment]
    test_data: pd.DataFrame = TabularDataset(training_data.to_dataframe("test"))  # ty:ignore[invalid-assignment]
    complete_data: pd.DataFrame = TabularDataset(training_data.to_dataframe(None))  # ty:ignore[invalid-assignment]

    # Initialize TabularPredictor
    print(f"\n🚀 Initializing AutoGluon TabularPredictor (preset='{settings.presets}')...")
    predictor = TabularPredictor(
        label="label",
        path=str(results_dir / "models"),
        verbosity=settings.verbosity,
    )

    # Train models
    print(f"\n⚡ Training models for {settings.time_limit / 60}min...")
    with stopwatch("AutoGluon training"):
        predictor.fit(
            train_data=train_data,
            time_limit=settings.time_limit,
            presets=settings.presets,
            num_gpus=1,
        )

    print("\n📊 Evaluating model performance...")
    leaderboard = predictor.leaderboard(test_data)
    feature_importance = predictor.feature_importance(test_data)
    metrics, confusion_matrix = _compute_metrics_and_confusion_matrix(predictor, train_data, test_data, complete_data)

    # ?: GPU inference is not yet implemented in AutoGluon, hence SHAP computation takes ages for large model ensembles,
    # as they are present in the higher quality presets. Disabling SHAP for now...
    # with stopwatch("Explaining model predictions with SHAP..."):
    #     explanation = _compute_shap_explanation(
    #         predictor, train_data, test_data, training_data.feature_names, training_data.target_labels
    #     )

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

    summary = Training(
        path=results_dir,
        dataset=dataset_ensemble,
        method=AutoML(time_budget=settings.time_limit, preset=settings.presets, hpo=False),
        task=settings.task,
        target=settings.target,
        model=predictor,
        model_type="autogluon",
        metrics=metrics,
        feature_importance=feature_importance,
        shap_explanation=None,
        predictions=preds,
        confusion_matrix=confusion_matrix,
        cv_results=None,
        leaderboard=leaderboard,
    )
    summary.save()
    return summary
