#!/usr/bin/env python3
"""C6 validation: HMM regime detection as feature + conditional confidence."""
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

sys.path.insert(0, str(ROOT / "scripts"))
from improvement_harness import (
    run_comparison, BASELINE_FEATURES, fetch_df, build_triple_barrier_labels,
    walk_forward, RF_PARAMS, MIN_TRAIN, STEP, LABEL_HORIZON, compute_metrics,
)
from models.predictor import _build_features, FEATURE_COLUMNS

STOCKS = ["2330", "0050", "2317", "2454", "2881"]


def add_hmm_regime(df: pd.DataFrame) -> pd.DataFrame:
    """
    Fit a 3-state GaussianHMM on TAIEX 5d return + 20d vol proxy.
    Add columns: hmm_regime (0/1/2), hmm_regime_bull (binary), hmm_regime_bear (binary).

    If TAIEX columns are missing or hmmlearn fails, add zeros (graceful fallback).
    """
    try:
        from hmmlearn import hmm

        obs_cols = []
        if "taiex_return_5d" in df.columns:
            obs_cols.append("taiex_return_5d")
        if "volatility_20d" in df.columns:
            obs_cols.append("volatility_20d")
        if "taiex_ma20_ratio" in df.columns:
            obs_cols.append("taiex_ma20_ratio")

        if len(obs_cols) < 2:
            raise ValueError("not enough observation columns")

        X_obs = df[obs_cols].fillna(0).values.astype(float)

        model = hmm.GaussianHMM(
            n_components=3, covariance_type="diag",
            n_iter=100, random_state=42,
        )
        model.fit(X_obs)
        regimes = model.predict(X_obs)

        # Label states by mean taiex_return_5d: highest = bull (2), lowest = bear (0)
        means = [X_obs[regimes == s, 0].mean() if (regimes == s).sum() > 0 else 0.0
                 for s in range(3)]
        order = np.argsort(means)
        remap = {order[0]: 0, order[1]: 1, order[2]: 2}
        labeled = np.array([remap[r] for r in regimes])

        df["hmm_regime"]      = labeled.astype(float)
        df["hmm_regime_bull"] = (labeled == 2).astype(float)
        df["hmm_regime_bear"] = (labeled == 0).astype(float)

    except Exception as e:
        df["hmm_regime"]      = 0.0
        df["hmm_regime_bull"] = 0.0
        df["hmm_regime_bear"] = 0.0

    return df


def main():
    BASELINE  = BASELINE_FEATURES
    HMM_FEATS = BASELINE + ["hmm_regime", "hmm_regime_bull", "hmm_regime_bear"]

    per_stock = {}
    agg = {"baseline": [], "hmm": []}

    hdr = f"{'Stock':>6}  {'Set':>10}  {'Acc%':>5}  {'Dir%':>5}  {'↑Prec%':>7}  {'Signals':>7}"
    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

        df = add_hmm_regime(df)
        feat  = _build_features(df)

        for col in ["hmm_regime", "hmm_regime_bull", "hmm_regime_bear"]:
            if col in df.columns:
                if len(df) == len(feat):
                    feat[col] = df[col].values
                else:
                    feat[col] = 0.0

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

        per_stock[stock_no] = {}
        for name, cols in [("baseline", BASELINE), ("hmm", HMM_FEATS)]:
            r = walk_forward(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}  {r.get('n_signals',0):>7}")

        if per_stock[stock_no].get("baseline") and per_stock[stock_no].get("hmm"):
            rb, rh = per_stock[stock_no]["baseline"], per_stock[stock_no]["hmm"]
            print(f"{'':>6}  {'Δ':>10}  {'':>5}  {rh['dir_accuracy']-rb['dir_accuracy']:>+5.1f}  "
                  f"{rh['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", "hmm"):
        rows = agg[name]
        s = {k: mf(rows, k) for k in ("accuracy", "dir_accuracy", "up_precision", "dn_precision", "n_signals")}
        agg_summary[name] = s
        print(f"  {name:>10}: acc={s['accuracy']}%  dir={s['dir_accuracy']}%  ↑prec={s['up_precision']}%")

    h = agg_summary.get("hmm", {})
    b = agg_summary.get("baseline", {})
    no_regress  = h.get("dir_accuracy", 0) >= b.get("dir_accuracy", 0) - 0.5 and \
                  h.get("up_precision", 0) >= b.get("up_precision", 0) - 0.5
    improvement = (h.get("dir_accuracy", 0) - b.get("dir_accuracy", 0) >= 1.5 or
                   h.get("up_precision", 0) - b.get("up_precision", 0) >= 1.5)
    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/c6_hmm_result.json", "w") as f:
        json.dump({"results": per_stock, "aggregate": agg_summary,
                   "passed": passed, "pass_criterion": "no_regress AND >=1.5pp improvement"}, f, indent=2)


if __name__ == "__main__":
    main()