"""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, L2SourceDataset, Model, TargetDataset, Task

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


EXPERIMENT_NAME = "feature_importance_era5-shoulder_arcticdem-v2"

all_features = None
univariate_features = [
    "x",
    "y",
    "era5_freezing_days_AMJ",
    "era5_freezing_days_JAS",
    "era5_freezing_days_trend_JAS",
    "era5_freezing_degree_days_AMJ",
    "era5_freezing_degree_days_JAS",
    "era5_freezing_degree_days_JFM",
    "era5_freezing_degree_days_trend_JAS",
    "era5_freezing_degree_days_trend_JFM",
    "era5_freezing_degree_days_trend_OND",
    "era5_precipitation_occurrences_AMJ",
    "era5_precipitation_occurrences_OND",
    "era5_sf_JAS",
    "era5_snowc_mean_trend_JAS",
    "era5_snowfall_occurrences_JAS",
    "era5_snowfall_occurrences_trend_JAS",
    "era5_sshf_JFM",
    "era5_sshf_OND",
    "era5_t2m_avg_AMJ",
    "era5_t2m_avg_JAS",
    "era5_t2m_avg_JFM",
    "era5_t2m_avg_trend_JFM",
    "era5_t2m_avg_trend_OND",
    "era5_t2m_max_AMJ",
    "era5_t2m_max_JAS",
    "era5_t2m_max_JFM",
    "era5_t2m_max_OND",
    "era5_t2m_max_trend_OND",
    "era5_t2m_mean_AMJ",
    "era5_t2m_mean_JAS",
    "era5_t2m_mean_JFM",
    "era5_t2m_mean_trend_JFM",
    "era5_t2m_mean_trend_OND",
    "era5_t2m_min_AMJ",
    "era5_t2m_min_trend_JAS",
    "era5_t2m_min_trend_JFM",
    "era5_t2m_range_AMJ",
    "era5_t2m_range_JAS",
    "era5_t2m_skew_JAS",
    "era5_thawing_days_AMJ",
    "era5_thawing_days_JAS",
    "era5_thawing_days_trend_JAS",
    "era5_thawing_degree_days_AMJ",
    "era5_thawing_degree_days_JAS",
    "era5_tp_AMJ",
    "era5_tp_JAS",
]
cluster_features = [
    "cell_area",
    "water_area",
    "x",
    "y",
    "arcticdem_aspect_median",
    "arcticdem_curvature_median",
    "arcticdem_dem_median",
    "arcticdem_slope_median",
    "era5_effective_snow_depth_trend_AMJ",
    "era5_effective_snow_depth_trend_JAS",
    "era5_effective_snow_depth_trend_JFM",
    "era5_effective_snow_depth_trend_OND",
    "era5_freezing_days_AMJ",
    "era5_freezing_days_JAS",
    "era5_freezing_days_JFM",
    "era5_freezing_days_OND",
    "era5_freezing_days_trend_AMJ",
    "era5_freezing_days_trend_OND",
    "era5_freezing_degree_days_trend_AMJ",
    "era5_freezing_degree_days_trend_JFM",
    "era5_freezing_degree_days_trend_OND",
    "era5_lblt_max_JAS",
    "era5_lblt_max_trend_AMJ",
    "era5_lblt_max_trend_JAS",
    "era5_lblt_max_trend_JFM",
    "era5_lblt_max_trend_OND",
    "era5_precipitation_occurrences_AMJ",
    "era5_precipitation_occurrences_JFM",
    "era5_precipitation_occurrences_trend_AMJ",
    "era5_precipitation_occurrences_trend_JAS",
    "era5_precipitation_occurrences_trend_JFM",
    "era5_precipitation_occurrences_trend_OND",
    "era5_sf_AMJ",
    "era5_sf_trend_AMJ",
    "era5_sf_trend_JAS",
    "era5_snowc_mean_JFM",
    "era5_snowc_mean_trend_AMJ",
    "era5_snowc_mean_trend_OND",
    "era5_sshf_AMJ",
    "era5_sshf_JFM",
    "era5_sshf_trend_AMJ",
    "era5_sshf_trend_JFM",
    "era5_t2m_avg_trend_AMJ",
    "era5_t2m_avg_trend_JAS",
    "era5_t2m_max_AMJ",
    "era5_t2m_max_trend_JAS",
    "era5_t2m_max_trend_JFM",
    "era5_t2m_max_trend_OND",
    "era5_t2m_min_JFM",
    "era5_t2m_min_trend_OND",
    "era5_t2m_range_JFM",
    "era5_t2m_range_trend_AMJ",
    "era5_t2m_range_trend_JFM",
    "era5_t2m_skew_trend_AMJ",
    "era5_t2m_skew_trend_JAS",
    "era5_t2m_skew_trend_JFM",
    "era5_t2m_skew_trend_OND",
]


@cli.default
def main(grid: Grid, target: TargetDataset, selection: Literal["none", "cluster", "univariate"] = "none"):
    """Feature Importance Experiment."""
    experiment_name = EXPERIMENT_NAME
    if selection != "none":
        experiment_name += f"_{selection=}"
    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}...")
        members: list[L2SourceDataset] = ["ArcticDEM", "ERA5-shoulder"]
        dimension_filters = {"ArcticDEM": {"aggregations": ["median"]}}
        is_era5_median_grid = (grid == "hex" and level in [3, 4]) or (grid == "healpix" and level in [6, 7])
        is_era5_mean_grid = (grid == "hex" and level in [5]) or (grid == "healpix" and level in [8, 9])
        if is_era5_median_grid:
            dimension_filters["ERA5-shoulder"] = {"aggregations": ["median"]}

        feature_filters = {
            "none": all_features,
            "cluster": cluster_features,
            "univariate": univariate_features,
        }[selection]
        if feature_filters is not None:
            if not any(f.startswith("arcticdem") for f in feature_filters):
                dimension_filters.pop("ArcticDEM", None)
                members.remove("ArcticDEM")
            if not any(f.startswith("era5") for f in feature_filters):
                dimension_filters.pop("ERA5-shoulder", None)
                members.remove("ERA5-shoulder")
            elif is_era5_median_grid:

                def _add_median(feature: str) -> str:
                    if not feature.startswith("era5"):
                        return feature
                    return (
                        feature.replace("JAS", "median_JAS")
                        .replace("AMJ", "median_AMJ")
                        .replace("JFM", "median_JFM")
                        .replace("OND", "median_OND")
                    )

                feature_filters = [_add_median(f) for f in feature_filters]

            elif is_era5_mean_grid:

                def _add_mean(feature: str) -> str:
                    if not feature.startswith("era5"):
                        return feature
                    return (
                        feature.replace("JAS", "mean_JAS")
                        .replace("AMJ", "mean_AMJ")
                        .replace("JFM", "mean_JFM")
                        .replace("OND", "mean_OND")
                    )

                feature_filters = [_add_mean(f) for f in feature_filters]

        dataset_ensemble = DatasetEnsemble(
            grid=grid,
            level=level,
            members=members,
            dimension_filters=dimension_filters,
            post_merge_filters=feature_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]
                scaler = "standard" if model in ["espa", "knn"] else "none"
                settings = HPOCVRunSettings(
                    n_iter=n_iter,
                    task=task,
                    target=target,
                    splitter=splitter,
                    model=model,
                    scaler=scaler,
                    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()
