"""Feature Importance Experiment."""

from typing import Literal, cast

import cyclopts
from stopuhr import stopwatch

from entropice.ml.autogluon import RunSettings as AutoGluonRunSettings
from entropice.ml.autogluon import train as train_autogluon
from entropice.ml.dataset import DatasetEnsemble
from entropice.ml.hpsearchcv import RunSettings as HPOCVRunSettings
from entropice.ml.hpsearchcv import hpsearch_cv
from entropice.utils.paths import RESULTS_DIR
from entropice.utils.types import Grid, Model, TargetDataset, Task

cli = cyclopts.App("entropice-feature-importance")

DEV = False

EXPERIMENT_NAME = "feature_importance_era5-shoulder_arcticdem-v2"
if DEV:
    EXPERIMENT_NAME = "tobis-final-tests"


@cli.default
def main(grid: Grid, target: TargetDataset, selection: Literal["none", "cluster", "univariate"] = "none"):
    """Feature Importance Experiment."""
    levels = [3, 4, 5, 6] if grid == "hex" else [6, 7, 8, 9, 10]
    if DEV:
        levels = [3, 6] if grid == "hex" else [6, 10]
    for level in levels:
        print(f"Running feature importance experiment for {grid} grid at level {level}...")
        dimension_filters = {"ArcticDEM": {"aggregations": ["median"]}}
        if (grid == "hex" and level in [3, 4]) or (grid == "healpix" and level in [6, 7]):
            dimension_filters["ERA5-shoulder"] = {"aggregations": ["median"]}
        dataset_ensemble = DatasetEnsemble(
            grid=grid, level=level, members=["ArcticDEM", "ERA5-shoulder"], dimension_filters=dimension_filters
        )

        for task in cast(list[Task], ["binary", "density"]):
            print(f"\nRunning for {task}...")

            # AutoGluon
            time_limit = 30 * 60  # 30 minutes
            if DEV:
                time_limit = 2 * 60  # 2 minutes
            presets = "extreme"
            if DEV:
                presets = "medium"
            settings = AutoGluonRunSettings(
                time_limit=time_limit,
                presets=presets,
                verbosity=2,
                task=task,
                target=target,
            )
            train_autogluon(dataset_ensemble, settings, experiment=EXPERIMENT_NAME)

            # HPOCV
            splitter = "stratified_shuffle" if task == "binary" else "kfold"
            models: list[Model] = ["xgboost", "rf", "knn"]
            if task == "binary":
                models.append("espa")
            for model in models:
                print(f"\nRunning HPOCV for model {model}...")
                n_iter = {
                    "espa": 300,
                    "xgboost": 300,
                    "rf": 100,  # RF is slow, so we reduce the number of iterations
                    "knn": 40,  # kNN hpspace is small, so we reduce the number of iterations
                }[model]
                if DEV:
                    n_iter = 3
                scaler = "standard" if model in ["espa", "knn"] else "none"
                normalize = scaler != "none"
                settings = HPOCVRunSettings(
                    n_iter=n_iter,
                    task=task,
                    target=target,
                    splitter=splitter,
                    model=model,
                    scaler=scaler,
                    normalize=normalize,
                )
                hpsearch_cv(dataset_ensemble, settings, experiment=EXPERIMENT_NAME)

    stopwatch.summary()
    times = stopwatch.export()
    times.to_parquet(RESULTS_DIR / EXPERIMENT_NAME / f"training_times_{target}_{grid}.parquet")
    print("Done.")


if __name__ == "__main__":
    cli()
