#!/usr/bin/env python3
"""C12: 1-day vs 5-day label horizon comparison."""
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
import pandas as pd
from sklearn.ensemble import RandomForestClassifier

from scripts.improvement_harness import (
    fetch_df, build_triple_barrier_labels, walk_forward,
    compute_metrics, BASELINE_FEATURES, DEFAULT_STOCKS,
    RF_PARAMS, MIN_TRAIN, STEP, LABEL_HORIZON,
)
from models.predictor import _build_features

PASS_DIR = 44.0
PASS_UP  = 54.0


def build_1day_labels(close: np.ndarray, vol_lookback: int = 20, pt_sl_ratio: float = 1.0) -> np.ndarray:
    """Triple barrier with 1-day vertical barrier instead of 5-day."""
    n = len(close)
    log_ret = np.diff(np.log(close + 1e-9))
    labels = np.full(n, np.nan)
    for i in range(n - 1):  # horizon=1
        start_v = max(0, i - vol_lookback)
        window = log_ret[start_v:i]
        vol = float(np.std(window)) if len(window) >= 5 else 0.015
        vol = max(vol, 0.001)
        upper = close[i] * (1.0 + vol * pt_sl_ratio)
        lower = close[i] * (1.0 - vol * pt_sl_ratio)
        c = close[i + 1]
        if c >= upper:
            labels[i] = 1
        elif c <= lower:
            labels[i] = -1
        else:
            labels[i] = 0
    return labels


def walk_forward_1d(feat_df: pd.DataFrame, label_arr: np.ndarray, cols: list) -> dict:
    LABEL_HORIZON_1D = 1
    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_1D <= 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_1D)
        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():
    hdr = f"{'Stock':>6}  {'5day_dir%':>9}  {'1day_dir%':>9}  {'Δdir':>5}  {'5day_↑prec%':>11}  {'1day_↑prec%':>11}  {'Δprec':>6}  {'1day_sigs':>9}"
    print(f"\n{hdr}\n{'-'*len(hdr)}")

    per_stock = {}
    rows_5d, rows_1d = [], []

    for stock_no in DEFAULT_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.set_index("date")["close"] if "date" in df.columns else df["close"]).values

        labels_5d = build_triple_barrier_labels(close)
        r5 = walk_forward(feat, labels_5d, BASELINE_FEATURES)

        labels_1d = build_1day_labels(close)
        r1 = walk_forward_1d(feat, labels_1d, BASELINE_FEATURES)

        per_stock[stock_no] = {"5day": r5, "1day": r1}
        if r5:
            rows_5d.append(r5)
        if r1:
            rows_1d.append(r1)

        d5 = r5.get("dir_accuracy", float("nan"))
        d1 = r1.get("dir_accuracy", float("nan"))
        p5 = r5.get("up_precision", float("nan"))
        p1 = r1.get("up_precision", float("nan"))
        dd = d1 - d5 if not (np.isnan(d1) or np.isnan(d5)) else float("nan")
        dp = p1 - p5 if not (np.isnan(p1) or np.isnan(p5)) else float("nan")
        sigs = r1.get("n_signals", 0)
        print(f"{stock_no:>6}  {d5:>9.1f}  {d1:>9.1f}  {dd:>+5.1f}  {p5:>11.1f}  {p1:>11.1f}  {dp:>+6.1f}  {sigs:>9}")

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

    agg5 = {k: _mean(rows_5d, k) for k in ("accuracy", "dir_accuracy", "up_precision", "dn_precision", "n_signals")}
    agg1 = {k: _mean(rows_1d, k) for k in ("accuracy", "dir_accuracy", "up_precision", "dn_precision", "n_signals")}

    dd_agg = agg1["dir_accuracy"] - agg5["dir_accuracy"]
    dp_agg = agg1["up_precision"] - agg5["up_precision"]

    print(f"\n{'=== AGGREGATE ==='}")
    print(f"  5day: dir={agg5['dir_accuracy']}%  ↑prec={agg5['up_precision']}%  signals={agg5['n_signals']}")
    print(f"  1day: dir={agg1['dir_accuracy']}%  ↑prec={agg1['up_precision']}%  signals={agg1['n_signals']}")
    print(f"  Δ   : dir={dd_agg:+.1f}pp  ↑prec={dp_agg:+.1f}pp")

    passed = bool((agg1["dir_accuracy"] >= PASS_DIR) and (agg1["up_precision"] >= PASS_UP))
    status = "PASSED" if passed else "FAILED"
    print(f"\n  Pass criterion: 1day dir ≥ {PASS_DIR} AND ↑prec ≥ {PASS_UP}")
    print(f"  C12 {status}")

    result = {
        "results": per_stock,
        "aggregate": {"5day": agg5, "1day": agg1},
        "passed": passed,
        "pass_criterion": f"1day dir_accuracy >= {PASS_DIR} AND up_precision >= {PASS_UP}",
    }

    out = ROOT / "docs" / "c12_1day_result.json"
    out.parent.mkdir(exist_ok=True)
    with open(out, "w") as f:
        json.dump(result, f, indent=2)
    print(f"\n  Written: {out}")

    print(f"\nC12 {status} — 1day dir: {agg1['dir_accuracy']}%, ↑prec: {agg1['up_precision']}%, signals/stock: {agg1['n_signals']}")


if __name__ == "__main__":
    main()