#!/usr/bin/env python3
"""Confirm train-learned factor side masks on held-out stocks.

The policy keeps only factor sides whose train-stock label lift is positive.
It is fixed before held-out replay and does not modify production defaults.
"""

from __future__ import annotations

import argparse
import html
import json
import sys
from pathlib import Path
from typing import Any

import numpy as np
import pandas as pd

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))

from scripts.confirm_fixed_sell_pattern_gate import _decision  # noqa: E402
from scripts.evaluate_sell_pattern_regime_gates import _prediction_rates  # noqa: E402
from scripts.optimize_factor_weights import compute_event_metrics, metric_deltas  # noqa: E402
from scripts.search_multi_factor_weight_config import MultiFactorConfig, _now_iso, apply_multi_factor_config  # noqa: E402


def _root_path(value: str | Path) -> Path:
    path = Path(value)
    return path if path.is_absolute() else ROOT / path


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


def _select_stocks(events: pd.DataFrame, *, offset: int, limit: int) -> list[str]:
    stocks = sorted(events["stock"].astype(str).unique().tolist())
    if offset:
        stocks = stocks[offset:]
    if limit and limit > 0:
        stocks = stocks[:limit]
    return stocks


def _split_float_csv(value: str) -> list[float]:
    return [float(item.strip()) for item in value.split(",") if item.strip()]


def _base_label_rates(events: pd.DataFrame) -> dict[str, float]:
    y = pd.to_numeric(events["y_true"], errors="coerce").fillna(0).to_numpy(dtype=int)
    return {
        "buy": float((y == 1).mean() * 100.0) if len(y) else 0.0,
        "sell": float((y == -1).mean() * 100.0) if len(y) else 0.0,
    }


def _side_lift(events: pd.DataFrame, factor: str, side: str, base_rates: dict[str, float]) -> dict[str, Any]:
    values = pd.to_numeric(events[f"factor__{factor}"], errors="coerce").fillna(0.0).to_numpy(dtype=float)
    y = pd.to_numeric(events["y_true"], errors="coerce").fillna(0).to_numpy(dtype=int)
    mask = values > 0.0 if side == "positive" else values < 0.0
    count = int(mask.sum())
    if count == 0:
        return {"count": 0, "precision": None, "lift_pp": None}
    if side == "positive":
        precision = float((y[mask] == 1).mean() * 100.0)
        lift = precision - base_rates["buy"]
    else:
        precision = float((y[mask] == -1).mean() * 100.0)
        lift = precision - base_rates["sell"]
    return {"count": count, "precision": round(precision, 4), "lift_pp": round(lift, 4)}


def _load_config(path: Path) -> MultiFactorConfig:
    payload = json.loads(path.read_text(encoding="utf-8"))
    config = payload.get("selected_config") or payload.get("config") or payload
    return MultiFactorConfig(**config)


def learn_policy(
    events: pd.DataFrame,
    config: MultiFactorConfig,
    *,
    min_lift_pp: float,
    min_count: int,
) -> dict[str, Any]:
    base_rates = _base_label_rates(events)
    factors = [factor for factor, weight in config.weights.items() if abs(float(weight)) > 0.0 and f"factor__{factor}" in events.columns]
    keep_positive: list[str] = []
    keep_negative: list[str] = []
    stats: dict[str, Any] = {}
    for factor in factors:
        positive = _side_lift(events, factor, "positive", base_rates)
        negative = _side_lift(events, factor, "negative", base_rates)
        if positive["count"] >= min_count and positive["lift_pp"] is not None and positive["lift_pp"] >= min_lift_pp:
            keep_positive.append(factor)
        if negative["count"] >= min_count and negative["lift_pp"] is not None and negative["lift_pp"] >= min_lift_pp:
            keep_negative.append(factor)
        stats[factor] = {"positive": positive, "negative": negative}
    return {
        "min_lift_pp": float(min_lift_pp),
        "min_count": int(min_count),
        "keep_positive": sorted(keep_positive),
        "keep_negative": sorted(keep_negative),
        "factor_stats": stats,
    }


def apply_policy(events: pd.DataFrame, config: MultiFactorConfig, policy: dict[str, Any]) -> pd.DataFrame:
    output = events.copy()
    keep_positive = set(policy["keep_positive"])
    keep_negative = set(policy["keep_negative"])
    for factor in config.weights:
        column = f"factor__{factor}"
        if column not in output:
            continue
        values = pd.to_numeric(output[column], errors="coerce").fillna(0.0)
        if factor not in keep_positive:
            values = values.where(values <= 0.0, 0.0)
        if factor not in keep_negative:
            values = values.where(values >= 0.0, 0.0)
        output[column] = values
    return output


def _evaluate(events: pd.DataFrame, config: MultiFactorConfig, policy: dict[str, Any]) -> dict[str, Any]:
    candidate_events = apply_policy(events, config, policy)
    baseline_metrics = compute_event_metrics(events, events["base_pred"].to_numpy(dtype=int))
    candidate_pred = apply_multi_factor_config(candidate_events, config)
    candidate_metrics = compute_event_metrics(events, candidate_pred)
    deltas = metric_deltas(baseline_metrics, candidate_metrics)
    return {
        "baseline": {"metrics": baseline_metrics, "prediction_rates": _prediction_rates(baseline_metrics)},
        "candidate": {"metrics": candidate_metrics, "prediction_rates": _prediction_rates(candidate_metrics)},
        "deltas": deltas,
    }


def _train_pass(deltas: dict[str, Any], args: argparse.Namespace) -> bool:
    return (
        deltas["accuracy_delta_pp"] >= args.min_accuracy_delta_pp
        and deltas["direction_accuracy_delta_pp"] >= args.min_direction_accuracy_delta_pp
        and deltas["buy_precision_delta_pp"] >= args.min_buy_precision_delta_pp
        and deltas["sell_precision_delta_pp"] >= args.min_sell_precision_delta_pp
    )


def _rank_train(row: dict[str, Any]) -> tuple[float, ...]:
    deltas = row["train"]["deltas"]
    return (
        float(row["train_passed"]),
        float(deltas["sell_precision_delta_pp"]),
        float(deltas["buy_precision_delta_pp"]),
        float(deltas["direction_accuracy_delta_pp"]),
        float(deltas["accuracy_delta_pp"]),
        -float(row["policy"]["min_lift_pp"]),
    )


def evaluate(args: argparse.Namespace) -> dict[str, Any]:
    events = pd.read_pickle(_root_path(args.events_cache))
    events["stock"] = events["stock"].astype(str)
    config = _load_config(_root_path(args.config))
    available_weights = {
        factor: float(weight)
        for factor, weight in config.weights.items()
        if f"factor__{factor}" in events.columns and abs(float(weight)) > 0.0
    }
    if not available_weights:
        raise ValueError("No non-zero config weights are available in the events cache.")
    config = MultiFactorConfig(weights=available_weights, hold_bias=float(config.hold_bias))

    train_stocks = _select_stocks(events, offset=args.train_stock_offset, limit=args.train_limit)
    heldout_stocks = _select_stocks(events, offset=args.heldout_stock_offset, limit=args.heldout_limit)
    train = events[events["stock"].isin(train_stocks)].reset_index(drop=True)
    heldout = events[events["stock"].isin(heldout_stocks)].reset_index(drop=True)

    rows = []
    for min_lift in _split_float_csv(args.min_lift_grid):
        policy = learn_policy(train, config, min_lift_pp=min_lift, min_count=args.min_count)
        train_eval = _evaluate(train, config, policy)
        rows.append(
            {
                "policy": policy,
                "train": train_eval,
                "train_passed": _train_pass(train_eval["deltas"], args),
            }
        )
    selected = max(rows, key=_rank_train)
    heldout_eval = _evaluate(heldout, config, selected["policy"])
    heldout_decision = _decision(
        heldout_eval["deltas"],
        min_signal_ratio=args.min_signal_ratio,
        base_metrics=heldout_eval["baseline"]["metrics"],
        candidate_metrics=heldout_eval["candidate"]["metrics"],
    )
    payload = {
        "experiment": "train_learned_factor_side_policy_confirmation",
        "generated_at": _now_iso(),
        "research_only": True,
        "production_defaults_modified": False,
        "events_cache": str(_root_path(args.events_cache)),
        "config": str(_root_path(args.config)),
        "validation_design": {
            "mode": "learn_side_mask_on_train_stocks_confirm_fixed_policy_on_heldout_stocks",
            "train_stock_offset": args.train_stock_offset,
            "train_limit": args.train_limit,
            "heldout_stock_offset": args.heldout_stock_offset,
            "heldout_limit": args.heldout_limit,
            "min_count": args.min_count,
            "min_lift_grid": _split_float_csv(args.min_lift_grid),
        },
        "train_stocks": train_stocks,
        "heldout_stocks": heldout_stocks,
        "policy_rows": sorted(rows, key=_rank_train, reverse=True),
        "selected_policy": selected,
        "heldout_confirmation": {**heldout_eval, "decision": heldout_decision},
    }
    output_json = _root_path(args.output_json)
    output_json.parent.mkdir(parents=True, exist_ok=True)
    output_json.write_text(json.dumps(payload, indent=2, ensure_ascii=False, default=_json_default))
    if args.output_html:
        render_html(payload, _root_path(args.output_html))
    return payload


def _fmt(value: Any, suffix: str = "") -> str:
    if value is None:
        return "-"
    if isinstance(value, float):
        return f"{value:.4f}{suffix}"
    return f"{value}{suffix}"


def _metric_row(label: str, item: dict[str, Any]) -> str:
    metrics = item["metrics"]
    rates = item["prediction_rates"]
    return (
        f"<tr><td>{html.escape(label)}</td><td>{_fmt(metrics['accuracy'], '%')}</td>"
        f"<td>{_fmt(metrics['direction_accuracy'], '%')}</td><td>{_fmt(metrics['buy_precision'], '%')}</td>"
        f"<td>{_fmt(metrics['sell_precision'], '%')}</td><td>{_fmt(rates['buy_rate'], '%')}</td>"
        f"<td>{_fmt(rates['hold_rate'], '%')}</td><td>{_fmt(rates['sell_rate'], '%')}</td></tr>"
    )


def render_html(payload: dict[str, Any], path: Path) -> None:
    selected = payload["selected_policy"]
    heldout = payload["heldout_confirmation"]
    decision_html = '<span class="pass">PASS</span>' if heldout["decision"]["passed"] else '<span class="fail">FAIL</span>'
    policy = selected["policy"]
    row_html = []
    for row in payload["policy_rows"]:
        deltas = row["train"]["deltas"]
        row_html.append(
            "<tr>"
            f"<td>{'PASS' if row['train_passed'] else 'FAIL'}</td>"
            f"<td>{_fmt(row['policy']['min_lift_pp'], 'pp')}</td>"
            f"<td>{len(row['policy']['keep_positive'])}</td><td>{len(row['policy']['keep_negative'])}</td>"
            f"<td>{_fmt(deltas['accuracy_delta_pp'], 'pp')}</td>"
            f"<td>{_fmt(deltas['direction_accuracy_delta_pp'], 'pp')}</td>"
            f"<td>{_fmt(deltas['buy_precision_delta_pp'], 'pp')}</td>"
            f"<td>{_fmt(deltas['sell_precision_delta_pp'], 'pp')}</td></tr>"
        )
    body = f"""<!doctype html>
<html lang="zh-Hant"><head><meta charset="utf-8">
<title>Train Learned Factor Side Policy Confirmation</title>
<style>
body{{font-family:-apple-system,BlinkMacSystemFont,"Segoe UI",sans-serif;margin:32px;color:#1f2933;line-height:1.5}}
h1,h2{{color:#102a43}}table{{border-collapse:collapse;width:100%;margin:16px 0 24px;font-size:13px}}
th,td{{border:1px solid #cbd2d9;padding:8px 10px;text-align:right;vertical-align:top}}th:first-child,td:first-child{{text-align:left}}th{{background:#f0f4f8}}
.pass{{color:#0b7285;font-weight:700}}.fail{{color:#b42318;font-weight:700}}code,pre{{background:#f0f4f8;padding:2px 4px;border-radius:4px}}pre{{white-space:pre-wrap;padding:12px}}
</style></head><body>
<h1>Train Learned Factor Side Policy Confirmation</h1>
<p><strong>Generated:</strong> {html.escape(payload['generated_at'])}</p>
<p><strong>Held-out decision:</strong> {decision_html}</p>
<h2>Selected Policy</h2>
<p><strong>Min lift:</strong> {_fmt(policy['min_lift_pp'], 'pp')} · <strong>Kept positive sides:</strong> {len(policy['keep_positive'])} · <strong>Kept negative sides:</strong> {len(policy['keep_negative'])}</p>
<pre>{html.escape(json.dumps({'keep_positive': policy['keep_positive'], 'keep_negative': policy['keep_negative']}, ensure_ascii=False, indent=2))}</pre>
<h2>Held-Out Confirmation</h2>
<table><thead><tr><th>Scope</th><th>Accuracy</th><th>Direction Acc</th><th>BUY Precision</th><th>SELL Precision</th><th>BUY Rate</th><th>HOLD Rate</th><th>SELL Rate</th></tr></thead><tbody>
{_metric_row('Baseline held-out', heldout['baseline'])}
{_metric_row('Candidate held-out', heldout['candidate'])}
<tr><td>Delta</td><td>{_fmt(heldout['deltas']['accuracy_delta_pp'], 'pp')}</td><td>{_fmt(heldout['deltas']['direction_accuracy_delta_pp'], 'pp')}</td><td>{_fmt(heldout['deltas']['buy_precision_delta_pp'], 'pp')}</td><td>{_fmt(heldout['deltas']['sell_precision_delta_pp'], 'pp')}</td><td></td><td></td><td></td></tr>
</tbody></table>
<h2>Train Policy Sweep</h2>
<table><thead><tr><th>Train Status</th><th>Min Lift</th><th>Keep +</th><th>Keep -</th><th>Acc Delta</th><th>Direction Delta</th><th>BUY Delta</th><th>SELL Delta</th></tr></thead><tbody>
{''.join(row_html)}
</tbody></table>
</body></html>"""
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(body)


def main() -> int:
    parser = argparse.ArgumentParser(description="Confirm train-learned factor side policy on held-out stocks.")
    parser.add_argument("--events-cache", default="docs/validation_runs/factor44_split_problem_v1_buy_factor_interaction_adaptive_macro_1000_events.pkl")
    parser.add_argument("--config", default="config/multi_factor_overlay.json")
    parser.add_argument("--train-stock-offset", type=int, default=0)
    parser.add_argument("--train-limit", type=int, default=1000)
    parser.add_argument("--heldout-stock-offset", type=int, default=1000)
    parser.add_argument("--heldout-limit", type=int, default=0)
    parser.add_argument("--min-count", type=int, default=1000)
    parser.add_argument("--min-lift-grid", default="0,0.5,1,2,3,5")
    parser.add_argument("--min-accuracy-delta-pp", type=float, default=0.0)
    parser.add_argument("--min-buy-precision-delta-pp", type=float, default=0.0)
    parser.add_argument("--min-sell-precision-delta-pp", type=float, default=0.0)
    parser.add_argument("--min-direction-accuracy-delta-pp", type=float, default=0.0)
    parser.add_argument("--min-signal-ratio", type=float, default=0.70)
    parser.add_argument("--output-json", default="docs/validation_runs/train_learned_side_policy_1000train_199heldout_20260612.json")
    parser.add_argument("--output-html", default="docs/validation_runs/train_learned_side_policy_1000train_199heldout_20260612.html")
    args = parser.parse_args()
    payload = evaluate(args)
    heldout = payload["heldout_confirmation"]
    print(f"Selected min lift: {payload['selected_policy']['policy']['min_lift_pp']}")
    print(f"Kept positive sides: {len(payload['selected_policy']['policy']['keep_positive'])}")
    print(f"Kept negative sides: {len(payload['selected_policy']['policy']['keep_negative'])}")
    print(f"Held-out decision: {'PASS' if heldout['decision']['passed'] else 'FAIL'}")
    print(f"Held-out baseline: {heldout['baseline']['metrics']}")
    print(f"Held-out candidate: {heldout['candidate']['metrics']}")
    print(f"Held-out deltas: {heldout['deltas']}")
    print(f"Saved JSON -> {_root_path(args.output_json)}")
    if args.output_html:
        print(f"Saved HTML -> {_root_path(args.output_html)}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())