"""
Module de sélection de variables pour les données AGB.
Utilise une approche combinée : filtrage (variance, corrélation),
importance par 4 modèles, agrégation Borda, et analyse de stabilité Bootstrap.
"""

import random
import warnings
from collections import defaultdict

import numpy as np
import pandas as pd
from catboost import CatBoostRegressor
from lightgbm import LGBMRegressor, early_stopping as lgb_early_stopping, log_evaluation
from sklearn.ensemble import GradientBoostingRegressor
from sklearn.feature_selection import VarianceThreshold
from sklearn.linear_model import Lasso, LassoCV
from sklearn.model_selection import KFold
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import RobustScaler

warnings.filterwarnings("ignore")

# Paramètres par défaut (peuvent être modifiés lors de l'appel)
DEFAULT_PARAMS = {
    "target": "agb",
    "n_features": 12,
    "n_splits": 5,
    "corr_threshold": 0.90,
    "var_threshold": 0.01,
    "borda_coverage": 0.80,
    "max_iter": 50000,
    "n_bootstrap": 50,
    "stability_threshold": 0.70,
    "stability_strategy": "strict",   # 'strict' ou 'hybrid'
}

# Colonnes à exclure de la sélection (ne pas les utiliser comme features)
DROP_COLS = [
    "inventory_date", "start_date", "end_date", "geometry",
    "agb", "cagb", "cagb_category", "HCS",
    "latitude", "longitude", "latitude_proj", "longitude_proj"
]


def get_importances_kfold(model_name, build_model_fn, X, y,
                          feature_names, n_splits=5, needs_scale_arg=False,
                          random_state=42):
    """
    Calcule les importances par validation croisée KFold,
    avec un score de rang réciproque.
    """
    kf = KFold(n_splits=n_splits, shuffle=True, random_state=random_state)
    importances = np.zeros(len(feature_names))

    for train_idx, val_idx in kf.split(X):
        X_tr, y_tr = X[train_idx], y[train_idx]
        X_val, y_val = X[val_idx], y[val_idx]

        model = build_model_fn()

        if needs_scale_arg:
            pipe = Pipeline([
                ("scaler", RobustScaler()),
                ("model", model)
            ])
            pipe.fit(X_tr, y_tr)
            fold_imp = np.abs(pipe.named_steps["model"].coef_)

        elif model_name == "LightGBM":
            model.fit(
                X_tr, y_tr,
                eval_set=[(X_val, y_val)],
                callbacks=[
                    lgb_early_stopping(stopping_rounds=30, verbose=False),
                    log_evaluation(-1)
                ]
            )
            fold_imp = model.feature_importances_

        elif model_name == "CatBoost":
            model.fit(
                X_tr, y_tr,
                eval_set=(X_val, y_val),
                use_best_model=True,
                verbose=False
            )
            fold_imp = model.get_feature_importance()

        else:  # GradientBoosting ou autres
            model.fit(X_tr, y_tr)
            fold_imp = model.feature_importances_

        fold_imp = fold_imp / (fold_imp.sum() + 1e-10)
        ranked_idx = np.argsort(fold_imp)[::-1]
        for rank_based, feat_idx in enumerate(ranked_idx):
            importances[feat_idx] += 1.0 / (rank_based + 1)

    return pd.Series(importances / n_splits, index=feature_names).sort_values(ascending=False)


def select_features(
    input_csv_path,
    output_csv_path=None,
    target=DEFAULT_PARAMS["target"],
    n_features=DEFAULT_PARAMS["n_features"],
    n_splits=DEFAULT_PARAMS["n_splits"],
    corr_threshold=DEFAULT_PARAMS["corr_threshold"],
    var_threshold=DEFAULT_PARAMS["var_threshold"],
    borda_coverage=DEFAULT_PARAMS["borda_coverage"],
    max_iter=DEFAULT_PARAMS["max_iter"],
    n_bootstrap=DEFAULT_PARAMS["n_bootstrap"],
    stability_threshold=DEFAULT_PARAMS["stability_threshold"],
    stability_strategy=DEFAULT_PARAMS["stability_strategy"],
    drop_cols=None,
    random_state=42
):
    """
    Exécute la sélection de variables sur le fichier CSV d'entrée et sauvegarde
    le jeu de données réduit aux features sélectionnées.

    Paramètres
    ----------
    input_csv_path : str
        Chemin du fichier CSV contenant le jeu de données enrichi (avec toutes les features).
    output_csv_path : str, optionnel
        Chemin de sauvegarde du fichier sélectionné. Par défaut, le nom du fichier
        d'entrée avec le suffixe '_selected.csv'.
    target : str
        Nom de la variable cible.
    n_features : int
        Nombre maximum de features à sélectionner.
    ... (autres paramètres)
    drop_cols : list, optionnel
        Liste des colonnes à exclure de la sélection (en plus de celles par défaut).
    random_state : int
        Graine aléatoire pour la reproductibilité.

    Retourne
    --------
    pandas.DataFrame
        Le DataFrame avec les colonnes sélectionnées + target.
    """
    if drop_cols is None:
        drop_cols = DROP_COLS.copy()
    else:
        # Fusionner avec la liste par défaut
        drop_cols = list(set(DROP_COLS) | set(drop_cols))

    if output_csv_path is None:
        output_csv_path = input_csv_path.replace('.csv', '_selected.csv')

    print("\n" + "="*60)
    print("   FEATURE SELECTION EXECUTION")
    print("="*60)
    print(f"Fichier d'entrée : {input_csv_path}")
    print(f"Fichier de sortie : {output_csv_path}")

    # 1. Lecture
    train_data = pd.read_csv(input_csv_path)
    print(f"Dataset : {train_data.shape[0]} observations x {train_data.shape[1]} features")

    # 2. Colonnes candidates (exclure drop_cols)
    feature_cols = [c for c in train_data.columns if c not in drop_cols and c != target]
    print(f"Feature candidates : {len(feature_cols)}")

    X_dev = train_data[feature_cols].copy()
    y_dev = train_data[target].copy()

    # 3. Variance threshold
    vt = VarianceThreshold(threshold=var_threshold)
    X_dev_vt = vt.fit_transform(X_dev)
    feature_cols_vt = np.array(feature_cols)[vt.get_support()].tolist()
    print(f"After variance threshold : {len(feature_cols_vt)} features")

    # 4. Correlation filter
    X_dev_vt_df = pd.DataFrame(X_dev_vt, columns=feature_cols_vt)
    target_corr = X_dev_vt_df.corrwith(pd.Series(np.array(y_dev), name='y')).abs()
    corr_matrix = X_dev_vt_df.corr().abs()
    upper = corr_matrix.where(np.triu(np.ones(corr_matrix.shape), k=1).astype(bool))

    to_drop = set()
    for col in upper.columns:
        correlated_with = upper.index[upper[col] > corr_threshold].tolist()
        for row in correlated_with:
            if target_corr.get(col, 0) >= target_corr.get(row, 0):
                to_drop.add(row)
            else:
                to_drop.add(col)

    cols_to_keep = [c for c in feature_cols_vt if c not in to_drop]
    col_to_idx = {col: i for i, col in enumerate(feature_cols_vt)}
    keep_idx = [col_to_idx[c] for c in cols_to_keep]
    X_dev_clean = X_dev_vt[:, keep_idx]
    y_dev_arr = np.array(y_dev)
    print(f"After correlation filter : {len(cols_to_keep)} features")
    if to_drop:
        print(f"  Dropped ({len(to_drop)}) : {sorted(to_drop)}")

    # 5. LassoCV pour trouver alpha optimal
    scaler_lasso = RobustScaler()
    X_dev_scaled = scaler_lasso.fit_transform(X_dev_clean)
    lasso_cv = LassoCV(cv=n_splits, random_state=random_state, n_jobs=-1, max_iter=max_iter)
    lasso_cv.fit(X_dev_scaled, y_dev_arr)
    best_alpha = lasso_cv.alpha_
    print(f"LassoCV best alpha : {best_alpha:.6f}")

    # 6. Configurer les modèles pour l'importance
    models_config = {
        "Lasso": (
            lambda: Lasso(alpha=best_alpha, max_iter=max_iter, random_state=random_state),
            True
        ),
        "LightGBM": (
            lambda: LGBMRegressor(random_state=random_state, n_jobs=-1, verbose=-1),
            False
        ),
        "CatBoost": (
            lambda: CatBoostRegressor(random_seed=random_state, verbose=0),
            False
        ),
        "GradientBoosting": (
            lambda: GradientBoostingRegressor(random_state=random_state),
            False
        ),
    }

    all_importances = {}
    print(f"\nCalcul des importances ({len(models_config)} modèles) ...")
    for name, (build_fn, needs_scale) in models_config.items():
        try:
            imp = get_importances_kfold(
                model_name=name,
                build_model_fn=build_fn,
                X=X_dev_clean,
                y=y_dev_arr,
                feature_names=cols_to_keep,
                n_splits=n_splits,
                needs_scale_arg=needs_scale,
                random_state=random_state
            )
            all_importances[name] = imp
        except Exception as e:
            print(f"  Ignoré {name} : {e}")

    if not all_importances:
        raise RuntimeError("Aucun modèle n'a fourni d'importances. Vérifiez vos données.")

    # 7. Agrégation Borda
    borda_scores = defaultdict(float)
    for name, imp in all_importances.items():
        for rank_0based, feat in enumerate(imp.index):
            borda_scores[feat] += (1.0 / (rank_0based + 1)) / len(cols_to_keep)
    borda_series = pd.Series(borda_scores).sort_values(ascending=False)

    # 8. Sélection automatique du nombre de features (couverture Borda)
    cumulative_ratio = borda_series.cumsum() / borda_series.sum()
    N_auto = int((cumulative_ratio < borda_coverage).sum()) + 1
    N_final = min(N_auto, n_features)
    print(f"Auto selection : {N_auto} features cover {borda_coverage:.0%} → N_final = {N_final}")
    top_features = borda_series.head(N_final).index.tolist()

    # 9. Bootstrap stability
    rng = np.random.default_rng(random_state)
    n_samples = X_dev_clean.shape[0]
    selection_counts = defaultdict(int)

    for b in range(n_bootstrap):
        idx = rng.integers(0, n_samples, size=n_samples)
        X_boot = X_dev_clean[idx]
        y_boot = y_dev_arr[idx]

        boot_importances = {}
        for name, (build_fn, needs_scale) in models_config.items():
            try:
                imp = get_importances_kfold(
                    model_name=name,
                    build_model_fn=build_fn,
                    X=X_boot,
                    y=y_boot,
                    feature_names=cols_to_keep,
                    n_splits=n_splits,
                    needs_scale_arg=needs_scale,
                    random_state=random_state
                )
                boot_importances[name] = imp
            except Exception:
                pass

        if not boot_importances:
            continue

        boot_borda = defaultdict(float)
        for name, imp in boot_importances.items():
            for rank_0, feat in enumerate(imp.index):
                boot_borda[feat] += (1.0 / (rank_0 + 1)) / len(cols_to_keep)

        boot_borda_series = pd.Series(boot_borda).sort_values(ascending=False)
        selected_boot = boot_borda_series.head(N_final).index.tolist()
        for feat in selected_boot:
            selection_counts[feat] += 1

        if (b + 1) % 10 == 0:
            print(f"Bootstrap {b + 1}/{n_bootstrap} done")

    stability_series = (
        pd.Series(selection_counts)
        .reindex(borda_series.index, fill_value=0)
        .sort_values(ascending=False)
    )
    stability_freq = stability_series / n_bootstrap

    # 10. Application du filtre de stabilité
    if stability_strategy == 'strict':
        top_features_final = [
            f for f in top_features
            if stability_freq.get(f, 0) >= stability_threshold
        ]
    elif stability_strategy == 'hybrid':
        hybrid_scores = borda_series * stability_freq.reindex(borda_series.index, fill_value=0)
        hybrid_scores = hybrid_scores.sort_values(ascending=False)
        top_features_final = hybrid_scores.head(N_final).index.tolist()
    else:
        raise ValueError(f"Unknown stability_strategy: {stability_strategy}")

    # 11. Rapport final
    print("\n" + "="*60)
    print("RÉSULTATS DE LA SÉLECTION")
    print("="*60)
    print(f"Nombre de features retenues : {len(top_features_final)}")
    print("Features sélectionnées :")
    for i, f in enumerate(top_features_final, 1):
        print(f"  {i:>2}. {f}  (stability = {stability_freq.get(f, 0):.0%})")

    # 12. Sauvegarde du dataset filtré
    train_selected = train_data[top_features_final + [target]].copy()
    train_selected.to_csv(output_csv_path, index=False)
    print(f"\nJeu de données sélectionné sauvegardé : {output_csv_path}")
    print(f"Shape : {train_selected.shape}")

    return train_selected


if __name__ == "__main__":
    # Exemple d'utilisation en ligne de commande (debug)
    import sys
    if len(sys.argv) < 2:
        print("Usage: python feature_selection.py <input_csv> [output_csv]")
        sys.exit(1)
    input_file = sys.argv[1]
    output_file = sys.argv[2] if len(sys.argv) > 2 else None
    select_features(input_file, output_file)