#!/usr/bin/env python3
"""
diagnose_simulation.py — 水力模拟结果快速诊断脚本

用途：对 runprojectreturndict 返回的 link_results + node_results 做多维度分析。
输出排序后的问题清单与统计摘要，供仿真诊断工作流使用。

用法：
    python diagnose_simulation.py <input_json> [options]

输入 JSON 格式（runprojectreturndict 返回结构）：
    {
      "link_results": [
        {"link_id": "...", "flow": float, "velocity": float,
         "headloss": float, "status": "...", "friction": float},
        ...
      ],
      "node_results": [
        {"node_id": "...", "demand": float, "head": float, "pressure": float},
        ...
      },
      ...  // 其他字段（如 tanks, pumps, valves 等）可选
    }

Options:
    --velocity-warn FLOAT    超速预警阈值 (默认 1.2 m/s)
    --velocity-critical FLOAT 超速严重阈值 (默认 2.0 m/s)
    --pressure-warn FLOAT    低压预警阈值 (默认 40 KPA)
    --pressure-critical FLOAT 低压严重阈值 (默认 28 KPA)
    --headloss-top INT       高水损 TOP N (默认 30)
    --flow-top INT           高流量主干管 TOP N (默认 20)
    --output FILE            输出结果到 JSON 文件 (可选)
    --summary-only           只输出统计摘要，不输出明细列表

输出：
    stdout 打印结构化诊断报告，含各维度问题清单和全局统计。
"""

import json
import sys
import argparse
from collections import Counter

# ── 默认阈值 ──────────────────────────────────────────────
DEFAULT_VELOCITY_WARN = 1.2       # m/s
DEFAULT_VELOCITY_CRITICAL = 2.0   # m/s
DEFAULT_PRESSURE_WARN = 40.0      # KPA
DEFAULT_PRESSURE_CRITICAL = 28.0  # KPA
DEFAULT_HEADLOSS_TOP = 30
DEFAULT_FLOW_TOP = 20


def load_input(path: str) -> dict:
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def analyze_links(link_results: list, args):
    """分析管段维度：超流速、高水损、零流量、Closed 状态。"""
    issues = {}

    # ── 超流速管段 ──
    over_velocity_warn = []
    over_velocity_critical = []
    zero_flow = []
    closed_links = []

    for link in link_results:
        v = link.get("velocity", 0) or 0
        q = abs(link.get("flow", 0) or 0)
        status = (link.get("status") or "").strip().lower()
        hl = link.get("headloss", 0) or 0

        # 流速分级
        if v > args.velocity_critical:
            over_velocity_critical.append(link)
        elif v > args.velocity_warn:
            over_velocity_warn.append(link)

        # 零流量（仅统计 Open 状态的）
        if q < 0.001 and status == "open":
            zero_flow.append(link)

        # Closed 管段
        if status != "open":
            closed_links.append(link)

    # 按流速降序排序
    over_velocity_critical.sort(key=lambda x: x.get("velocity", 0), reverse=True)
    over_velocity_warn.sort(key=lambda x: x.get("velocity", 0), reverse=True)

    issues["over_velocity_critical"] = {
        "count": len(over_velocity_critical),
        "threshold": f">{args.velocity_critical} m/s",
        "items": over_velocity_critical[:20],
        "top_velocity": over_velocity_critical[0]["velocity"] if over_velocity_critical else 0,
    }
    issues["over_velocity_warn"] = {
        "count": len(over_velocity_warn),
        "threshold": f">{args.velocity_warn} m/s",
        "items": over_velocity_warn[:20],
    }

    # ── 高水损管段（按绝对值降序） ──
    high_headloss = sorted(link_results, key=lambda x: abs(x.get("headloss", 0) or 0), reverse=True)
    issues["high_headloss"] = {
        "count": len([x for x in link_results if abs(x.get("headloss", 0) or 0) > 0.5]),
        "items": high_headloss[:args.headloss_top],
    }

    # ── 零流量管段 ──
    issues["zero_flow_open"] = {
        "count": len(zero_flow),
    }

    # ── Closed 管段 ──
    issues["closed_links"] = {
        "count": len(closed_links),
        "items": closed_links[:20],
    }

    # ── 高流量主干管 ──
    high_flow = sorted(link_results, key=lambda x: abs(x.get("flow", 0) or 0), reverse=True)
    issues["high_flow_trunk"] = {
        "items": high_flow[:args.flow_top],
    }

    return issues


def analyze_nodes(node_results: list, args):
    """分析节点维度：低压、高压。"""
    issues = {}

    pressure_critical = []
    pressure_warn = []
    high_pressure = []

    for node in node_results:
        p = node.get("pressure", 0) or 0
        if p < args.pressure_critical:
            pressure_critical.append(node)
        elif p < args.pressure_warn:
            pressure_warn.append(node)
        if p > 200:
            high_pressure.append(node)

    pressure_critical.sort(key=lambda x: x.get("pressure", 0))
    pressure_warn.sort(key=lambda x: x.get("pressure", 0))

    total = len(node_results)

    issues["pressure_critical"] = {
        "count": len(pressure_critical),
        "pct": round(len(pressure_critical) / total * 100, 1) if total else 0,
        "threshold": f"<{args.pressure_critical} KPA",
        "items": pressure_critical[:20],
    }
    issues["pressure_warn"] = {
        "count": len(pressure_warn),
        "pct": round(len(pressure_warn) / total * 100, 1) if total else 0,
        "threshold": f"<{args.pressure_warn} KPA (>= {args.pressure_critical})",
        "items": pressure_warn[:20],
    }
    issues["high_pressure"] = {
        "count": len(high_pressure),
        "items": high_pressure[:20],
    }

    return issues, total


def global_stats(link_results, node_results):
    """计算全局统计量。"""
    velocities = [x.get("velocity", 0) or 0 for x in link_results]
    flows = [abs(x.get("flow", 0) or 0) for x in link_results]
    headlosses = [abs(x.get("headloss", 0) or 0) for x in link_results]
    pressures = [x.get("pressure", 0) or 0 for x in node_results]
    demands = [x.get("demand", 0) or 0 for x in node_results]

    # 状态分布
    status_counter = Counter((x.get("status") or "").strip().lower() for x in link_results)

    stats = {
        "link_count": len(link_results),
        "node_count": len(node_results),
        "velocity": {
            "max": round(max(velocities), 4),
            "mean": round(sum(velocities) / len(velocities), 4) if velocities else 0,
            "p95": sorted(velocities)[int(len(velocities) * 0.95)] if velocities else 0,
        },
        "pressure": {
            "max": round(max(pressures), 2),
            "min": round(min(pressures), 2),
            "mean": round(sum(pressures) / len(pressures), 2) if pressures else 0,
            "p5": sorted(pressures)[int(len(pressures) * 0.05)] if pressures else 0,
        },
        "flow": {
            "max": round(max(flows), 4),
        },
        "headloss": {
            "max": round(max(headlosses), 4),
        },
        "status_distribution": dict(status_counter.most_common()),
    }
    return stats


def format_report(stats, link_issues, node_issues, node_total, summary_only=False):
    """输出格式化的诊断报告。"""
    lines = []
    lines.append("=" * 60)
    lines.append("  水力模拟诊断报告")
    lines.append("=" * 60)
    lines.append("")

    # ── 全局统计 ──
    lines.append("【一、全局统计】")
    lines.append(f"  管段总数: {stats['link_count']}")
    lines.append(f"  节点总数: {stats['node_count']}")
    lines.append(f"  流速: 最大 {stats['velocity']['max']} m/s, "
                  f"平均 {stats['velocity']['mean']} m/s, P95 {stats['velocity']['p95']} m/s")
    lines.append(f"  压力: 最大 {stats['pressure']['max']} KPA, "
                  f"最小 {stats['pressure']['min']} KPA, "
                  f"平均 {stats['pressure']['mean']} KPA, P5 {stats['pressure']['p5']} KPA")
    lines.append(f"  最大流量: {stats['flow']['max']} LPS")
    lines.append(f"  最大水损: {stats['headloss']['max']} m")
    lines.append(f"  状态分布: {stats['status_distribution']}")
    lines.append("")

    # ── 超流速 ──
    crit_v = link_issues["over_velocity_critical"]
    warn_v = link_issues["over_velocity_warn"]
    lines.append("【二、超流速管段】")
    lines.append(f"  🔴 严重超速 (> {crit_v['threshold']}): {crit_v['count']} 条, "
                  f"最高 {crit_v['top_velocity']} m/s")
    if crit_v["count"] > 0 and not summary_only:
        for item in crit_v["items"][:10]:
            lines.append(f"    - {item['link_id']}: {item['velocity']} m/s, "
                          f"flow={item.get('flow', 'N/A')} LPS, status={item.get('status', 'N/A')}")
    lines.append(f"  🟡 流速预警 (> {warn_v['threshold']}): {warn_v['count']} 条")
    lines.append("")

    # ── 高水损 ──
    hl = link_issues["high_headloss"]
    lines.append("【三、高水损管段 TOP】")
    lines.append(f"  水损 > 0.5m 的管段: {hl['count']} 条")
    if hl["items"] and not summary_only:
        for item in hl["items"][:10]:
            lines.append(f"    - {item['link_id']}: headloss={item.get('headloss', 0)} m, "
                          f"velocity={item.get('velocity', 0)} m/s")
    lines.append("")

    # ── 低压节点 ──
    pc = node_issues["pressure_critical"]
    pw = node_issues["pressure_warn"]
    lines.append("【四、低压节点】")
    lines.append(f"  🔴 严重低压 ({pc['threshold']}): {pc['count']} 个 "
                  f"({pc['pct']}%)")
    if pc["count"] > 0 and not summary_only:
        for item in pc["items"][:10]:
            lines.append(f"    - {item['node_id']}: {item['pressure']} KPA, "
                          f"demand={item.get('demand', 0)} LPS, head={item.get('head', 0)} m")
    lines.append(f"  🟡 低压预警 ({pw['threshold']}): {pw['count']} 个 "
                  f"({pw['pct']}%)")
    lines.append("")

    # ── 其他发现 ──
    lines.append("【五、其他发现】")
    zf = link_issues["zero_flow_open"]
    cl = link_issues["closed_links"]
    hp = node_issues["high_pressure"]
    lines.append(f"  - 零流量 (Open 状态但 flow≈0): {zf['count']} 条")
    lines.append(f"  - Closed 管段/阀门: {cl['count']} 条")
    if hp["count"] > 0:
        lines.append(f"  - 超压节点 (> 200 KPA): {hp['count']} 个")
    lines.append("")

    # ── 高流量主干 ──
    hft = link_issues["high_flow_trunk"]
    lines.append("【六、高流量主干管 TOP】")
    if hft["items"] and not summary_only:
        for item in hft["items"][:10]:
            lines.append(f"  - {item['link_id']}: |flow|={abs(item.get('flow', 0))} LPS, "
                          f"velocity={item.get('velocity', 0)} m/s")
    lines.append("")

    lines.append("=" * 60)
    return "\n".join(lines)


def main():
    parser = argparse.ArgumentParser(description="水力模拟结果快速诊断")
    parser.add_argument("input", help="输入 JSON 文件路径 (runprojectreturndict 结果)")
    parser.add_argument("--velocity-warn", type=float, default=DEFAULT_VELOCITY_WARN)
    parser.add_argument("--velocity-critical", type=float, default=DEFAULT_VELOCITY_CRITICAL)
    parser.add_argument("--pressure-warn", type=float, default=DEFAULT_PRESSURE_WARN)
    parser.add_argument("--pressure-critical", type=float, default=DEFAULT_PRESSURE_CRITICAL)
    parser.add_argument("--headloss-top", type=int, default=DEFAULT_HEADLOSS_TOP)
    parser.add_argument("--flow-top", type=int, default=DEFAULT_FLOW_TOP)
    parser.add_argument("--output", help="输出 JSON 结果到文件")
    parser.add_argument("--summary-only", action="store_true", help="仅输出统计摘要")
    args = parser.parse_args()

    data = load_input(args.input)
    link_results = data.get("link_results", [])
    node_results = data.get("node_results", [])

    if not link_results and not node_results:
        print("错误: 输入数据中既无 link_results 也无 node_results", file=sys.stderr)
        sys.exit(1)

    # 执行分析
    stats = global_stats(link_results, node_results)
    link_issues = analyze_links(link_results, args)
    node_issues, node_total = analyze_nodes(node_results, args)

    # 输出报告
    report = format_report(stats, link_issues, node_issues, node_total, args.summary_only)
    print(report)

    # 输出 JSON（可选）
    if args.output:
        output_data = {
            "stats": stats,
            "link_issues": {
                k: {kk: vv for kk, vv in v.items() if kk != "items"}
                for k, v in link_issues.items()
            },
            "node_issues": {
                k: {kk: vv for kk, vv in v.items() if kk != "items"}
                for k, v in node_issues.items()
            },
        }
        with open(args.output, "w", encoding="utf-8") as f:
            json.dump(output_data, f, ensure_ascii=False, indent=2)
        print(f"详细结果已输出到: {args.output}")


if __name__ == "__main__":
    main()