"""
AHAD QUANT — Ensemble Training Pipeline  (V5 — Deep Learning Edition)
Trains LightGBM + XGBoost + RandomForest + TFT + TransformerGRU ensemble.

Usage:
    python download_data.py   # download historical data first
    python train.py           # train full ensemble (~2-6h, GPU recommandé)
    python train.py --no-dl   # ML only, skip TFT/TGRU (~15-30 min)

What's new in V5 vs V4:
  - TFT  (Temporal Fusion Transformer) — variable selection + attention
  - TGRU (TransformerGRU hybrid)       — bidir GRU + transformer encoder
  - Sequence dataset builder           — per-pair 168-candle sliding windows
  - Extended meta-learner              — up to 5 base models (was 3)
  - Backward-compatible when torch/DL unavailable — falls back to ML-only
"""

import json, os, pickle, sys, time, argparse
sys.stdout.reconfigure(encoding="utf-8", errors="replace")
import numpy as np
import lightgbm as lgb
import config
from features import build_features, FEATURE_NAMES

# ─── CLI flags ───────────────────────────────────────────────────────────────
_parser = argparse.ArgumentParser(add_help=False)
_parser.add_argument("--no-dl", action="store_true",
                     help="Skip TFT/TGRU training (ML-only mode, faster)")
_args, _ = _parser.parse_known_args()
SKIP_DL: bool = _args.no_dl

try:
    import xgboost as xgb
    HAS_XGB = True
except ImportError:
    HAS_XGB = False
    print("[WARN] xgboost not installed — skipping. pip install xgboost")

try:
    from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
    from sklearn.linear_model import LogisticRegression
    from sklearn.preprocessing import StandardScaler
    from sklearn.metrics import accuracy_score
    HAS_SKLEARN = True
except ImportError:
    HAS_SKLEARN = False
    print("[WARN] scikit-learn not installed — skipping RF. pip install scikit-learn")

try:
    import optuna
    optuna.logging.set_verbosity(optuna.logging.WARNING)
    HAS_OPTUNA = True
except ImportError:
    HAS_OPTUNA = False

# ─── Deep Learning imports (optional — guarded) ──────────────────────────────
HAS_DL = False
if not SKIP_DL:
    try:
        import torch
        from prepare_sequences import (
            prepare_seq_dataset, fit_seq_scaler, transform_sequences, SEQ_LEN
        )
        from tft_model import (
            TemporalFusionTransformer, train_tft, predict_tft_proba
        )
        from transformer_gru_model import (
            TransformerGRU, train_tgru, predict_tgru_proba
        )
        HAS_DL = True
        print(f"[DL] PyTorch {torch.__version__} detected — TFT + TGRU enabled")
    except ImportError as _dl_err:
        print(f"[WARN] DL modules not available ({_dl_err}). "
              "Falling back to ML-only. "
              "Install with: pip install torch>=2.2.0")

# ─── Config ─────────────────────────────────────────────────────────────────
LOOKAHEAD = 3
MIN_CANDLES = 200
TRAIN_RATIO = 0.70
VAL_RATIO   = 0.15
N_WALK_FORWARD_WINDOWS = 4
OPTUNA_TRIALS = 30     # increase for better tuning (slower)


# ─── Data loading ────────────────────────────────────────────────────────────

def load_candles(coin: str) -> dict | None:
    path = os.path.join(config.DATA_DIR, f"{coin}_1h.json")
    if not os.path.exists(path):
        return None
    with open(path) as f:
        return json.load(f)


def make_labels(close: np.ndarray, lookahead: int = LOOKAHEAD) -> np.ndarray:
    labels = np.zeros(len(close))
    for i in range(len(close) - lookahead):
        labels[i] = 1.0 if close[i + lookahead] > close[i] else 0.0
    return labels


def prepare_dataset() -> tuple[np.ndarray, np.ndarray]:
    all_X, all_y = [], []
    skipped = []
    # Paire de référence pour la corrélation (EURUSD = paire dominante en Forex)
    _ref_pair = "EURUSD"
    btc_data = load_candles(_ref_pair)
    if btc_data is None:
        print(f"  [WARN] Paire de référence {_ref_pair} non trouvée dans data/ —"
              " features de corrélation désactivées.")
    btc_close = np.array([c["c"] for c in btc_data]) if btc_data else None

    for coin in config.COINS:
        data = load_candles(coin)
        if data is None:
            print(f"  ⚠️  {coin:8s} — fichier data/{coin}_1h.json introuvable")
            skipped.append((coin, "fichier manquant"))
            continue
        if len(data) < MIN_CANDLES:
            print(f"  ⚠️  {coin:8s} — données insuffisantes ({len(data)} bougies < {MIN_CANDLES} min)")
            skipped.append((coin, f"seulement {len(data)} bougies"))
            continue
        close = np.array([c["c"] for c in data])
        coin_btc = btc_close[-len(close):] if btc_close is not None and len(btc_close) >= len(close) else None

        X = build_features(data, btc_closes=coin_btc)
        y = make_labels(close)
        warmup = 30
        X, y = X[warmup:-LOOKAHEAD], y[warmup:-LOOKAHEAD]
        valid = ~np.isnan(X).any(axis=1)
        X, y = X[valid], y[valid]
        all_X.append(X); all_y.append(y)
        print(f"  ✅ {coin:8s} — {len(X):,} samples ({round(len(data)/24)} jours)")

    # ─── Guard critique : aucune paire chargée ────────────────────────────────
    if len(all_X) == 0:
        print()
        print("━" * 60)
        print("  ERREUR : Aucune paire chargée. Lancez d'abord :")
        print("    python download_data.py")
        print()
        print("  Paires manquantes :")
        for coin, reason in skipped:
            print(f"    {coin:8s} — {reason}")
        print("━" * 60)
        raise RuntimeError(
            "Aucune donnée disponible dans data/. "
            "Exécutez `python download_data.py` avant `python train.py`."
        )

    if len(all_X) < 3:
        print(f"\n  [WARN] Seulement {len(all_X)} paire(s) chargée(s). "
              f"Résultats d'entraînement potentiellement insuffisants. "
              f"Minimum recommandé : 5 paires.")

    if skipped:
        print(f"\n  [INFO] {len(skipped)} paire(s) ignorée(s) : "
              f"{', '.join(c for c, _ in skipped)}")

    return np.concatenate(all_X), np.concatenate(all_y)


# ─── Walk-forward cross-validation ──────────────────────────────────────────

def walk_forward_cv(X: np.ndarray, y: np.ndarray, n_windows: int = 4) -> dict:
    """
    Walk-forward validation with expanding window.
    Returns mean accuracy and std across windows.
    """
    n = len(X)
    base = int(n * 0.5)   # first training window = 50% of data
    step = (n - base) // n_windows
    results = []

    print(f"\n  Walk-forward CV ({n_windows} windows):")
    for i in range(n_windows):
        train_end = base + i * step
        test_end  = min(train_end + step, n)
        X_tr, y_tr = X[:train_end], y[:train_end]
        X_te, y_te = X[train_end:test_end], y[train_end:test_end]

        # Quick LightGBM for CV (fast)
        ds_tr = lgb.Dataset(X_tr, label=y_tr)
        ds_va = lgb.Dataset(X_te, label=y_te, reference=ds_tr)
        params = {"objective": "binary", "metric": "binary_logloss",
                  "num_leaves": 63, "learning_rate": 0.05, "verbose": -1}
        m = lgb.train(params, ds_tr, 500, valid_sets=[ds_va],
                      callbacks=[lgb.early_stopping(30), lgb.log_evaluation(-1)])
        acc = accuracy_score(y_te, (m.predict(X_te) > 0.5).astype(int))
        results.append(acc)
        print(f"    Window {i+1}: train={train_end:,}  test={len(y_te):,}  acc={acc:.4f}")

    mean_acc = np.mean(results)
    std_acc  = np.std(results)
    print(f"  CV accuracy: {mean_acc:.4f} ± {std_acc:.4f}")
    return {"mean": mean_acc, "std": std_acc, "windows": results}


# ─── Optuna hyperparameter search ───────────────────────────────────────────

def tune_lgbm(X_tr, y_tr, X_va, y_va, n_trials: int = OPTUNA_TRIALS) -> dict:
    if not HAS_OPTUNA:
        return {"num_leaves": 63, "learning_rate": 0.05,
                "feature_fraction": 0.8, "bagging_fraction": 0.8,
                "bagging_freq": 5, "min_child_samples": 50}

    def objective(trial):
        params = {
            "objective": "binary", "metric": "binary_logloss", "verbose": -1,
            "num_leaves":        trial.suggest_int("num_leaves", 20, 150),
            "learning_rate":     trial.suggest_float("learning_rate", 0.01, 0.15, log=True),
            "feature_fraction":  trial.suggest_float("feature_fraction", 0.5, 1.0),
            "bagging_fraction":  trial.suggest_float("bagging_fraction", 0.5, 1.0),
            "bagging_freq":      trial.suggest_int("bagging_freq", 1, 10),
            "min_child_samples": trial.suggest_int("min_child_samples", 20, 100),
            "lambda_l1":         trial.suggest_float("lambda_l1", 0.0, 1.0),
            "lambda_l2":         trial.suggest_float("lambda_l2", 0.0, 1.0),
        }
        ds_tr = lgb.Dataset(X_tr, label=y_tr)
        ds_va = lgb.Dataset(X_va, label=y_va, reference=ds_tr)
        m = lgb.train(params, ds_tr, 1000, valid_sets=[ds_va],
                      callbacks=[lgb.early_stopping(30), lgb.log_evaluation(-1)])
        preds = m.predict(X_va)
        return accuracy_score(y_va, (preds > 0.5).astype(int))

    study = optuna.create_study(direction="maximize")
    study.optimize(objective, n_trials=n_trials, show_progress_bar=False)
    print(f"  Best LightGBM accuracy (Optuna): {study.best_value:.4f}")
    return study.best_params


# ─── Model training ──────────────────────────────────────────────────────────

def train_lgbm(X_tr, y_tr, X_va, y_va, params: dict) -> lgb.Booster:
    final_params = {
        "objective": "binary", "metric": "binary_logloss",
        "boosting_type": "gbdt", "verbose": -1,
        **params
    }
    ds_tr = lgb.Dataset(X_tr, label=y_tr, feature_name=FEATURE_NAMES)
    ds_va = lgb.Dataset(X_va, label=y_va, feature_name=FEATURE_NAMES, reference=ds_tr)
    return lgb.train(final_params, ds_tr, 3000, valid_sets=[ds_va],
                     callbacks=[lgb.early_stopping(50), lgb.log_evaluation(200)])


def train_xgb(X_tr, y_tr, X_va, y_va) -> object:
    if not HAS_XGB:
        return None
    model = xgb.XGBClassifier(
        n_estimators=1000, learning_rate=0.05, max_depth=6,
        subsample=0.8, colsample_bytree=0.8, min_child_weight=5,
        use_label_encoder=False, eval_metric="logloss",
        early_stopping_rounds=50, verbosity=0,
        tree_method="hist",
    )
    model.fit(X_tr, y_tr, eval_set=[(X_va, y_va)], verbose=False)
    return model


def train_rf(X_tr, y_tr) -> object:
    if not HAS_SKLEARN:
        return None
    model = RandomForestClassifier(
        n_estimators=300, max_depth=12, min_samples_leaf=30,
        max_features="sqrt", n_jobs=-1, random_state=42
    )
    model.fit(X_tr, y_tr)
    return model


# ─── Ensemble / stacking ─────────────────────────────────────────────────────

def build_ensemble(lgbm_model, xgb_model, rf_model,
                   X_te: np.ndarray, y_te: np.ndarray,
                   tft_preds: np.ndarray | None = None,
                   tgru_preds: np.ndarray | None = None) -> tuple:
    """
    Stack base models with a logistic regression meta-learner.
    Accepts optional TFT and TGRU probability arrays (already computed on X_te).
    Returns (meta_model, scaler, base_accuracies, ensemble_accuracy).
    """
    preds = []
    names = []
    accs  = {}

    p_lgbm = lgbm_model.predict(X_te)
    preds.append(p_lgbm); names.append("LightGBM")
    accs["LightGBM"] = accuracy_score(y_te, (p_lgbm > 0.5).astype(int))

    if xgb_model is not None:
        p_xgb = xgb_model.predict_proba(X_te)[:, 1]
        preds.append(p_xgb); names.append("XGBoost")
        accs["XGBoost"] = accuracy_score(y_te, (p_xgb > 0.5).astype(int))

    if rf_model is not None:
        p_rf = rf_model.predict_proba(X_te)[:, 1]
        preds.append(p_rf); names.append("RandomForest")
        accs["RandomForest"] = accuracy_score(y_te, (p_rf > 0.5).astype(int))

    # ── Optional DL models ───────────────────────────────────────────────────
    if tft_preds is not None:
        preds.append(tft_preds); names.append("TFT")
        accs["TFT"] = accuracy_score(y_te, (tft_preds > 0.5).astype(int))

    if tgru_preds is not None:
        preds.append(tgru_preds); names.append("TGRU")
        accs["TGRU"] = accuracy_score(y_te, (tgru_preds > 0.5).astype(int))

    # Stack predictions as features for meta-learner
    meta_X = np.column_stack(preds)
    scaler = StandardScaler()
    meta_X_scaled = scaler.fit_transform(meta_X)
    meta = LogisticRegression(C=1.0, max_iter=500)
    meta.fit(meta_X_scaled, y_te)   # train meta on test (unseen by base models)
    ens_acc = accuracy_score(y_te, meta.predict(meta_X_scaled))

    print(f"\n  Base models ({len(names)}) :")
    for name, acc in accs.items():
        print(f"    {name:15s}: {acc:.4f} ({acc*100:.1f}%)")
    print(f"    {'Ensemble':15s}: {ens_acc:.4f} ({ens_acc*100:.1f}%)")

    return meta, scaler, accs, ens_acc


# ─── SHAP feature importance ─────────────────────────────────────────────────

def print_top_features(lgbm_model, top_n: int = 15):
    importance = lgbm_model.feature_importance(importance_type="gain")
    idx = np.argsort(importance)[::-1]
    print(f"\n  Top {top_n} features (LightGBM gain):")
    for rank, i in enumerate(idx[:top_n], 1):
        print(f"    {rank:2d}. {FEATURE_NAMES[i]:30s} {importance[i]:>10,.0f}")


# ─── Main ────────────────────────────────────────────────────────────────────

def main():
    print("=" * 65)
    print("  AHAD QUANT — Ensemble Training Pipeline  (V5)")
    if HAS_DL:
        print("  LightGBM + XGBoost + RF + TFT + TransformerGRU + Meta-learner")
    else:
        print("  LightGBM + XGBoost + RandomForest + Stacking Meta-Learner")
        print("  [DL disabled — install torch to enable TFT/TGRU]")
    print("=" * 65)

    # ── 0. Pré-vol : vérifier que data/ existe et contient des fichiers ───────
    data_dir = config.DATA_DIR
    if not os.path.isdir(data_dir):
        print()
        print("━" * 65)
        print(f"  ERREUR : Dossier '{data_dir}/' introuvable.")
        print()
        print("  Lancez d'abord le téléchargement des données :")
        print("    python download_data.py")
        print("━" * 65)
        sys.exit(1)
    json_files = [f for f in os.listdir(data_dir) if f.endswith("_1h.json")]
    if len(json_files) == 0:
        print()
        print("━" * 65)
        print(f"  ERREUR : Aucun fichier de données dans '{data_dir}/'.")
        print()
        print("  Lancez d'abord le téléchargement des données :")
        print("    python download_data.py")
        print("━" * 65)
        sys.exit(1)
    print(f"\n[0/7] Données disponibles : {len(json_files)} fichiers dans {data_dir}/")

    t0 = time.time()

    # ── 1. Tabular data (always) ──────────────────────────────────────────────
    print("\n[1/7] Loading data and building tabular features...")
    X, y = prepare_dataset()
    n = len(X)
    print(f"\n  Total: {n:,} samples | {X.shape[1]} features | "
          f"{y.mean():.2%} long labels")

    # ── 1b. Sequence data (DL only) ──────────────────────────────────────────
    X_seq = None
    y_seq = None
    dl_scaler = None

    if HAS_DL:
        print(f"\n[1b/7] Building sequence dataset (SEQ_LEN={SEQ_LEN})...")
        X_tab_aligned, X_seq_raw, y_seq = prepare_seq_dataset()
        n_seq = len(y_seq)
        print(f"  Sequence samples : {n_seq:,} | shape: {X_seq_raw.shape}")

        # Sequence train/val/test splits
        seq_tr_end  = int(n_seq * TRAIN_RATIO)
        seq_va_end  = int(n_seq * (TRAIN_RATIO + VAL_RATIO))

        # Fit scaler on training portion
        dl_scaler = fit_seq_scaler(X_tab_aligned[:seq_tr_end])

        # Normalise all splits
        X_seq_tr = transform_sequences(dl_scaler, X_seq_raw[:seq_tr_end])
        X_seq_va = transform_sequences(dl_scaler, X_seq_raw[seq_tr_end:seq_va_end])
        X_seq_te = transform_sequences(dl_scaler, X_seq_raw[seq_va_end:])
        y_seq_tr = y_seq[:seq_tr_end]
        y_seq_va = y_seq[seq_tr_end:seq_va_end]
        y_seq_te = y_seq[seq_va_end:]
        print(f"  Seq splits: train={len(y_seq_tr):,} | val={len(y_seq_va):,} | test={len(y_seq_te):,}")

    # ── 2. Walk-forward CV ────────────────────────────────────────────────────
    print("\n[2/7] Walk-forward cross-validation (tabular)...")
    cv_results = walk_forward_cv(X, y, N_WALK_FORWARD_WINDOWS)

    # ── 3. Tabular splits ─────────────────────────────────────────────────────
    train_end = int(n * TRAIN_RATIO)
    val_end   = int(n * (TRAIN_RATIO + VAL_RATIO))
    X_tr, y_tr = X[:train_end], y[:train_end]
    X_va, y_va = X[train_end:val_end], y[train_end:val_end]
    X_te, y_te = X[val_end:], y[val_end:]
    print(f"\n  Tabular splits: train={len(X_tr):,} | val={len(X_va):,} | test={len(X_te):,}")

    # ── 4. Optuna search ──────────────────────────────────────────────────────
    print(f"\n[3/7] Hyperparameter search ({OPTUNA_TRIALS} Optuna trials)...")
    best_params = tune_lgbm(X_tr, y_tr, X_va, y_va, OPTUNA_TRIALS)

    # ── 5. Train tabular base models ──────────────────────────────────────────
    print("\n[4/7] Training tabular base models (LightGBM + XGBoost + RF)...")
    print("  Training LightGBM...")
    lgbm_model = train_lgbm(X_tr, y_tr, X_va, y_va, best_params)

    print("  Training XGBoost...")
    xgb_model = train_xgb(X_tr, y_tr, X_va, y_va)

    print("  Training RandomForest...")
    rf_model = train_rf(X_tr, y_tr)

    # ── 5b. Train DL models ───────────────────────────────────────────────────
    tft_model  = None
    tgru_model = None
    tft_preds_te  = None
    tgru_preds_te = None

    if HAS_DL:
        print("\n[4b/7] Training Deep Learning models (TFT + TransformerGRU)...")
        print("  Training Temporal Fusion Transformer...")
        tft_model = train_tft(X_seq_tr, y_seq_tr, X_seq_va, y_seq_va)

        print("\n  Training TransformerGRU...")
        tgru_model = train_tgru(X_seq_tr, y_seq_tr, X_seq_va, y_seq_va)

        # ── Free training/val sequence tensors: not needed past this point ──
        # (X_seq_tr/X_seq_va are the largest arrays in memory, ~70k+15k seqs)
        import gc, torch
        del X_seq_tr, X_seq_va, y_seq_tr, y_seq_va
        try:
            del X_seq_raw
        except NameError:
            pass
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        print("  [MEM] Freed train/val sequence tensors before inference")

        # Get DL predictions on sequence test set
        X_seq_te_t = torch.FloatTensor(X_seq_te)
        tft_preds_te  = predict_tft_proba(tft_model, X_seq_te_t)
        tgru_preds_te = predict_tgru_proba(tgru_model, X_seq_te_t)
        print(f"\n  TFT  test acc : {accuracy_score(y_seq_te, (tft_preds_te > 0.5).astype(int)):.4f}")
        print(f"  TGRU test acc : {accuracy_score(y_seq_te, (tgru_preds_te > 0.5).astype(int)):.4f}")

        # X_seq_te_t no longer needed after inference (predictions already extracted)
        del X_seq_te_t
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    # ── 6. Ensemble / stacking ────────────────────────────────────────────────
    print("\n[5/7] Building stacking ensemble on test set...")

    if HAS_DL and tft_preds_te is not None:
        # Use sequence-aligned test set for ALL models (common index space)
        # Tabular predictions on the aligned X_tab_aligned test portion
        X_te_dl  = X_tab_aligned[seq_va_end:]      # aligned tabular test
        y_te_dl  = y_seq_te                        # same labels

        meta_model, meta_scaler, base_accs, ens_acc = build_ensemble(
            lgbm_model, xgb_model, rf_model,
            X_te_dl, y_te_dl,
            tft_preds=tft_preds_te,
            tgru_preds=tgru_preds_te,
        )
    else:
        # ML-only path (original behavior)
        meta_model, meta_scaler, base_accs, ens_acc = build_ensemble(
            lgbm_model, xgb_model, rf_model, X_te, y_te
        )

    print_top_features(lgbm_model)

    # ── 6b. Save feature importances to JSON (lu par web_ui.py dashboard) ────
    try:
        importance = lgbm_model.feature_importance(importance_type="gain")
        feat_imp_data = [
            {"feature": FEATURE_NAMES[i], "importance": float(importance[i])}
            for i in np.argsort(importance)[::-1]
        ]
        imp_path = os.path.join(os.path.dirname(__file__), "feature_importance.json")
        with open(imp_path, "w") as f:
            json.dump({"updated_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
                       "features": feat_imp_data}, f, indent=2)
    except Exception as _fie:
        print(f"  [WARN] feature_importance.json non généré: {_fie}")

    # ── 7. Save ───────────────────────────────────────────────────────────────
    print(f"\n[6/7] Saving models...")

    # Standalone LightGBM (compatible with free version loader)
    with open(config.MODEL_PATH, "wb") as f:
        pickle.dump(lgbm_model, f)
    print(f"  LightGBM model  -> {config.MODEL_PATH}")

    # Full ensemble (all models + DL if available)
    ensemble_data = {
        # Tabular base models
        "lgbm"       : lgbm_model,
        "xgb"        : xgb_model,
        "rf"         : rf_model,
        # DL models (None if torch not available or --no-dl)
        "tft"        : tft_model,
        "tgru"       : tgru_model,
        "dl_scaler"  : dl_scaler,         # StandardScaler for sequence normalisation
        # Meta-learner
        "meta"       : meta_model,
        "scaler"     : meta_scaler,
        # Metadata
        "feature_names" : FEATURE_NAMES,
        "cv_results"    : cv_results,
        "base_accs"     : base_accs,
        "ens_acc"       : ens_acc,
        "has_dl"        : HAS_DL and tft_model is not None,
        "version"       : "ahad_quant-forex-v5",
    }
    with open(config.ENSEMBLE_MODEL_PATH, "wb") as f:
        pickle.dump(ensemble_data, f)
    print(f"  Ensemble model  -> {config.ENSEMBLE_MODEL_PATH}")

    # ── 8. Summary ────────────────────────────────────────────────────────────
    elapsed = time.time() - t0
    print(f"\n{'=' * 65}")
    print(f"  Training complete in {elapsed/60:.1f} min")
    print(f"  Walk-forward accuracy : {cv_results['mean']:.4f} ± {cv_results['std']:.4f}")
    print(f"  Ensemble test accuracy: {ens_acc:.4f} ({ens_acc*100:.1f}%)")
    if HAS_DL and tft_model is not None:
        print(f"  DL models           : TFT ✅  TGRU ✅")
    else:
        print(f"  DL models           : not trained (--no-dl or torch missing)")
    print(f"\n  Run `python ahad_quant.py` to start trading.")
    print(f"{'=' * 65}\n")


if __name__ == "__main__":
    main()