#!/usr/bin/env python3
"""Audit multi-factor calculation coverage and value ranges on real data."""

from __future__ import annotations

import argparse
import html
import json
import sys
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

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 numpy as np
import pandas as pd

from models.predictor import _build_features
from scripts.optimize_factor_weights import load_stock_codes_file
from scripts.search_multi_factor_weight_config import (
    DEFAULT_FACTORS,
    SUPPORTED_FACTORS,
    add_raw_factor_inputs,
    build_signed_factor_frame,
    fetch_validation_df_months,
    load_twstock_universe_codes,
)


def _now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


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


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


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 summarize_factor_series(values: pd.Series) -> dict[str, Any]:
    clean = pd.to_numeric(values, errors="coerce")
    finite = clean[np.isfinite(clean)]
    n = int(len(clean))
    finite_n = int(len(finite))
    if finite_n == 0:
        return {
            "n": n,
            "finite_n": 0,
            "nonfinite_count": n,
            "zero_pct": 1.0,
            "nonzero_pct": 0.0,
            "positive_pct": 0.0,
            "negative_pct": 0.0,
            "min": None,
            "max": None,
            "mean": None,
            "std": None,
            "p01": None,
            "p99": None,
            "out_of_range_count": 0,
            "all_zero": True,
        }

    zero_mask = finite.abs() <= 1e-12
    return {
        "n": n,
        "finite_n": finite_n,
        "nonfinite_count": int(n - finite_n),
        "zero_pct": round(float(zero_mask.mean()), 4),
        "nonzero_pct": round(float((~zero_mask).mean()), 4),
        "positive_pct": round(float((finite > 1e-12).mean()), 4),
        "negative_pct": round(float((finite < -1e-12).mean()), 4),
        "min": round(float(finite.min()), 6),
        "max": round(float(finite.max()), 6),
        "mean": round(float(finite.mean()), 6),
        "std": round(float(finite.std(ddof=0)), 6),
        "p01": round(float(finite.quantile(0.01)), 6),
        "p99": round(float(finite.quantile(0.99)), 6),
        "out_of_range_count": int((finite.abs() > 1.000001).sum()),
        "all_zero": bool(zero_mask.all()),
    }


def build_factor_warnings(summary: dict[str, Any]) -> list[str]:
    warnings: list[str] = []
    for factor, stats in summary.items():
        if stats["nonfinite_count"]:
            warnings.append(f"{factor}: has {stats['nonfinite_count']} non-finite values")
        if stats["out_of_range_count"]:
            warnings.append(f"{factor}: has {stats['out_of_range_count']} values outside [-1, 1]")
        if stats["all_zero"]:
            warnings.append(f"{factor}: all values are zero; this factor is inert for this data slice")
        elif stats["nonzero_pct"] < 0.02:
            warnings.append(f"{factor}: non-zero coverage is only {stats['nonzero_pct']:.2%}")
    return warnings


def audit_stock(
    stock: str,
    *,
    factors: list[str],
    months: int,
    with_institutional: bool,
    with_margin: bool,
    with_securities_lending: bool,
) -> tuple[pd.DataFrame, dict[str, Any]]:
    df = fetch_validation_df_months(
        stock,
        months=months,
        with_institutional=with_institutional,
        with_margin=with_margin,
        with_securities_lending=with_securities_lending,
    )
    if df is None or df.empty:
        return pd.DataFrame(), {"stock": stock, "rows": 0, "error": "empty history"}

    features = _build_features(df).replace([np.inf, -np.inf], np.nan).fillna(0.0)
    features = add_raw_factor_inputs(features, df).replace([np.inf, -np.inf], np.nan).fillna(0.0)
    factor_frame = build_signed_factor_frame(features, factors)
    factor_frame.insert(0, "stock", stock)
    factor_frame.insert(1, "date", df["date"].astype(str).to_numpy() if "date" in df.columns else np.arange(len(df)))
    return factor_frame, {"stock": stock, "rows": int(len(df)), "factor_rows": int(len(factor_frame))}


def render_html_report(payload: dict[str, Any], output: Path) -> None:
    rows = []
    for factor, stats in payload["aggregate"].items():
        rows.append(
            "<tr>"
            f"<td>{html.escape(factor)}</td>"
            f"<td>{stats['nonzero_pct']:.2%}</td>"
            f"<td>{stats['positive_pct']:.2%}</td>"
            f"<td>{stats['negative_pct']:.2%}</td>"
            f"<td>{html.escape(str(stats['min']))}</td>"
            f"<td>{html.escape(str(stats['max']))}</td>"
            f"<td>{html.escape(str(stats['mean']))}</td>"
            f"<td>{stats['out_of_range_count']}</td>"
            f"<td>{'YES' if stats['all_zero'] else 'NO'}</td>"
            "</tr>"
        )
    warnings = "".join(f"<li>{html.escape(warning)}</li>" for warning in payload.get("warnings", [])) or "<li>None</li>"
    body = f"""<!doctype html>
<html lang="zh-Hant">
<head>
  <meta charset="utf-8">
  <title>Multi-Factor Calculation Audit</title>
  <style>
    body {{ font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; margin: 24px; color: #17202a; }}
    table {{ border-collapse: collapse; width: 100%; font-size: 13px; }}
    th, td {{ border-bottom: 1px solid #e5e7eb; padding: 8px; text-align: left; }}
    th {{ background: #f8fafc; }}
    code {{ white-space: pre-wrap; }}
  </style>
</head>
<body>
  <h1>Multi-Factor Calculation Audit</h1>
  <p>Generated: {html.escape(str(payload.get('generated_at')))}</p>
  <p>Stocks covered: {len(payload.get('covered_stocks', []))}/{len(payload.get('stocks', []))}</p>
  <h2>Warnings</h2>
  <ul>{warnings}</ul>
  <h2>Aggregate Factor Stats</h2>
  <table>
    <tr><th>Factor</th><th>Non-zero</th><th>Positive</th><th>Negative</th><th>Min</th><th>Max</th><th>Mean</th><th>Out of range</th><th>All zero</th></tr>
    {''.join(rows)}
  </table>
</body>
</html>
"""
    output.parent.mkdir(parents=True, exist_ok=True)
    output.write_text(body)


def run(
    *,
    stocks: list[str],
    factors: list[str],
    months: int,
    with_institutional: bool,
    with_margin: bool,
    with_securities_lending: bool,
    output_json: Path,
    output_html: Path | None,
) -> dict[str, Any]:
    unsupported = [factor for factor in factors if factor not in SUPPORTED_FACTORS]
    if unsupported:
        raise ValueError(f"unsupported factors: {', '.join(unsupported)}")

    frames: list[pd.DataFrame] = []
    stock_summaries: list[dict[str, Any]] = []
    for idx, stock in enumerate(stocks, 1):
        print(f"[{idx}/{len(stocks)}] audit {stock}", flush=True)
        try:
            frame, stock_summary = audit_stock(
                stock,
                factors=factors,
                months=months,
                with_institutional=with_institutional,
                with_margin=with_margin,
                with_securities_lending=with_securities_lending,
            )
        except Exception as exc:
            frame = pd.DataFrame()
            stock_summary = {"stock": stock, "rows": 0, "error": str(exc)}
        stock_summaries.append(stock_summary)
        if not frame.empty:
            frames.append(frame)

    covered_stocks = [summary["stock"] for summary in stock_summaries if summary.get("factor_rows", 0) > 0]
    if frames:
        all_factors = pd.concat(frames, ignore_index=True)
        aggregate = {
            factor: summarize_factor_series(all_factors[factor])
            for factor in factors
        }
    else:
        aggregate = {factor: summarize_factor_series(pd.Series(dtype=float)) for factor in factors}

    payload = {
        "experiment": "multi_factor_calculation_audit",
        "generated_at": _now_iso(),
        "stocks": stocks,
        "covered_stocks": covered_stocks,
        "factors": factors,
        "months": months,
        "data_options": {
            "with_institutional": with_institutional,
            "with_margin": with_margin,
            "with_securities_lending": with_securities_lending,
        },
        "stock_summaries": stock_summaries,
        "aggregate": aggregate,
        "warnings": build_factor_warnings(aggregate),
    }
    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 output_html:
        render_html_report(payload, output_html)
    print(f"Covered stocks: {len(covered_stocks)}/{len(stocks)}")
    print(f"Warnings: {len(payload['warnings'])}")
    print(f"Saved JSON -> {output_json}")
    if output_html:
        print(f"Saved HTML -> {output_html}")
    return payload


def main() -> int:
    parser = argparse.ArgumentParser(description="Audit multi-factor calculation coverage and ranges.")
    parser.add_argument("--stocks", default="")
    parser.add_argument("--stocks-file", default="")
    parser.add_argument("--universe", choices=["default", "extended", "twstock"], default="default")
    parser.add_argument("--limit", type=int, default=5)
    parser.add_argument("--factors", default=",".join(DEFAULT_FACTORS))
    parser.add_argument("--months", type=int, default=36)
    parser.add_argument("--no-institutional", action="store_true")
    parser.add_argument("--with-margin", action=argparse.BooleanOptionalAction, default=True)
    parser.add_argument("--with-securities-lending", action=argparse.BooleanOptionalAction, default=True)
    parser.add_argument("--output-json", default="docs/validation_runs/multi_factor_calculation_audit.json")
    parser.add_argument("--output-html", default="docs/validation_runs/multi_factor_calculation_audit.html")
    args = parser.parse_args()

    if args.stocks:
        stocks = _split_csv(args.stocks)
    elif args.stocks_file:
        stocks = load_stock_codes_file(_root_path(args.stocks_file))[: args.limit if args.limit else None]
    elif args.universe == "twstock":
        stocks = load_twstock_universe_codes(args.limit)
    elif args.universe == "extended":
        from scripts.improvement_harness import EXTENDED_STOCKS

        stocks = EXTENDED_STOCKS[: args.limit]
    else:
        from scripts.improvement_harness import DEFAULT_STOCKS

        stocks = DEFAULT_STOCKS[: args.limit]

    payload = run(
        stocks=stocks,
        factors=_split_csv(args.factors),
        months=args.months,
        with_institutional=not args.no_institutional,
        with_margin=args.with_margin,
        with_securities_lending=args.with_securities_lending,
        output_json=_root_path(args.output_json),
        output_html=_root_path(args.output_html) if args.output_html else None,
    )
    return 0 if payload["covered_stocks"] else 1


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