#!/usr/bin/env python3
"""C16: Stacking ensemble — RF + XGB + LGBM base models, XGBoost meta-learner.
Instead of averaging base model probabilities, a meta-learner learns the optimal
combination from held-out base predictions.

Meta-train: last 60 rows of each training window (before test period).
Meta-features: [rf_p(-1), rf_p(0), rf_p(1), xgb_p(-1..1), lgbm_p(-1..1)] = 9 dims.
Meta-learner: XGBoostClassifier (light params, fast).

Pass gate: dir_accuracy > 43.5% AND up_precision > 54%.
"""
import sys
from pathlib import Path
ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))
try:
    from dotenv import load_dotenv; load_dotenv(ROOT / ".env")
except ImportError:
    pass

import warnings; warnings.filterwarnings("ignore")
import json
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from xgboost import XGBClassifier
from lightgbm import LGBMClassifier

sys.path.insert(0, str(ROOT / "scripts"))
from improvement_harness import (
    BASELINE_FEATURES, DEFAULT_STOCKS, RF_PARAMS,
    fetch_df, build_triple_barrier_labels, compute_metrics,
    walk_forward,
    MIN_TRAIN, STEP, LABEL_HORIZON,
)
from models.predictor import _build_features

STOCKS = DEFAULT_STOCKS
CLASSES = np.array([-1, 0, 1])
META_WINDOW = 60  # rows used to train meta-learner per step


def _fit_base(X_tr, y_v):
    """Train RF + XGB + LGBM, return fitted list."""
    rf = RandomForestClassifier(
        n_estimators=200, max_depth=6, min_samples_leaf=10,
        class_weight="balanced", random_state=42, n_jobs=-1
    )
    xgb = XGBClassifier(
        n_estimators=150, max_depth=4, learning_rate=0.05,
        subsample=0.8, colsample_bytree=0.8,
        use_label_encoder=False, eval_metric="mlogloss",
        random_state=42, verbosity=0, n_jobs=-1
    )
    lgbm = LGBMClassifier(
        n_estimators=150, max_depth=4, learning_rate=0.05,
        subsample=0.8, colsample_bytree=0.8,
        class_weight="balanced", random_state=42, verbose=-1, n_jobs=-1
    )
    rf.fit(X_tr, y_v)
    xgb.fit(X_tr, y_v + 1)   # XGB needs 0-indexed labels
    lgbm.fit(X_tr, y_v)
    return rf, xgb, lgbm


def _base_proba(rf, xgb, lgbm, X) -> np.ndarray:
    """Return (n, 9) meta-features: 3 probs × 3 models."""
    def _align(clf, X_in, shift=0):
        p = clf.predict_proba(X_in)
        out = np.zeros((len(X_in), 3))
        cls = list(clf.classes_)
        for j, c in enumerate(CLASSES):
            idx = cls.index(c + shift) if (c + shift) in cls else -1
            if idx >= 0:
                out[:, j] = p[:, idx]
        return out

    rf_p   = _align(rf,   X, shift=0)
    xgb_p  = _align(xgb,  X, shift=1)   # XGB labels were shifted +1
    lgbm_p = _align(lgbm, X, shift=0)
    return np.hstack([rf_p, xgb_p, lgbm_p])  # (n, 9)


def walk_forward_stacked(feat_df, label_arr, cols):
    avail  = [c for c in cols if c in feat_df.columns]
    X_all  = feat_df[avail].fillna(0).values
    n      = len(feat_df)
    y_true_all, y_pred_avg_all, y_pred_stack_all = [], [], []

    cutoff = MIN_TRAIN
    while cutoff + STEP + LABEL_HORIZON <= n:
        y_tr  = label_arr[:cutoff]
        valid = ~np.isnan(y_tr)
        y_v   = y_tr[valid].astype(int)
        if len(y_v) < 40 or len(np.unique(y_v)) < 2:
            cutoff += STEP; continue

        # ── Base train window (everything before meta window) ──────────────────
        meta_start = max(0, valid.sum() - META_WINDOW)
        # Indices in the original (unfiltered) space
        valid_idx = np.where(valid)[0]
        if len(valid_idx) <= META_WINDOW + 10:
            cutoff += STEP; continue

        base_idx = valid_idx[:meta_start]
        meta_idx = valid_idx[meta_start:]

        X_base, y_base = X_all[base_idx], y_v[:meta_start]
        X_meta, y_meta = X_all[meta_idx], y_v[meta_start:]

        if len(np.unique(y_base)) < 2:
            cutoff += STEP; continue

        try:
            rf, xgb, lgbm = _fit_base(X_base, y_base)
        except Exception:
            cutoff += STEP; continue

        # ── Meta-learner ──────────────────────────────────────────────────────
        meta_feats = _base_proba(rf, xgb, lgbm, X_meta)
        meta_clf = XGBClassifier(
            n_estimators=50, max_depth=3, learning_rate=0.1,
            use_label_encoder=False, eval_metric="mlogloss",
            random_state=42, verbosity=0, n_jobs=1
        )
        try:
            meta_clf.fit(meta_feats, y_meta + 1)  # shift for XGB
        except Exception:
            cutoff += STEP; continue

        # ── Test window ───────────────────────────────────────────────────────
        test_end = min(cutoff + STEP, n - LABEL_HORIZON)
        X_te = X_all[cutoff:test_end]
        y_te = label_arr[cutoff:test_end]
        valid_te = ~np.isnan(y_te)
        if valid_te.sum() == 0:
            cutoff += STEP; continue

        X_te_v = X_te[valid_te]

        # Average ensemble prediction
        avg_p = (
            _align_proba(rf,   X_te_v, shift=0) +
            _align_proba(xgb,  X_te_v, shift=1) +
            _align_proba(lgbm, X_te_v, shift=0)
        ) / 3
        y_avg = CLASSES[np.argmax(avg_p, axis=1)]

        # Stacked prediction
        te_meta_feats = _base_proba(rf, xgb, lgbm, X_te_v)
        stack_proba = meta_clf.predict_proba(te_meta_feats)
        # meta_clf classes are 0,1,2 → map to -1,0,1
        y_stack = CLASSES[np.argmax(stack_proba, axis=1)]

        y_true_all.extend(y_te[valid_te].tolist())
        y_pred_avg_all.extend(y_avg.tolist())
        y_pred_stack_all.extend(y_stack.tolist())
        cutoff += STEP

    if not y_true_all:
        return {}, {}
    yt = np.array(y_true_all)
    return (
        compute_metrics(yt, np.array(y_pred_avg_all)),
        compute_metrics(yt, np.array(y_pred_stack_all)),
    )


def _align_proba(clf, X, shift=0):
    p = clf.predict_proba(X)
    out = np.zeros((len(X), 3))
    cls = list(clf.classes_)
    for j, c in enumerate(CLASSES):
        idx = cls.index(c + shift) if (c + shift) in cls else -1
        if idx >= 0:
            out[:, j] = p[:, idx]
    return out


CURRENT_FEATURES = [f for f in BASELINE_FEATURES if f not in {
    "macd_cross_up", "macd_cross_down", "price_volume_div",
    "foreign_net_vol_ratio", "trust_net_vol_ratio", "dealer_net_vol_ratio",
    "institutional_net_vol_ratio", "institutional_5d_net_vol_ratio",
    "institutional_20d_zscore", "foreign_trust_alignment", "institutional_streak",
}]

PASS_DIR_ACC = 43.5
PASS_UP_PREC = 54.0


def _mean(rows, field):
    vals = [r[field] for r in rows if isinstance(r.get(field), (int, float))
            and not np.isnan(r.get(field, float("nan")))]
    return round(sum(vals) / len(vals), 1) if vals else float("nan")


def main():
    hdr = f"{'Stock':>6}  {'Model':>10}  {'Acc%':>5}  {'Dir%':>5}  {'↑Prec%':>7}  {'Signals':>7}"
    print(f"\n{hdr}\n{'-'*len(hdr)}")

    agg_avg, agg_stack = [], []
    results = {}

    for stock_no in STOCKS:
        print(f"  computing {stock_no}...", end="\r", flush=True)
        df = fetch_df(stock_no)
        if df is None or df.empty:
            print(f"{stock_no:>6}  no data"); continue

        feat   = _build_features(df)
        close  = df["close"].values if "date" not in df.columns else df.set_index("date")["close"].values
        labels = build_triple_barrier_labels(close)

        r_avg, r_stack = walk_forward_stacked(feat, labels, CURRENT_FEATURES)
        if not r_avg:
            continue

        agg_avg.append(r_avg)
        agg_stack.append(r_stack)
        results[stock_no] = {"avg": r_avg, "stack": r_stack}

        for name, r in [("avg_ens", r_avg), ("stacked", r_stack)]:
            prefix = f"{stock_no:>6}" if name == "avg_ens" else f"{'':>6}"
            print(f"{prefix}  {name:>10}  {r.get('accuracy',0):>5.1f}  "
                  f"{r.get('dir_accuracy',0):>5.1f}  {r.get('up_precision',0):>7.1f}  "
                  f"{r.get('n_signals',0):>7}")

        da = r_stack.get("dir_accuracy", 0) - r_avg.get("dir_accuracy", 0)
        dp = r_stack.get("up_precision", 0) - r_avg.get("up_precision", 0)
        print(f"{'':>6}  {'':>10}  {'Δ':>5}  {da:>+5.1f}  {dp:>+7.1f}")

    agg_a = {"dir_accuracy": _mean(agg_avg,   "dir_accuracy"),
             "up_precision": _mean(agg_avg,   "up_precision")}
    agg_s = {"dir_accuracy": _mean(agg_stack, "dir_accuracy"),
             "up_precision": _mean(agg_stack, "up_precision")}

    print(f"\n=== AGGREGATE ===")
    print(f"  avg_ensemble: dir={agg_a['dir_accuracy']}%  ↑prec={agg_a['up_precision']}%")
    print(f"  stacked:      dir={agg_s['dir_accuracy']}%  ↑prec={agg_s['up_precision']}%")

    dir_delta  = (agg_s["dir_accuracy"] or 0) - (agg_a["dir_accuracy"] or 0)
    prec_delta = (agg_s["up_precision"] or 0) - (agg_a["up_precision"] or 0)

    passed = (
        (agg_s["dir_accuracy"] or 0) >= PASS_DIR_ACC and
        (agg_s["up_precision"]  or 0) >= PASS_UP_PREC and
        dir_delta >= -0.5 and prec_delta >= -0.5
    )

    print(f"\n── Results ──────────────────────────────────────")
    print(f"  Δ dir_acc={dir_delta:+.1f}pp  Δ up_prec={prec_delta:+.1f}pp")
    print(f"  Pass gate: dir>{PASS_DIR_ACC}% AND up_prec>{PASS_UP_PREC}%")
    print(f"  Result: {'PASSED' if passed else 'FAILED'}")

    out = {"experiment": "C16_stacking", "aggregate": {"avg": agg_a, "stacked": agg_s},
           "results": results, "passed": passed}
    Path("docs").mkdir(exist_ok=True)
    with open("docs/c16_result.json", "w") as f:
        json.dump(out, f, indent=2, default=str)
    print("  Saved → docs/c16_result.json")


if __name__ == "__main__":
    main()