#!/usr/bin/env python3
"""C19: Adaptive refined labeling.

Hypothesis: In low-vol regimes (20d vol < median), standard 1.0σ barriers are
too wide relative to signal, generating noisy UP labels. Tightening to 0.7σ
filters true signals from noise, improving up_precision.

Label scheme:
  low-vol  (vol < median): pt_sl = LOW_RATIO  (0.7 — tighter)
  high-vol (vol ≥ median): pt_sl = HIGH_RATIO (1.0 — unchanged baseline)

Same CURRENT_FEATURES, only labels change.
"""
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 scripts.improvement_harness import (
    fetch_df, walk_forward, compute_metrics,
    CURRENT_FEATURES, DEFAULT_STOCKS, EXTENDED_STOCKS,
    PASS_DIR_ACC, PASS_UP_PREC,
    LABEL_HORIZON, VOL_LOOKBACK,
)
from models.predictor import _build_features

# Regime ratios to test: (low_ratio, high_ratio)
VARIANTS = [
    (0.7, 1.0),   # primary hypothesis: tighten low-vol only
    (0.7, 1.3),   # tighten low, widen high
    (0.6, 1.0),   # more aggressive tightening
]


def build_adaptive_labels(
    close: np.ndarray,
    low_ratio: float = 0.7,
    high_ratio: float = 1.0,
) -> np.ndarray:
    n = len(close)
    log_ret = np.diff(np.log(close + 1e-9))

    vols = []
    for i in range(n):
        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
        vols.append(max(vol, 0.001))

    vol_threshold = float(np.median(vols))

    labels = np.full(n, np.nan)
    for i in range(n - LABEL_HORIZON):
        vol = vols[i]
        pt_sl = low_ratio if vol < vol_threshold else high_ratio
        upper = close[i] * (1.0 + vol * pt_sl)
        lower = close[i] * (1.0 - vol * pt_sl)
        label = 0
        for j in range(1, LABEL_HORIZON + 1):
            c = close[i + j]
            if c >= upper:
                label = 1; break
            if c <= lower:
                label = -1; break
        labels[i] = label
    return labels


def build_baseline_labels(close: np.ndarray) -> np.ndarray:
    """Replicate harness build_triple_barrier_labels for fair comparison."""
    n = len(close)
    log_ret = np.diff(np.log(close + 1e-9))
    labels = np.full(n, np.nan)
    for i in range(n - LABEL_HORIZON):
        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 * 1.0)
        lower = close[i] * (1.0 - vol * 1.0)
        label = 0
        for j in range(1, LABEL_HORIZON + 1):
            c = close[i + j]
            if c >= upper: label = 1; break
            if c <= lower: label = -1; break
        labels[i] = label
    return labels


def label_stats(labels: np.ndarray) -> str:
    valid = labels[~np.isnan(labels)].astype(int)
    n = len(valid)
    if n == 0:
        return "empty"
    up = (valid == 1).sum()
    dn = (valid == -1).sum()
    ho = (valid == 0).sum()
    return f"UP={up/n:.1%} DN={dn/n:.1%} HOLD={ho/n:.1%}"


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--extended", action="store_true")
    args = parser.parse_args()

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

    print(f"\n=== C19: Adaptive refined labeling [{tag}] ===")
    print(f"  Variants: {VARIANTS}\n")

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

    # Run all variants
    variant_results = {v: [] for v in VARIANTS}
    base_results = []
    per_stock = {}

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

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

        base_labels = build_baseline_labels(close)
        m_base = walk_forward(feat, base_labels, CURRENT_FEATURES)
        base_results.append(m_base)

        per_stock[stock_no] = {"baseline": m_base, "variants": {}}
        variant_lines = []
        for (lr, hr) in VARIANTS:
            adap_labels = build_adaptive_labels(close, lr, hr)
            m = walk_forward(feat, adap_labels, CURRENT_FEATURES)
            variant_results[(lr, hr)].append(m)
            per_stock[stock_no]["variants"][f"{lr},{hr}"] = m
            d = m.get("dir_accuracy", float("nan"))
            u = m.get("up_precision", float("nan"))
            variant_lines.append(f"({lr},{hr}) dir={d}% ↑prec={u}%")

        b_dir = m_base.get("dir_accuracy", float("nan"))
        b_up  = m_base.get("up_precision", float("nan"))
        print(f"base dir={b_dir}% ↑prec={b_up}%  →  " + "  |  ".join(variant_lines))

    print()
    base_agg = {
        "dir_accuracy": _avg(base_results, "dir_accuracy"),
        "up_precision": _avg(base_results, "up_precision"),
    }
    print(f"  Baseline avg: dir={base_agg['dir_accuracy']}%  ↑prec={base_agg['up_precision']}%")

    best_variant = None
    best_agg     = None
    best_score   = -999

    for (lr, hr), results in variant_results.items():
        agg = {
            "dir_accuracy": _avg(results, "dir_accuracy"),
            "up_precision": _avg(results, "up_precision"),
        }
        passed = agg["dir_accuracy"] >= PASS_DIR_ACC and agg["up_precision"] >= PASS_UP_PREC
        flag   = "PASS ✓" if passed else "FAIL ✗"
        print(f"  ({lr},{hr}) avg:   dir={agg['dir_accuracy']}%  ↑prec={agg['up_precision']}%  {flag}")
        # Pick best by sum of metrics, only if dir passes gate
        score = agg["up_precision"] if agg["dir_accuracy"] >= PASS_DIR_ACC else agg["up_precision"] - 10
        if score > best_score:
            best_score   = score
            best_variant = (lr, hr)
            best_agg     = agg

    passed = best_agg["dir_accuracy"] >= PASS_DIR_ACC and best_agg["up_precision"] >= PASS_UP_PREC
    print(f"\n  Best variant: low_ratio={best_variant[0]}, high_ratio={best_variant[1]}")
    print(f"  Gate (dir≥{PASS_DIR_ACC}% AND ↑prec≥{PASS_UP_PREC}%): {'PASS ✓' if passed else 'FAIL ✗'}")

    result = {
        "experiment":    "C19",
        "description":   "Adaptive triple-barrier labeling (regime-conditional pt_sl)",
        "variants":      [{"low_ratio": lr, "high_ratio": hr} for (lr, hr) in VARIANTS],
        "best_variant":  {"low_ratio": best_variant[0], "high_ratio": best_variant[1]},
        "stocks":        stocks,
        "aggregate":     {"baseline": base_agg, "best_c19": best_agg},
        "all_variants":  {
            f"{lr},{hr}": {
                "dir_accuracy": _avg(variant_results[(lr, hr)], "dir_accuracy"),
                "up_precision": _avg(variant_results[(lr, hr)], "up_precision"),
            }
            for (lr, hr) in VARIANTS
        },
        "passed":        passed,
        "pass_gate":     {"dir_accuracy": PASS_DIR_ACC, "up_precision": PASS_UP_PREC},
        "per_stock":     per_stock,
    }

    out = ROOT / f"docs/c19_result{suffix}.json"
    out.parent.mkdir(exist_ok=True)
    out.write_text(json.dumps(result, indent=2))
    print(f"\n  Saved: {out}")
    return 0 if passed else 1


if __name__ == "__main__":
    sys.exit(main())