"""Training helper to create and configure ML models with hyperparameter search."""

from dataclasses import dataclass, field
from typing import TypedDict

import cupy as cp
import pandas as pd
import scipy.stats
import torch
import xarray as xr
from autogluon.tabular import TabularPredictor
from cuml.ensemble import RandomForestClassifier, RandomForestRegressor
from cuml.neighbors import KNeighborsClassifier, KNeighborsRegressor
from entropy import ESPAClassifier
from scipy.stats._distn_infrastructure import rv_continuous_frozen, rv_discrete_frozen
from sklearn.model_selection import (
    GroupKFold,
    GroupShuffleSplit,
    KFold,
    ShuffleSplit,
    StratifiedGroupKFold,
    StratifiedKFold,
    StratifiedShuffleSplit,
)
from sklearn.pipeline import Pipeline
from xgboost.sklearn import XGBClassifier, XGBRegressor

from entropice.ml.dataset import TrainingSet
from entropice.utils.types import Splitter, Task


class Distribution(TypedDict):
    """Distribution specification for hyperparameter optimization."""

    distribution: str
    low: float | int
    high: float | int


type HPConfig = dict[str, Distribution | list]
type SupportedModel = (
    ESPAClassifier
    | XGBClassifier
    | XGBRegressor
    | RandomForestClassifier
    | RandomForestRegressor
    | KNeighborsClassifier
    | KNeighborsRegressor
)


def get_search_space(hp_config: HPConfig) -> dict[str, list | rv_continuous_frozen | rv_discrete_frozen]:
    """Convert the HPConfig into a search space dictionary usable by sklearn's RandomizedSearchCV."""
    search_space = {}
    for key, dist in hp_config.items():
        if isinstance(dist, list):
            search_space[key] = dist
            continue
        assert hasattr(scipy.stats, dist["distribution"]), (
            f"Unknown distribution type for {key}: {dist['distribution']}"
        )
        distfn = getattr(scipy.stats, dist["distribution"])
        # Add edge-case for uniform distribution, as low-high is different there
        # https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.uniform.html#scipy.stats.uniform
        # Using the parameters loc and scale, one obtains the uniform distribution on [loc, loc + scale].
        if dist["distribution"] == "uniform":
            search_space[key] = distfn(loc=dist["low"], scale=dist["high"] - dist["low"])
        else:
            search_space[key] = distfn(dist["low"], dist["high"])
    return search_space


def is_classifier(model: Pipeline | SupportedModel | TabularPredictor) -> bool:
    """Check if the model is a classifier."""
    if isinstance(model, TabularPredictor):
        return model.problem_type in ["binary", "multiclass"]
    if isinstance(model, Pipeline):
        # Check the last step of the pipeline for the model type
        return is_classifier(model.steps[-1][1])
    return isinstance(
        model,
        (ESPAClassifier, XGBClassifier, RandomForestClassifier, KNeighborsClassifier),
    )


def is_regressor(model: Pipeline | SupportedModel | TabularPredictor) -> bool:
    """Check if the model is a regressor."""
    if isinstance(model, TabularPredictor):
        return model.problem_type in ["regression"]
    if isinstance(model, Pipeline):
        # Check the last step of the pipeline for the model type
        return is_regressor(model.steps[-1][1])
    return isinstance(
        model,
        (XGBRegressor, RandomForestRegressor, KNeighborsRegressor),
    )


@dataclass(frozen=True)
class ModelHPOConfig:
    """Model - Hyperparameter Optimization Configuration."""

    model: SupportedModel
    hp_config: HPConfig
    fit_params: dict[str, object] = field(default_factory=dict)

    @property
    def search_space(self) -> dict[str, list | rv_continuous_frozen | rv_discrete_frozen]:
        """Convert the HPConfig into a search space dictionary usable by sklearn's RandomizedSearchCV."""
        return get_search_space(self.hp_config)


# Hardcode Search Settings for now
espa_hpconfig: HPConfig = {
    "eps_cl": {"distribution": "loguniform", "low": 1e-11, "high": 1e-6},
    "eps_e": {"distribution": "loguniform", "low": 1e4, "high": 1e8},
    "initial_K": {"distribution": "randint", "low": 50, "high": 800},
}
xgboost_hpconfig: HPConfig = {
    # Learning & Regularization
    "learning_rate": {"distribution": "loguniform", "low": 0.01, "high": 0.3},
    "n_estimators": {"distribution": "randint", "low": 100, "high": 500},
    # Tree Structure (critical for overfitting control)
    "max_depth": {"distribution": "randint", "low": 3, "high": 10},
    # "min_child_weight": {"distribution": "randint", "low": 1, "high": 10},
    # Feature Sampling (important for high-dimensional data)
    "colsample_bytree": {"distribution": "uniform", "low": 0.3, "high": 1.0},
    # "colsample_bylevel": {"distribution": "uniform", "low": 0.3, "high": 1.0},
    # Row Sampling
    # "subsample": {"distribution": "uniform", "low": 0.5, "high": 1.0},
    # Regularization
    # "reg_alpha": {"distribution": "loguniform", "low": 1e-8, "high": 1.0},  # L1
    "reg_lambda": {"distribution": "loguniform", "low": 1e-8, "high": 10.0},  # L2
}
rf_hpconfig: HPConfig = {
    # Tree Structure
    "max_depth": {"distribution": "randint", "low": 10, "high": 50},
    "n_estimators": {"distribution": "randint", "low": 50, "high": 300},
    # Split Criteria (critical for >100 features)
    "max_features": {"distribution": "uniform", "low": 0.1, "high": 0.5},  # sqrt(100) ≈ 10% of features
    "min_samples_split": {"distribution": "randint", "low": 2, "high": 20},
    # "min_samples_leaf": {"distribution": "randint", "low": 1, "high": 10},
    # Regularization (cuML specific)
    # "min_impurity_decrease": {"distribution": "loguniform", "low": 1e-7, "high": 1e-2},
    # Bootstrap
    # "max_samples": {"distribution": "uniform", "low": 0.5, "high": 1.0},
}
knn_hpconfig: HPConfig = {
    "n_neighbors": {"distribution": "randint", "low": 10, "high": 200},
    "weights": ["uniform", "distance"],  # Categorical parameter
}


def get_model_hpo_config(model: str, task: Task, **model_kwargs) -> ModelHPOConfig:
    """Create a model and its hyperparameter optimization configuration based on the specified model type and task.

    Args:
        model (str): The type of model to create. Supported values are "espa", "xgboost", "rf", and "knn".
        task (Task): The type of task to perform. Supported values are "classifier" and "regressor".
    Kwargs:
        model_kwargs: Additional keyword arguments to pass to the model constructor.

    Returns:
        ModelHPOConfig: A dataclass containing the model instance, hyperparameter configuration, and any fit parameters.

    """
    model_type = "classifier" if task in ("binary", "count_regimes", "density_regimes") else "regressor"
    match (model, model_type):
        case ("espa", "classifier"):
            clf = ESPAClassifier(20, 0.1, 0.1, robust=True, **model_kwargs)
            fit_params = {"max_iter": 300}
            return ModelHPOConfig(clf, espa_hpconfig, fit_params)
        case ("xgboost", "classifier"):
            clf = XGBClassifier(
                objective="multi:softprob" if task != "binary" else "binary:logistic",
                eval_metric="mlogloss" if task != "binary" else "logloss",
                tree_method="hist",
                max_depth=10,
                device="gpu",
                **model_kwargs,
            )
            return ModelHPOConfig(clf, xgboost_hpconfig)
        case ("xgboost", "regressor"):
            reg = XGBRegressor(
                objective="reg:squarederror",
                eval_metric="rmse",
                tree_method="hist",
                max_depth=10,
                device="gpu",
                **model_kwargs,
            )
            return ModelHPOConfig(reg, xgboost_hpconfig)
        case ("rf", "classifier"):
            clf = RandomForestClassifier(split_criterion="entropy", **model_kwargs)
            return ModelHPOConfig(clf, rf_hpconfig)
        case ("rf", "regressor"):
            reg = RandomForestRegressor(split_criterion="poisson", **model_kwargs)
            return ModelHPOConfig(reg, rf_hpconfig)
        case ("knn", "classifier"):
            clf = KNeighborsClassifier(**model_kwargs)
            return ModelHPOConfig(clf, knn_hpconfig)
        case ("knn", "regressor"):
            reg = KNeighborsRegressor(**model_kwargs)
            return ModelHPOConfig(reg, knn_hpconfig)
        case _:
            raise ValueError(f"Unsupported model/task combination: {model}/{task}")


def extract_espa_state(model: ESPAClassifier, training_data: TrainingSet) -> xr.Dataset:
    """Extract the inner state of a trained ESPAClassifier as an xarray Dataset."""
    # Annotate the state with xarray metadata
    boxes = list(range(model.K_))
    box_centers = xr.DataArray(
        model.S_.cpu().numpy(),
        dims=["feature", "box"],
        coords={"feature": training_data.feature_names, "box": boxes},
        name="box_centers",
        attrs={"description": "Centers of the boxes in feature space."},
    )
    box_assignments = xr.DataArray(
        model.Lambda_.cpu().numpy(),
        dims=["class", "box"],
        coords={"class": training_data.target_labels, "box": boxes},
        name="box_assignments",
        attrs={"description": "Assignments of samples to boxes."},
    )
    feature_weights = xr.DataArray(
        model.W_.cpu().numpy(),
        dims=["feature"],
        coords={"feature": training_data.feature_names},
        name="feature_weights",
        attrs={"description": "Weights for each feature."},
    )
    state = xr.Dataset(
        {
            "box_centers": box_centers,
            "box_assignments": box_assignments,
            "feature_weights": feature_weights,
        },
        attrs={
            "description": "Inner state of the best ESPAClassifier from RandomizedSearchCV.",
        },
    )
    return state


def extract_espa_feature_importance(model: ESPAClassifier, training_data: TrainingSet) -> pd.DataFrame:
    """Extract feature importance from a trained ESPAClassifier as a pandas DataFrame."""
    weights = model.W_
    if isinstance(weights, torch.Tensor):
        weights = weights.cpu().numpy()
    elif isinstance(weights, cp.ndarray):
        weights = weights.get()
    return pd.DataFrame(
        {"importance": weights},
        index=training_data.feature_names,
    )


def extract_xgboost_feature_importance(model: XGBClassifier | XGBRegressor, training_data: TrainingSet) -> pd.DataFrame:
    """Extract feature importance from a trained XGBoost model as a pandas DataFrame."""
    booster = model.get_booster()
    fi = {}
    for metric in ["weight", "gain", "cover", "total_gain", "total_cover"]:
        scores = booster.get_score(importance_type=metric)
        fi[metric] = [scores.get(f"f{i}", 0.0) for i in range(len(training_data.feature_names))]
    return pd.DataFrame(fi, index=training_data.feature_names)


def extract_rf_feature_importance(
    model: RandomForestClassifier | RandomForestRegressor, training_data: TrainingSet
) -> pd.DataFrame:
    """Extract feature importance from a trained RandomForest model as a pandas DataFrame."""
    # Extract Random Forest-specific information
    # Note: cuML's RandomForestClassifier doesn't expose individual trees (estimators_)
    # like sklearn does, so we can only extract feature importances and model parameters

    # Feature importances (Gini importance)
    feature_importances = {"gini": model.feature_importances_}
    return pd.DataFrame(feature_importances, index=training_data.feature_names)


def get_splitter(
    splitter: Splitter, n_splits: int
) -> (
    KFold
    | StratifiedKFold
    | ShuffleSplit
    | StratifiedShuffleSplit
    | GroupKFold
    | StratifiedGroupKFold
    | GroupShuffleSplit
):
    """Get a scikit-learn splitter object based on the specified splitter type."""
    match splitter:
        case "kfold":
            return KFold(n_splits=n_splits, shuffle=True, random_state=42)
        case "shuffle":
            return ShuffleSplit(n_splits=n_splits, test_size=0.2, random_state=42)
        case "stratified":
            return StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=42)
        case "stratified_shuffle":
            return StratifiedShuffleSplit(n_splits=n_splits, test_size=0.2, random_state=42)
        case "group":
            return GroupKFold(n_splits=n_splits, shuffle=True, random_state=42)
        case "stratified_group":
            return StratifiedGroupKFold(n_splits=n_splits, shuffle=True, random_state=42)
        case "shuffle_group":
            return GroupShuffleSplit(n_splits=n_splits, test_size=0.2, random_state=42)
        case _:
            raise ValueError(f"Unsupported splitter type: {splitter}")
