#!/usr/bin/env python3
"""C1 validation: 40f baseline vs 40f + fracdiff features."""
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 sklearn.ensemble import RandomForestClassifier

from models.predictor import FEATURE_COLUMNS

BASELINE = FEATURE_COLUMNS  # 40 features
FRACDIFF = FEATURE_COLUMNS + ["fracdiff_close_35", "fracdiff_close_40"]

STOCKS = ["2330", "0050", "2317", "2454", "2881"]
LABEL_HORIZON = 5
LABEL_THRESH = 0.02
MIN_TRAIN = 240
STEP = 21
RF = dict(n_estimators=200, max_depth=6, min_samples_leaf=10,
          class_weight="balanced", random_state=42, n_jobs=-1)

def fetch_df(stock_no):
    from services.predictor_service import _fetch_with_cache
    from indicators.technical import add_all_indicators, add_cross_asset_tw
    from data.institutional_flow import add_institutional_flow
    from data.margin_flow import add_margin_flow
    from data.fetcher import fetch_cross_asset_tw
    df = _fetch_with_cache(stock_no, months=24)
    if df is None or df.empty: return None
    df = add_all_indicators(df)
    df = add_institutional_flow(df, stock_no)
    df = add_margin_flow(df, stock_no)
    start, end = str(df["date"].min()), str(df["date"].max())
    result = fetch_cross_asset_tw(start, end)
    taiex, usdtwd = result[0], result[1]
    sox = result[2] if len(result) > 2 else None
    tnx = result[3] if len(result) > 3 else None
    df = add_cross_asset_tw(df, taiex, usdtwd, sox_close=sox, tnx_close=tnx)
    return df

def make_labels(close):
    fwd = close.shift(-LABEL_HORIZON)
    ret = (fwd - close) / close
    return np.where(ret > LABEL_THRESH, 1, np.where(ret < -LABEL_THRESH, -1, 0))

def walk_forward(feat, cols):
    avail = [c for c in cols if c in feat.columns]
    labels = feat["_label"].values
    X_all = feat[avail].fillna(0).values
    n = len(feat)
    y_true_all, y_pred_all = [], []
    cutoff = MIN_TRAIN
    while cutoff + STEP + LABEL_HORIZON <= n:
        X_tr, y_tr = X_all[:cutoff], labels[:cutoff]
        test_end = min(cutoff + STEP, n - LABEL_HORIZON)
        X_te, y_te = X_all[cutoff:test_end], labels[cutoff:test_end]
        if len(np.unique(y_tr)) < 2 or len(X_te) == 0:
            cutoff += STEP; continue
        clf = RandomForestClassifier(**RF)
        clf.fit(X_tr, y_tr)
        y_pred = clf.predict(X_te)
        y_true_all.extend(y_te.tolist())
        y_pred_all.extend(y_pred.tolist())
        cutoff += STEP
    if not y_true_all: return {}
    y_true = np.array(y_true_all)
    y_pred = np.array(y_pred_all)
    dir_mask = y_pred != 0
    acc = (y_true == y_pred).mean() * 100
    dir_acc = (y_true[dir_mask] == y_pred[dir_mask]).mean() * 100 if dir_mask.sum() > 0 else float("nan")
    up_mask = y_pred == 1
    up_prec = (y_true[up_mask] == 1).mean() * 100 if up_mask.sum() > 0 else float("nan")
    return {"accuracy": round(acc, 1), "dir_accuracy": round(dir_acc, 1),
            "up_precision": round(up_prec, 1), "n_predictions": len(y_true),
            "n_features": len(avail)}

def main():
    from models.predictor import _build_features
    results = {}
    hdr = f"{'Stock':>6}  {'Set':>10}  {'Acc%':>5}  {'Dir%':>5}  {'UpPrec%':>7}  {'Feat':>4}"
    print(f"\n{hdr}\n{'-'*len(hdr)}")
    agg = {"baseline": [], "fracdiff": []}
    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: continue
        feat = _build_features(df)
        # Pass fracdiff columns through from df (not built by _build_features)
        for col in ("fracdiff_close_35", "fracdiff_close_40"):
            if col in df.columns:
                feat[col] = df[col].values
        close = df.set_index("date")["close"] if "date" in df.columns else df["close"]
        feat["_label"] = make_labels(close)
        results[stock_no] = {}
        for name, cols in [("baseline", BASELINE), ("fracdiff", FRACDIFF)]:
            r = walk_forward(feat, cols)
            results[stock_no][name] = r
            if r:
                print(f"{stock_no:>6}  {name:>10}  {r['accuracy']:>5.1f}  {r['dir_accuracy']:>5.1f}  {r['up_precision']:>7.1f}  {r['n_features']:>4}")
                agg[name].append(r)
        if results[stock_no].get("baseline") and results[stock_no].get("fracdiff"):
            rb, rf = results[stock_no]["baseline"], results[stock_no]["fracdiff"]
            dd = rf["dir_accuracy"] - rb["dir_accuracy"]
            print(f"{'':>6}  {'delta':>10}  {'':>5}  {dd:>+5.1f}")
        print()

    def mean_field(rows, f):
        vals = [r[f] for r in rows if not (isinstance(r.get(f), float) and 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", "fracdiff"):
        rows = agg[name]
        s = {k: mean_field(rows, k) for k in ("accuracy", "dir_accuracy", "up_precision")}
        agg_summary[name] = s
        print(f"  {name:>10}: acc={s['accuracy']}%  dir={s['dir_accuracy']}%  up_prec={s['up_precision']}%")

    winner = "fracdiff" if agg_summary["fracdiff"]["dir_accuracy"] > agg_summary["baseline"]["dir_accuracy"] else "baseline"
    passed = agg_summary["fracdiff"]["dir_accuracy"] >= 31.5
    print(f"\n  Winner: {winner}  |  passed={passed}")

    out = {"results": results, "aggregate": agg_summary, "winner": winner,
           "pass_threshold": 31.5, "passed": passed}
    Path("docs").mkdir(exist_ok=True)
    with open("docs/c1_fracdiff_result.json", "w") as f:
        json.dump(out, f, indent=2)
    return out

if __name__ == "__main__":
    main()