#!/usr/bin/env python3
"""C2 validation: BASELINE_FEATURES vs BASELINE + PTT sentiment."""
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 numpy as np
import pandas as pd
import json
from pathlib import Path
from sklearn.ensemble import RandomForestClassifier

sys.path.insert(0, str(ROOT / "scripts"))
from improvement_harness import (
    BASELINE_FEATURES, fetch_df, build_triple_barrier_labels, compute_metrics, RF_PARAMS,
)
from models.predictor import _build_features
from data.ptt_sentiment import add_ptt_sentiment, scrape_ptt_sentiment

STOCKS      = ["2330", "0050", "2317", "2454", "2881"]
PTT_FEATS   = BASELINE_FEATURES + ["ptt_sentiment_1d", "ptt_sentiment_5d_ma"]
MIN_TRAIN   = 60    # reduced: PTT window ~180 days
STEP        = 10
LABEL_HORIZON = 5


def walk_forward_ptt(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_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]
        if len(y_v) < 10 or len(np.unique(y_v)) < 2:
            cutoff += STEP; continue
        clf = RandomForestClassifier(**RF_PARAMS)
        clf.fit(X_all[:cutoff][valid], y_v.astype(int))
        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
        y_pred = clf.predict(X_all[cutoff:test_end][valid_te])
        y_true_all.extend(y_te[valid_te].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 main():
    # Pre-warm PTT cache (scrape once, reuse)
    print("Scraping PTT sentiment data (180 days)...", flush=True)
    try:
        scrape_ptt_sentiment(days_back=180)
    except Exception as e:
        print(f"  PTT scrape failed: {e} — proceeding with zeros")

    per_stock = {}
    agg = {"baseline": [], "ptt": []}
    hdr = f"{'Stock':>6}  {'Set':>10}  {'Acc%':>5}  {'Dir%':>5}  {'↑Prec%':>7}  {'PTT rows':>8}"
    print(f"\n{hdr}\n{'-'*len(hdr)}")

    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

        # Add PTT sentiment
        df = add_ptt_sentiment(df, stock_no)

        # Use only last 180 days (PTT window)
        if "date" in df.columns:
            cutoff_date = pd.Timestamp.now() - pd.Timedelta(days=180)
            df_recent = df[pd.to_datetime(df["date"]) >= cutoff_date].copy()
        else:
            df_recent = df.tail(130).copy()

        if len(df_recent) < MIN_TRAIN + STEP + LABEL_HORIZON:
            print(f"{stock_no:>6}  insufficient recent data ({len(df_recent)} rows)")
            continue

        # Count non-zero PTT rows
        ptt_rows = int((df_recent.get("ptt_sentiment_1d", pd.Series([0])) != 0).sum())

        feat   = _build_features(df_recent)
        # Copy PTT cols into feat (not built by _build_features)
        for col in ["ptt_sentiment_1d", "ptt_sentiment_5d_ma"]:
            if col in df_recent.columns:
                feat[col] = df_recent[col].values[:len(feat)]

        close  = (df_recent.set_index("date")["close"] if "date" in df_recent.columns
                  else df_recent["close"]).values
        labels = build_triple_barrier_labels(close)

        per_stock[stock_no] = {}
        for name, cols in [("baseline", BASELINE_FEATURES), ("ptt", PTT_FEATS)]:
            r = walk_forward_ptt(feat, labels, cols)
            per_stock[stock_no][name] = r
            if r:
                agg[name].append(r)
                prefix = f"{stock_no:>6}" if name == "baseline" else f"{'':>6}"
                print(f"{prefix}  {name:>10}  {r['accuracy']:>5.1f}  {r['dir_accuracy']:>5.1f}  "
                      f"{r['up_precision']:>7.1f}  {ptt_rows if name=='ptt' else '':>8}")

        if per_stock[stock_no].get("baseline") and per_stock[stock_no].get("ptt"):
            rb, rp = per_stock[stock_no]["baseline"], per_stock[stock_no]["ptt"]
            print(f"{'':>6}  {'Δ':>10}  {'':>5}  {rp['dir_accuracy']-rb['dir_accuracy']:>+5.1f}  "
                  f"{rp['up_precision']-rb['up_precision']:>+7.1f}")
        print()

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

    print("=== AGGREGATE ===")
    agg_summary = {}
    for name in ("baseline","ptt"):
        s = {k: mf(agg[name], k) for k in ("accuracy","dir_accuracy","up_precision","dn_precision")}
        agg_summary[name] = s
        print(f"  {name:>10}: acc={s['accuracy']}%  dir={s['dir_accuracy']}%  ↑prec={s['up_precision']}%")

    b, p = agg_summary.get("baseline",{}), agg_summary.get("ptt",{})
    no_regress  = (p.get("dir_accuracy",0) >= b.get("dir_accuracy",0) - 0.5 and
                   p.get("up_precision",0) >= b.get("up_precision",0) - 0.5)
    improvement = (p.get("dir_accuracy",0) - b.get("dir_accuracy",0) >= 1.0 or
                   p.get("up_precision",0) - b.get("up_precision",0) >= 1.0)
    passed = no_regress and improvement
    print(f"\n  Pass: {'YES' if passed else 'NO'}  (no_regress={no_regress}, improvement={improvement})")

    Path("docs").mkdir(exist_ok=True)
    with open("docs/c2_ptt_result.json","w") as f:
        json.dump({"results":per_stock,"aggregate":agg_summary,
                   "passed":passed,"pass_criterion":"no_regress AND ≥1.0pp on shorter 180d window"}, f, indent=2)


if __name__ == "__main__":
    main()