#!/usr/bin/env python3
"""C18: DoubleEnsemble sample reweighting — target up_precision.

Two-stage walk-forward per window:
  Stage 1: fast RF to score each training sample's difficulty.
  Weight:  upweight false-UP samples (true≠UP, pred=UP) and boundary cases.
  Stage 2: retrain RF + LGBM with sample_weight; blend 1:1.

Rationale: RF/LGBM support sample_weight natively.  By penalising the model
for false-UP predictions on training data it sees fewer of them on test data,
directly improving up_precision without changing label quality or features.
"""
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))

import warnings; warnings.filterwarnings("ignore")
import json
import argparse
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.preprocessing import StandardScaler
from lightgbm import LGBMClassifier

from scripts.improvement_harness import (
    fetch_df, build_triple_barrier_labels, compute_metrics,
    CURRENT_FEATURES, DEFAULT_STOCKS, EXTENDED_STOCKS,
    PASS_DIR_ACC, PASS_UP_PREC,
    MIN_TRAIN, STEP, LABEL_HORIZON, RF_PARAMS,
)
from models.predictor import _build_features

CLASSES = [-1, 0, 1]   # DOWN, HOLD, UP
UP_IDX  = 2            # column index for P(UP) in proba array

# Stage-1 RF is lighter — we just need decent in-sample discrimination
STAGE1_PARAMS = {**RF_PARAMS, "n_estimators": 50}

LGBM_PARAMS = dict(
    n_estimators=150, max_depth=6, learning_rate=0.05,
    class_weight="balanced", random_state=42, n_jobs=1, verbose=-1,
)

# Reweighting hyper-parameters — searched over alpha in main()
FALSE_UP_ALPHA   = 2.0   # multiplier added for false-UP samples
NEAR_UP_BETA     = 0.8   # multiplier for near-miss UP (high P(UP) but true≠UP)
NEAR_UP_THRESH   = 0.35  # P(UP) threshold to count as near-miss


def _compute_weights(y_true, y_pred, proba, alpha, beta, thresh):
    """Return per-sample weights emphasising false-UP and near-UP boundary cases."""
    weights = np.ones(len(y_true))
    up_prob = proba[:, UP_IDX]

    for i in range(len(y_true)):
        if y_true[i] != 1 and y_pred[i] == 1:
            # False positive for UP on training data — most important to fix
            weights[i] = 1.0 + alpha
        elif y_true[i] != 1 and up_prob[i] >= thresh:
            # Near-false-positive: model almost called UP but didn't
            weights[i] = 1.0 + beta * up_prob[i]
    return weights


def _proba_ordered(clf, X, classes=CLASSES):
    """Return proba columns ordered as CLASSES=[-1,0,1], zero-filling missing classes."""
    raw   = clf.predict_proba(X)
    order = list(clf.classes_)
    p = np.zeros((len(X), len(classes)))
    for col_i, cls in enumerate(classes):
        if cls in order:
            p[:, col_i] = raw[:, order.index(cls)]
    return p


def walk_forward_reweighted(feat_df, label_arr, cols, alpha, beta, thresh):
    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_all = [], []

    cutoff = MIN_TRAIN
    while cutoff + STEP + LABEL_HORIZON <= n:
        train_end = cutoff - LABEL_HORIZON
        if train_end < MIN_TRAIN - LABEL_HORIZON:
            cutoff += STEP; continue

        y_tr  = label_arr[:train_end]
        valid = ~np.isnan(y_tr)
        y_v   = y_tr[valid].astype(int)
        if len(y_v) < 10 or len(np.unique(y_v)) < 2:
            cutoff += STEP; continue

        X_tr = X_all[:train_end][valid]
        scaler = StandardScaler()
        X_tr_s = scaler.fit_transform(X_tr)

        # ── Stage 1: fast RF for difficulty scoring ──────────────────────────
        s1 = RandomForestClassifier(**STAGE1_PARAMS)
        s1.fit(X_tr_s, y_v)

        s1_proba = _proba_ordered(s1, X_tr_s)
        s1_pred  = np.array([CLASSES[i] for i in s1_proba.argmax(axis=1)])
        weights  = _compute_weights(y_v, s1_pred, s1_proba, alpha, beta, thresh)

        # ── Stage 2: retrain RF + LGBM with sample weights ───────────────────
        rf   = RandomForestClassifier(**RF_PARAMS)
        rf.fit(X_tr_s, y_v, sample_weight=weights)

        lgbm = LGBMClassifier(**LGBM_PARAMS)
        lgbm.fit(X_tr_s, y_v, sample_weight=weights)

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

        X_te_s = scaler.transform(X_all[cutoff:test_end][valid_te])

        rf_p   = _proba_ordered(rf,   X_te_s)
        lgbm_p = _proba_ordered(lgbm, X_te_s)
        blended = (rf_p + lgbm_p) / 2.0
        y_pred  = np.array([CLASSES[i] for i in blended.argmax(axis=1)])

        y_true_all.extend(y_te[valid_te].astype(int).tolist())
        y_pred_all.extend(y_pred.tolist())
        cutoff += STEP

    if not y_true_all:
        return {}
    return compute_metrics(np.array(y_true_all), np.array(y_pred_all))


def run_stocks(stocks, alpha, beta=NEAR_UP_BETA, thresh=NEAR_UP_THRESH, verbose=True):
    results, metrics_list = {}, []
    for stock_no in stocks:
        if verbose:
            print(f"  {stock_no}...", end=" ", flush=True)
        df = fetch_df(stock_no)
        if df is None or df.empty:
            if verbose: print("no data")
            continue
        feat   = _build_features(df)
        close  = (df.set_index("date")["close"] if "date" in df.columns else df["close"]).values
        labels = build_triple_barrier_labels(close)
        m = walk_forward_reweighted(feat, labels, CURRENT_FEATURES, alpha, beta, thresh)
        results[stock_no] = m
        metrics_list.append(m)
        if verbose and m:
            print(f"dir={m['dir_accuracy']}% ↑prec={m['up_precision']}%")
    return results, metrics_list


def _avg(metrics_list, key):
    vals = [m[key] for m in metrics_list if m and not np.isnan(m.get(key, float("nan")))]
    return round(float(np.mean(vals)), 1) if vals else float("nan")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--extended", action="store_true")
    parser.add_argument("--alpha-search", action="store_true",
                        help="Grid-search alpha on 5-stock first")
    args = parser.parse_args()

    stocks = EXTENDED_STOCKS if args.extended else DEFAULT_STOCKS
    tag    = "12-stock" if args.extended else "5-stock"

    if args.alpha_search and not args.extended:
        print("\n=== C18: alpha search [5-stock] ===")
        best_alpha, best_up = FALSE_UP_ALPHA, -1.0
        for a in [0.5, 1.0, 1.5, 2.0, 2.5, 3.0]:
            _, ml = run_stocks(DEFAULT_STOCKS, alpha=a, verbose=False)
            up = _avg(ml, "up_precision")
            dr = _avg(ml, "dir_accuracy")
            print(f"  alpha={a:.1f}  dir={dr}%  ↑prec={up}%")
            if up > best_up and dr >= 40.0:
                best_up, best_alpha = up, a
        print(f"\n  Best alpha = {best_alpha}  (↑prec={best_up}%)")
        alpha = best_alpha
    else:
        alpha = FALSE_UP_ALPHA

    print(f"\n=== C18: DoubleEnsemble reweight [alpha={alpha}] [{tag}] ===")
    per_stock, metrics_list = run_stocks(stocks, alpha=alpha)

    avg_dir  = _avg(metrics_list, "dir_accuracy")
    avg_up   = _avg(metrics_list, "up_precision")
    passed   = avg_dir >= PASS_DIR_ACC and avg_up >= PASS_UP_PREC

    print(f"\n  Avg: dir={avg_dir}%  ↑prec={avg_up}%")
    print(f"  Gate (dir≥{PASS_DIR_ACC}% AND ↑prec≥{PASS_UP_PREC}%): {'PASS ✓' if passed else 'FAIL ✗'}")

    result = {
        "experiment":  "C18",
        "description": "DoubleEnsemble sample reweighting (false-UP upweight)",
        "alpha":       alpha, "beta": NEAR_UP_BETA, "thresh": NEAR_UP_THRESH,
        "stocks":      stocks,
        "aggregate":   {"dir_accuracy": avg_dir, "up_precision": avg_up},
        "per_stock":   per_stock,
        "passed":      passed,
        "pass_gate":   {"dir_accuracy": PASS_DIR_ACC, "up_precision": PASS_UP_PREC},
    }

    suffix = "_12stock" if args.extended else ""
    out    = ROOT / f"docs/c18_result{suffix}.json"
    out.parent.mkdir(exist_ok=True)
    out.write_text(json.dumps(result, indent=2))
    print(f"\n  Saved: {out}")


if __name__ == "__main__":
    main()