#!/usr/bin/env python3
"""C18: rolling confidence gate for higher overall accuracy.

The gate is calibrated only on past data inside each walk-forward window:
train -> calibration -> next test slice. Low-confidence directional predictions
are converted to HOLD. This targets overall accuracy while reporting the signal
count trade-off explicitly.
"""

from __future__ import annotations

import argparse
import json
import os
import sys
from collections import Counter
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
from sklearn.ensemble import RandomForestClassifier

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


CLASSES = np.array([-1, 0, 1])
THRESHOLDS = tuple(round(v, 2) for v in np.arange(0.34, 0.72, 0.02))
CAL_WINDOW = 84
MIN_SIGNAL_RATE = 0.35
PASS_ACCURACY_DELTA = 5.0


def _align_proba(model, X: np.ndarray) -> np.ndarray:
    prob = model.predict_proba(X)
    out = np.zeros((len(X), len(CLASSES)))
    model_classes = list(model.classes_)
    for idx, klass in enumerate(CLASSES):
        if klass in model_classes:
            out[:, idx] = prob[:, model_classes.index(klass)]
    return out


def _predict_with_gate(prob: np.ndarray, threshold: float) -> np.ndarray:
    raw = CLASSES[np.argmax(prob, axis=1)]
    confidence = prob.max(axis=1)
    return np.where(confidence >= threshold, raw, 0)


def choose_threshold(
    prob: np.ndarray,
    y_true: np.ndarray,
    *,
    thresholds: tuple[float, ...] = THRESHOLDS,
    min_signal_rate: float = MIN_SIGNAL_RATE,
) -> tuple[float, dict]:
    """Choose a confidence threshold from past calibration data only."""
    candidates = []
    fallback = []
    for threshold in thresholds:
        pred = _predict_with_gate(prob, threshold)
        signal_rate = float(np.mean(pred != 0))
        accuracy = float(np.mean(pred == y_true))
        dir_mask = pred != 0
        dir_accuracy = float(np.mean(pred[dir_mask] == y_true[dir_mask])) if dir_mask.any() else 0.0
        row = {
            "threshold": threshold,
            "accuracy": accuracy,
            "dir_accuracy": dir_accuracy,
            "signal_rate": signal_rate,
        }
        fallback.append(row)
        if signal_rate >= min_signal_rate:
            candidates.append(row)

    pool = candidates or fallback
    best = max(pool, key=lambda row: (row["accuracy"], row["dir_accuracy"], row["signal_rate"]))
    return float(best["threshold"]), {
        "accuracy": round(best["accuracy"] * 100, 1),
        "dir_accuracy": round(best["dir_accuracy"] * 100, 1),
        "signal_rate": round(best["signal_rate"], 3),
    }


def walk_forward_confidence_gate(feat_df, labels: np.ndarray, cols: list[str]) -> tuple[dict, dict]:
    avail = [col for col in cols if col in feat_df.columns]
    X_all = feat_df[avail].fillna(0).values
    y_true_all, y_pred_base_all, y_pred_gate_all = [], [], []
    threshold_counts: Counter[float] = Counter()
    cal_rows = []

    cutoff = MIN_TRAIN
    n = len(feat_df)
    while cutoff + STEP + LABEL_HORIZON <= n:
        # Embargo the horizon immediately before the test slice so training
        # labels never depend on prices inside the next test window.
        train_end = cutoff - LABEL_HORIZON
        if train_end < MIN_TRAIN:
            cutoff += STEP
            continue

        y_train_full = labels[:train_end]
        valid_idx = np.where(~np.isnan(y_train_full))[0]
        if len(valid_idx) < MIN_TRAIN or len(np.unique(y_train_full[valid_idx])) < 2:
            cutoff += STEP
            continue

        cal_start = max(0, len(valid_idx) - CAL_WINDOW)
        train_idx = valid_idx[:cal_start]
        cal_idx = valid_idx[cal_start:]
        if len(train_idx) < 120 or len(cal_idx) < 30 or len(np.unique(y_train_full[train_idx])) < 2:
            cutoff += STEP
            continue

        cal_model = RandomForestClassifier(**RF_PARAMS)
        cal_model.fit(X_all[train_idx], y_train_full[train_idx].astype(int))
        cal_prob = _align_proba(cal_model, X_all[cal_idx])
        threshold, cal_metric = choose_threshold(cal_prob, y_train_full[cal_idx].astype(int))
        threshold_counts[threshold] += 1
        cal_rows.append(cal_metric)

        full_model = RandomForestClassifier(**RF_PARAMS)
        full_model.fit(X_all[valid_idx], y_train_full[valid_idx].astype(int))

        test_end = min(cutoff + STEP, n - LABEL_HORIZON)
        y_test = labels[cutoff:test_end]
        valid_test = ~np.isnan(y_test)
        if valid_test.sum() == 0:
            cutoff += STEP
            continue

        test_prob = _align_proba(full_model, X_all[cutoff:test_end][valid_test])
        y_pred_base = CLASSES[np.argmax(test_prob, axis=1)]
        y_pred_gate = _predict_with_gate(test_prob, threshold)

        y_true_all.extend(y_test[valid_test].tolist())
        y_pred_base_all.extend(y_pred_base.tolist())
        y_pred_gate_all.extend(y_pred_gate.tolist())
        cutoff += STEP

    if not y_true_all:
        return {}, {}

    y_true = np.array(y_true_all)
    baseline = compute_metrics(y_true, np.array(y_pred_base_all))
    gated = compute_metrics(y_true, np.array(y_pred_gate_all))
    gated["threshold_counts"] = {str(k): v for k, v in sorted(threshold_counts.items())}
    if cal_rows:
        gated["avg_cal_accuracy"] = round(float(np.mean([r["accuracy"] for r in cal_rows])), 1)
        gated["avg_cal_signal_rate"] = round(float(np.mean([r["signal_rate"] for r in cal_rows])), 3)
    return baseline, gated


def _mean(rows: list[dict], field: str) -> float:
    vals = [
        row[field]
        for row in rows
        if isinstance(row.get(field), (int, float)) and not np.isnan(row.get(field, float("nan")))
    ]
    return round(sum(vals) / len(vals), 1) if vals else float("nan")


def _aggregate(rows: list[dict]) -> dict:
    return {
        "accuracy": _mean(rows, "accuracy"),
        "dir_accuracy": _mean(rows, "dir_accuracy"),
        "up_precision": _mean(rows, "up_precision"),
        "dn_precision": _mean(rows, "dn_precision"),
        "n_signals": _mean(rows, "n_signals"),
        "n_predictions": _mean(rows, "n_predictions"),
    }


def _json_default(value):
    if isinstance(value, np.integer):
        return int(value)
    if isinstance(value, np.floating):
        return float(value)
    if isinstance(value, np.bool_):
        return bool(value)
    raise TypeError(f"Object of type {type(value).__name__} is not JSON serializable")


def run(stocks: list[str], output_path: Path) -> dict:
    results = {}
    agg_base, agg_gate = [], []

    hdr = f"{'Stock':>6}  {'Model':>12}  {'Acc%':>5}  {'Dir%':>5}  {'Up%':>6}  {'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

        feat = _build_features(df)
        close = (df.set_index("date")["close"] if "date" in df.columns else df["close"]).values
        labels = build_triple_barrier_labels(close)
        baseline, gated = walk_forward_confidence_gate(feat, labels, BASELINE_FEATURES)
        if not baseline:
            continue

        results[stock_no] = {"baseline": baseline, "confidence_gate": gated}
        agg_base.append(baseline)
        agg_gate.append(gated)

        for name, row in (("baseline", baseline), ("conf_gate", gated)):
            prefix = f"{stock_no:>6}" if name == "baseline" else f"{'':>6}"
            print(
                f"{prefix}  {name:>12}  {row['accuracy']:>5.1f}  {row['dir_accuracy']:>5.1f}  "
                f"{row['up_precision']:>6.1f}  {row['n_signals']:>7}"
            )
        print(f"{'':>6}  {'delta':>12}  {gated['accuracy'] - baseline['accuracy']:>+5.1f}")

    aggregate = {"baseline": _aggregate(agg_base), "confidence_gate": _aggregate(agg_gate)}
    acc_delta = aggregate["confidence_gate"]["accuracy"] - aggregate["baseline"]["accuracy"]
    signal_delta = aggregate["confidence_gate"]["n_signals"] - aggregate["baseline"]["n_signals"]
    passed = bool(acc_delta >= PASS_ACCURACY_DELTA)

    print("\n=== AGGREGATE ===")
    for name, row in aggregate.items():
        print(
            f"  {name:>15}: acc={row['accuracy']}%  dir={row['dir_accuracy']}%  "
            f"up={row['up_precision']}%  signals={row['n_signals']}"
        )
    print(f"\n  Delta accuracy: {acc_delta:+.1f}pp")
    print(f"  Delta signals:  {signal_delta:+.1f} per stock")
    print(f"  Pass +{PASS_ACCURACY_DELTA:.1f}pp overall accuracy: {'YES' if passed else 'NO'}")

    result = {
        "experiment": "C18_confidence_gate",
        "description": "Rolling no-lookahead confidence gate: low-confidence directional predictions become HOLD.",
        "stocks": stocks,
        "pass_criterion": {"accuracy_delta_pp": PASS_ACCURACY_DELTA},
        "aggregate": aggregate,
        "accuracy_delta_pp": round(acc_delta, 1),
        "signal_delta_per_stock": round(signal_delta, 1),
        "passed": passed,
        "results": results,
    }
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(json.dumps(result, indent=2, default=_json_default), encoding="utf-8")
    print(f"  Saved -> {output_path}")
    return result


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--extended", action="store_true", help="Use 12-stock extended validation set")
    parser.add_argument("--stocks", help="Comma-separated stock codes")
    parser.add_argument("--full-data", action="store_true", help="Enable slower institutional/margin fetches")
    parser.add_argument("--output", default=str(ROOT / "docs" / "c18_confidence_gate_result.json"))
    args = parser.parse_args()

    if not args.full_data:
        os.environ.setdefault("ENABLE_INSTITUTIONAL_FLOW", "0")
        os.environ.setdefault("ENABLE_MARGIN_FLOW", "0")
        os.environ.setdefault("INSTITUTIONAL_FLOW_DAYS", "0")

    if args.stocks:
        stocks = [code.strip() for code in args.stocks.split(",") if code.strip()]
    else:
        stocks = EXTENDED_STOCKS if args.extended else DEFAULT_STOCKS
    run(stocks, Path(args.output))


if __name__ == "__main__":
    main()