"""Feature Importance Experiment."""

from typing import 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")

EXPERIMENT_NAME = "feature_importance_era5-shoulder_arcticdem-v2"


@cli.default
def main(grid: Grid, target: TargetDataset):
    """Feature Importance Experiment."""
    levels = [3, 4, 5, 6] if grid == "hex" else [6, 7, 8, 9, 10]
    for level in levels:
        print(f"Running feature importance experiment for {grid} grid at level {level}...")
        dimension_filters = {"AlphaEarth": {"agg": ["median"]}}
        dataset_ensemble = DatasetEnsemble(
            grid=grid, level=level, members=["AlphaEarth"], dimension_filters=dimension_filters
        )

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

            # AutoGluon
            time_limit = 30 * 60  # 30 minutes
            presets = "extreme"
            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]
                settings = HPOCVRunSettings(
                    n_iter=n_iter,
                    task=task,
                    target=target,
                    splitter=splitter,
                    model=model,
                    # AlphaEarth Embeddings are already normalized unit vectors
                    scaler="none",
                    normalize=False,
                )
                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()
