from __future__ import annotations

import logging
from typing import Any

import numpy as np
import pandas as pd

from trader_training.io_utils import read_parquet, run_root, write_json, write_text
from trader_training.pm import _pm_frame, _simulate_open_trades, _threshold_candidates, _trade_metrics
from trader_training.schemas import FIT_SPLIT, LATEST_STRESS_SPLIT, TUNE_SPLIT, VALIDATION_LOCKED_SPLIT


DIAGNOSTIC_SPLITS = (FIT_SPLIT, TUNE_SPLIT, VALIDATION_LOCKED_SPLIT, LATEST_STRESS_SPLIT)
PM_EVAL_SPLITS = (TUNE_SPLIT, VALIDATION_LOCKED_SPLIT, LATEST_STRESS_SPLIT)


def diagnose_training_run(args: Any) -> None:
    root = run_root(args)
    label_summary = _label_summary(root)
    pm_summary = _pm_summary(root)
    payload = {
        "run_id": args.run_id,
        "label_summary": label_summary,
        "pm_summary": pm_summary,
        "conclusion": _diagnostic_conclusion(pm_summary),
    }
    write_json(root / "diagnostics" / "training_failure_diagnostics.json", _jsonable(payload))
    write_text(root / "diagnostics" / "training_failure_diagnostics.md", _markdown_report(payload))
    logging.info(
        "trader.training.diagnostics_written runId=%s conclusion=%s path=%s",
        args.run_id,
        payload["conclusion"]["status"],
        root / "diagnostics" / "training_failure_diagnostics.md",
    )


def _label_summary(root) -> dict[str, Any]:
    direction = read_parquet(root / "label" / "direction_labels.parquet")
    entry = read_parquet(root / "label" / "entry_labels.parquet")
    summary: dict[str, Any] = {}
    for split_id in DIAGNOSTIC_SPLITS:
        direction_split = direction[direction["split_id"].eq(split_id)].copy()
        entry_split = entry[entry["split_id"].eq(split_id)].copy()
        item: dict[str, Any] = {"direction": {}, "entry": {}}
        if not direction_split.empty:
            item["direction"] = {
                "rows": len(direction_split),
                "label_ratio": direction_split["direction_label"].value_counts(normalize=True).round(6).to_dict(),
                "future_return_bps_quantile": _quantiles(direction_split["future_return_bps"], (0.01, 0.05, 0.25, 0.5, 0.75, 0.95, 0.99)),
            }
        if not entry_split.empty:
            grouped = entry_split.groupby("side", observed=False)
            item["entry"] = {
                "rows": len(entry_split),
                "target_rate_by_side": grouped["entry_target"].mean().round(6).to_dict(),
                "edge_mean_by_side": grouped["expected_net_edge_bps"].mean().round(6).to_dict(),
                "edge_quantile_by_side": {
                    str(side): _quantiles(group["expected_net_edge_bps"], (0.05, 0.5, 0.95))
                    for side, group in grouped
                },
            }
        summary[split_id] = item
    return summary


def _pm_summary(root) -> dict[str, Any]:
    summary: dict[str, Any] = {}
    for split_id in PM_EVAL_SPLITS:
        frame = _pm_frame(root, split_id)
        item = {
            "rows": len(frame),
            "score_distribution": _score_distribution(frame),
            "gate_funnel": _gate_funnel(frame),
            "relaxed_variants": _relaxed_variants(frame),
            "top_bucket_edge": _top_bucket_edge(frame),
            "grid_search_any_trade": _grid_trade_summary(frame),
        }
        summary[split_id] = item
    return summary


def _score_distribution(frame: pd.DataFrame) -> dict[str, dict[str, float]]:
    columns = [
        "long_prob",
        "short_prob",
        "long_entry_prob",
        "short_entry_prob",
        "market_risk_prob",
        "pred_long_expected_net_edge_bps",
        "pred_short_expected_net_edge_bps",
        "actual_long_expected_net_edge_bps",
        "actual_short_expected_net_edge_bps",
    ]
    return {column: _quantiles(frame[column], (0.0, 0.05, 0.5, 0.95, 1.0)) for column in columns}


def _gate_funnel(frame: pd.DataFrame) -> dict[str, Any]:
    thresholds = {
        "long_open_prob": 0.54,
        "short_open_prob": 0.54,
        "min_entry_prob": 0.50,
        "max_market_risk_prob": 0.55,
        "min_expected_edge_bps": 1.0,
        "min_direction_margin": 0.02,
    }
    long_steps = {
        "long_prob >= 0.54": frame["long_prob"] >= thresholds["long_open_prob"],
        "long_prob - short_prob >= 0.02": (frame["long_prob"] - frame["short_prob"]) >= thresholds["min_direction_margin"],
        "long_entry_prob >= 0.50": frame["long_entry_prob"] >= thresholds["min_entry_prob"],
        "market_risk_prob <= 0.55": frame["market_risk_prob"] <= thresholds["max_market_risk_prob"],
        "pred_long_expected_net_edge_bps >= 1.0": frame["pred_long_expected_net_edge_bps"] >= thresholds["min_expected_edge_bps"],
    }
    short_steps = {
        "short_prob >= 0.54": frame["short_prob"] >= thresholds["short_open_prob"],
        "short_prob - long_prob >= 0.02": (frame["short_prob"] - frame["long_prob"]) >= thresholds["min_direction_margin"],
        "short_entry_prob >= 0.50": frame["short_entry_prob"] >= thresholds["min_entry_prob"],
        "market_risk_prob <= 0.55": frame["market_risk_prob"] <= thresholds["max_market_risk_prob"],
        "pred_short_expected_net_edge_bps >= 1.0": frame["pred_short_expected_net_edge_bps"] >= thresholds["min_expected_edge_bps"],
    }
    return {
        "thresholds": thresholds,
        "long": _cumulative_gate_counts(long_steps, len(frame)),
        "short": _cumulative_gate_counts(short_steps, len(frame)),
    }


def _cumulative_gate_counts(steps: dict[str, pd.Series], total_rows: int) -> dict[str, Any]:
    mask = np.ones(total_rows, dtype=bool)
    cumulative = []
    single = {}
    for name, step in steps.items():
        values = step.to_numpy(dtype=bool)
        single[name] = int(values.sum())
        mask &= values
        cumulative.append({"gate": name, "rows_after_gate": int(mask.sum())})
    return {"single_gate_pass": single, "cumulative": cumulative}


def _relaxed_variants(frame: pd.DataFrame) -> dict[str, Any]:
    variants = {
        "no_risk_no_edge": {"prob": 0.54, "entry": 0.50, "margin": 0.02, "risk": 1.0, "edge": -99.0},
        "entry_only_55": {"prob": 0.0, "entry": 0.55, "margin": -99.0, "risk": 1.0, "edge": -99.0},
        "direction_only_54": {"prob": 0.54, "entry": 0.0, "margin": 0.02, "risk": 1.0, "edge": -99.0},
        "very_loose": {"prob": 0.50, "entry": 0.45, "margin": 0.0, "risk": 1.0, "edge": -99.0},
    }
    result: dict[str, Any] = {}
    for name, thresholds in variants.items():
        trades = _variant_trades(frame, thresholds)
        result[name] = _plain_trade_metrics(trades)
    return result


def _variant_trades(frame: pd.DataFrame, thresholds: dict[str, float]) -> pd.DataFrame:
    long_mask = (
        (frame["long_prob"] >= thresholds["prob"])
        & ((frame["long_prob"] - frame["short_prob"]) >= thresholds["margin"])
        & (frame["long_entry_prob"] >= thresholds["entry"])
        & (frame["market_risk_prob"] <= thresholds["risk"])
        & (frame["pred_long_expected_net_edge_bps"] >= thresholds["edge"])
    )
    short_mask = (
        (frame["short_prob"] >= thresholds["prob"])
        & ((frame["short_prob"] - frame["long_prob"]) >= thresholds["margin"])
        & (frame["short_entry_prob"] >= thresholds["entry"])
        & (frame["market_risk_prob"] <= thresholds["risk"])
        & (frame["pred_short_expected_net_edge_bps"] >= thresholds["edge"])
    )
    long = frame.loc[long_mask].copy()
    long["side"] = "LONG"
    long["actual_edge_bps"] = long["actual_long_expected_net_edge_bps"]
    short = frame.loc[short_mask].copy()
    short["side"] = "SHORT"
    short["actual_edge_bps"] = short["actual_short_expected_net_edge_bps"]
    return pd.concat([long, short], ignore_index=True)


def _plain_trade_metrics(trades: pd.DataFrame) -> dict[str, Any]:
    if trades.empty:
        return {"rows": 0, "win_rate": 0.0, "avg_actual_edge_bps": 0.0}
    return {
        "rows": len(trades),
        "win_rate": float((trades["actual_edge_bps"].astype(float) > 0).mean()),
        "avg_actual_edge_bps": float(trades["actual_edge_bps"].astype(float).mean()),
        "side_counts": trades["side"].value_counts().to_dict(),
    }


def _top_bucket_edge(frame: pd.DataFrame) -> dict[str, Any]:
    side = np.where(frame["long_prob"] >= frame["short_prob"], "LONG", "SHORT")
    side_prob = np.where(side == "LONG", frame["long_prob"], frame["short_prob"])
    side_edge = np.where(side == "LONG", frame["actual_long_expected_net_edge_bps"], frame["actual_short_expected_net_edge_bps"])
    direction_frame = pd.DataFrame({"score": side_prob, "actual_edge_bps": side_edge, "side": side})
    direction_top = {}
    for fraction in (0.01, 0.02, 0.05, 0.10):
        top = direction_frame.sort_values("score", ascending=False).head(max(1, int(len(direction_frame) * fraction)))
        direction_top[str(fraction)] = _plain_trade_metrics(top.rename(columns={"actual_edge_bps": "actual_edge_bps"}))
    return {
        "direction_top_score": direction_top,
        "long_entry_prob_deciles": _decile_edge(frame, "long_entry_prob", "actual_long_expected_net_edge_bps", "long_entry_target"),
        "short_entry_prob_deciles": _decile_edge(frame, "short_entry_prob", "actual_short_expected_net_edge_bps", "short_entry_target"),
    }


def _decile_edge(frame: pd.DataFrame, score_col: str, edge_col: str, target_col: str) -> list[dict[str, Any]]:
    sample = frame[[score_col, edge_col, target_col]].dropna().copy()
    if sample.empty:
        return []
    sample["bucket"] = pd.qcut(sample[score_col].rank(method="first"), 10, labels=False) + 1
    rows = []
    for bucket, group in sample.groupby("bucket", observed=False):
        rows.append(
            {
                "bucket": int(bucket),
                "rows": len(group),
                "score_min": float(group[score_col].min()),
                "score_max": float(group[score_col].max()),
                "hit_rate": float(group[target_col].mean()),
                "avg_actual_edge_bps": float(group[edge_col].mean()),
            }
        )
    return rows


def _grid_trade_summary(frame: pd.DataFrame) -> dict[str, Any]:
    nonzero = 0
    best_by_count = None
    best_metrics = None
    for thresholds in _threshold_candidates():
        trades = _simulate_open_trades(frame, thresholds)
        metrics = _trade_metrics(trades)
        if metrics["trade_count"] > 0:
            nonzero += 1
        if best_metrics is None or metrics["trade_count"] > best_metrics["trade_count"]:
            best_by_count = thresholds
            best_metrics = metrics
    return {
        "candidate_count": len(_threshold_candidates()),
        "candidates_with_trade": nonzero,
        "best_by_trade_count": best_by_count,
        "best_metrics": best_metrics,
    }


def _diagnostic_conclusion(pm_summary: dict[str, Any]) -> dict[str, Any]:
    tune = pm_summary.get(TUNE_SPLIT, {})
    gate = tune.get("gate_funnel", {})
    long_single = gate.get("long", {}).get("single_gate_pass", {})
    short_single = gate.get("short", {}).get("single_gate_pass", {})
    pred_edge_blocked = (
        long_single.get("pred_long_expected_net_edge_bps >= 1.0", 0) == 0
        and short_single.get("pred_short_expected_net_edge_bps >= 1.0", 0) == 0
    )
    relaxed = tune.get("relaxed_variants", {})
    any_relaxed_positive = any(item.get("avg_actual_edge_bps", 0.0) > 0 for item in relaxed.values())
    if pred_edge_blocked and not any_relaxed_positive:
        return {
            "status": "MODEL_SIGNAL_NOT_TRADABLE",
            "plain_reason": "Entry 预测的净收益基本都是负数；即使放松风险和收益门槛，选出来的样本平均仍亏。",
            "next_action": "优先重查 Entry 标签和价格计划，再考虑更强模型；不要直接放松 PM 阈值上线。",
        }
    if pred_edge_blocked:
        return {
            "status": "ENTRY_EDGE_GATE_BLOCKED",
            "plain_reason": "PM 没有交易主要是 Entry 预测净收益过低。",
            "next_action": "重训 Entry 或调整价格计划后再搜索 PM 阈值。",
        }
    return {
        "status": "NEEDS_MANUAL_REVIEW",
        "plain_reason": "没有发现单一硬挡板，需要人工继续看各模型分数和回测明细。",
        "next_action": "查看 diagnostics 报告中的漏斗和放松阈值结果。",
    }


def _quantiles(series: pd.Series, points: tuple[float, ...]) -> dict[str, float]:
    values = pd.to_numeric(series, errors="coerce").replace([np.inf, -np.inf], np.nan).dropna()
    if values.empty:
        return {}
    result = values.quantile(list(points)).round(6).to_dict()
    return {str(key): float(value) for key, value in result.items()}


def _markdown_report(payload: dict[str, Any]) -> str:
    lines = [
        "# Trader Training Failure Diagnostics",
        "",
        f"- run_id: `{payload['run_id']}`",
        f"- status: `{payload['conclusion']['status']}`",
        f"- 结论: {payload['conclusion']['plain_reason']}",
        f"- 下一步: {payload['conclusion']['next_action']}",
        "",
        "## 标签分布",
        "",
    ]
    for split_id, item in payload["label_summary"].items():
        direction = item.get("direction", {})
        entry = item.get("entry", {})
        lines.append(f"### {split_id}")
        lines.append("")
        if direction:
            lines.append(f"- Direction 行数: {direction['rows']}")
            lines.append(f"- Direction 标签比例: `{direction['label_ratio']}`")
            lines.append(f"- 45 分钟未来收益分位: `{direction['future_return_bps_quantile']}`")
        if entry:
            lines.append(f"- Entry 行数: {entry['rows']}")
            lines.append(f"- Entry 命中率: `{entry['target_rate_by_side']}`")
            lines.append(f"- Entry 平均净收益: `{entry['edge_mean_by_side']}`")
        lines.append("")

    lines.extend(["## PM 挡单漏斗", ""])
    for split_id, item in payload["pm_summary"].items():
        lines.append(f"### {split_id}")
        lines.append("")
        lines.append(f"- 样本数: {item['rows']}")
        lines.append(f"- 网格里有交易的候选数: {item['grid_search_any_trade']['candidates_with_trade']} / {item['grid_search_any_trade']['candidate_count']}")
        lines.append("")
        for side in ("long", "short"):
            lines.append(f"#### {side.upper()}")
            lines.append("")
            rows = item["gate_funnel"][side]["cumulative"]
            lines.append("| 条件 | 剩余样本 |")
            lines.append("| --- | ---: |")
            for row in rows:
                lines.append(f"| {row['gate']} | {row['rows_after_gate']} |")
            lines.append("")
        lines.append("#### 放松条件后的结果")
        lines.append("")
        lines.append("| 方案 | 样本数 | 胜率 | 平均真实净收益bps |")
        lines.append("| --- | ---: | ---: | ---: |")
        for name, metrics in item["relaxed_variants"].items():
            lines.append(
                f"| {name} | {metrics['rows']} | {metrics['win_rate']:.4f} | {metrics['avg_actual_edge_bps']:.4f} |"
            )
        lines.append("")
    return "\n".join(lines) + "\n"


def _jsonable(value: Any) -> Any:
    if isinstance(value, dict):
        return {str(key): _jsonable(item) for key, item in value.items()}
    if isinstance(value, list):
        return [_jsonable(item) for item in value]
    if isinstance(value, tuple):
        return [_jsonable(item) for item in value]
    if isinstance(value, (np.integer,)):
        return int(value)
    if isinstance(value, (np.floating,)):
        return float(value)
    if isinstance(value, np.ndarray):
        return value.tolist()
    return value