#!/usr/bin/env python3
"""
check_fort9.py — 解析 {stage}.9 收敛轨迹，判数值收敛；检测 STOP / NaN / 雪崩。

复刻 DCTS check_fort9（conv_check.rs:68）的核心逻辑。
退出码：0=收敛 /1=未收敛或异常 /2=文件缺失或无有效数据。

⚠️ **只对已完成的运行使用本脚本**：TLUSTY 是边算边写 fort.9 的（每次迭代的
深度行逐行写出），运行中实时读取，最新一拍的 max_relc 读到的是半写状态，数值
不可信（实测：运行中读到 iter6=0.477、iter9=0.00176 貌似收敛，进程退出后同一
迭代实为 1060 / 5.07e4）。等 run.sh 完全结束、进程退出后再跑本脚本。

判据：
  - 每次 ITER 的 max_relc = 该迭代所有深度行 |MAXIMUM| 的最大值
  - 末次迭代 max_relc < chmax（默认 0.001）且有限 = 收敛
  - 单调下降=健康；骤涨几个数量级=雪崩
  - fort.6（stdout）里的 "STOP in SOLVE after ITER N" = 求解器发散中止

用法：
    python3 check_fort9.py outputs/nl.9 --teff 60000
    python3 check_fort9.py outputs/nc.9 --chmax 0.001 --stdout outputs/nc.6
"""
import argparse
import os
import re
import sys

# 复刻 conv_check.rs:44 parse_fortran_float
_NO_E_EXP_RE = re.compile(r"^([+-]?[\d.]+)([+-]\d+)$")


def parse_fortran_float(s):
    s = s.strip()
    try:
        return float(s)
    except ValueError:
        pass
    m = _NO_E_EXP_RE.match(s)
    if m:
        try:
            return float(f"{m.group(1)}E{m.group(2)}")
        except ValueError:
            pass
    return None


# 复刻 conv_check.rs:107 FORT9_RE：iter depth 5浮点 2整数
FORT9_RE = re.compile(
    r"^\s*(\d+)\s+(\d+)\s+([-+\dE.]+)\s+([-+\dE.]+)\s+([-+\dE.]+)\s+([-+\dE.]+)\s+([-+\dE.]+)\s+(\d+)\s+(\d+)\s*$"
)
# NaN 模式（conv_check.rs:18）
NAN_RE = re.compile(r"(?i)(\bnan\b|\binf(?:inity)?\b|\*{3,}|[eE]\+(?:3\d{2}|[4-9]\d{2,}))")
# STOP 模式（conv_check.rs SOLVER_STOP_RE）
STOP_RE = re.compile(r"STOP\s+in\s+SOLVE\s+after\s+ITER\s+(\d+)", re.IGNORECASE)


def parse_fort9(path, chmax):
    """返回 (iters, last_iter, last_max_relc, converged, n_depths_last)。
    iters = [(iter_no, max_relc, n_depths), ...]"""
    iters = []
    cur_iter = None
    cur_max = 0.0
    cur_n = 0
    last_iter = None

    with open(path, "r", errors="replace") as f:
        for line in f:
            m = FORT9_RE.match(line)
            if not m:
                continue
            try:
                it = int(m.group(1))
            except ValueError:
                continue
            maximum = parse_fortran_float(m.group(7))
            if maximum is None:
                continue
            if cur_iter != it:
                if cur_iter is not None:
                    iters.append((cur_iter, cur_max, cur_n))
                cur_iter, cur_max, cur_n = it, 0.0, 0
            cur_n += 1
            am = abs(maximum)
            if am > cur_max:
                cur_max = am
            last_iter = it
    if cur_iter is not None:
        iters.append((cur_iter, cur_max, cur_n))

    if not iters or last_iter is None:
        return [], None, float("inf"), False, 0

    last_max_relc = iters[-1][1]
    converged = (last_max_relc < chmax) and (last_max_relc == last_max_relc) and (abs(last_max_relc) != float("inf"))
    return iters, last_iter, last_max_relc, converged, iters[-1][2]


def detect_avalanche(iters):
    """检测雪崩：相邻迭代 max_relc 涨 >1e3 倍。返回告警列表。"""
    warns = []
    for i in range(1, len(iters)):
        prev = iters[i - 1][1]
        cur = iters[i][1]
        if prev > 0 and cur / prev > 1e3:
            warns.append(f"  iter {iters[i-1][0]}→{iters[i][0]}: max_relc {prev:.3g} → {cur:.3g}（涨 {cur/prev:.0g}×，雪崩）")
    return warns


def scan_stop(stdout_path):
    """扫 fort.6 找 STOP in SOLVE。返回 iter 号或 None。"""
    if not stdout_path or not os.path.exists(stdout_path):
        return None
    last = None
    with open(stdout_path, "r", errors="replace") as f:
        for line in f:
            m = STOP_RE.search(line)
            if m:
                last = int(m.group(1))
    return last


def scan_nan(path):
    if not path or not os.path.exists(path):
        return 0
    n = 0
    with open(path, "r", errors="replace") as f:
        for line in f:
            if NAN_RE.search(line):
                n += 1
    return n


def main():
    ap = argparse.ArgumentParser(description="fort.9 收敛轨迹分析")
    ap.add_argument("fort9", help="{stage}.9 路径")
    ap.add_argument("--teff", type=float, default=None, help="Teff（上下文，非必填）")
    ap.add_argument("--chmax", type=float, default=0.001, help="收敛阈值（默认 0.001）")
    ap.add_argument("--stdout", default=None, help="对应 {stage}.6 路径，用于 STOP 检测（默认自动推断）")
    ap.add_argument("--expected-niter", type=int, default=None,
                    help="配置的 NITER（来自 nst）。若末次 iter 远小于此且未收敛，提示疑似被 timeout 截断")
    args = ap.parse_args()

    if not os.path.exists(args.fort9):
        print(f"⚠ 文件不存在: {args.fort9}")
        return 2

    # 自动推断 stdout（同目录同名 .6）：把末尾 .9 换成 .6
    stdout = args.stdout
    if stdout is None:
        if args.fort9.endswith(".9"):
            cand = args.fort9[:-2] + ".6"
        else:
            cand = args.fort9 + "6"
        stdout = cand if os.path.exists(cand) else None

    if args.chmax <= 0:
        print(f"❌ 非法 chmax={args.chmax}（须为正有限数）")
        return 1

    iters, last_iter, last_max_relc, converged, n_depths = parse_fort9(args.fort9, args.chmax)
    if not iters:
        print(f"⚠ {args.fort9} 无有效迭代数据")
        return 2

    stop_iter = scan_stop(stdout)
    nan9 = scan_nan(args.fort9)
    # 截断检测：① 附近有 .TIMEOUT 标记文件（run.sh 在 rc=124 时写到 outputs/{stage}.TIMEOUT）；
    #          ② 末次 iter 远小于 --expected-niter 且未收敛（疑似 timeout）。
    import glob
    timeout_marker = None
    d = os.path.dirname(os.path.abspath(args.fort9))
    for pat in (os.path.join(d, "*.TIMEOUT"),
                os.path.join(d, "outputs", "*.TIMEOUT"),
                os.path.join(d, "..", "outputs", "*.TIMEOUT")):
        hits = glob.glob(pat)
        if hits:
            timeout_marker = ",".join(os.path.basename(h) for h in hits)
            break
    likely_truncated = (timeout_marker is not None or
                        (args.expected_niter and last_iter < args.expected_niter - 1 and not converged and stop_iter is None))

    print(f"=== 收敛轨迹: {args.fort9} (chmax={args.chmax}" +
          (f", teff={args.teff}" if args.teff else "") + ") ===")
    print(f"  迭代次数:     {len(iters)}（末次 ITER {last_iter}）")
    print(f"  末 max_relc:  {last_max_relc:.4g}  → {'✅ < chmax 收敛' if converged else '❌ ≥ chmax 未收敛'}")
    print(f"  末次深度数:   {n_depths}")
    if stop_iter is not None:
        print(f"  STOP:         ❌ STOP in SOLVE after ITER {stop_iter}（求解器发散中止）")
    else:
        print(f"  STOP:         无")
    print(f"  fort.9 NaN行: {nan9}")
    if likely_truncated:
        detail = f"（标记文件: {timeout_marker}）" if timeout_marker else (
            f"（仅 iter {last_iter}/{args.expected_niter}，疑似 timeout 截断）" if args.expected_niter else "")
        print(f"  截断:         ⚠ 疑似被 timeout 杀死 {detail}——轨迹不完整，结论需谨慎")

    # 轨迹（最多显示首/末若干 + 全部若 ≤15）
    print("  轨迹 (iter: max_relc):")
    show = iters if len(iters) <= 20 else iters[:8] + [None] + iters[-5:]
    for it in show:
        if it is None:
            print("    ...")
        else:
            print(f"    iter {it[0]:>3}: {it[1]:.4g}  ({it[2]} 深度)")

    warns = detect_avalanche(iters)
    if warns:
        print("  ⚠ 雪崩告警:")
        for w in warns:
            print(w)

    # 伪收敛启发式：max_relc 恰为 0 或极小且迭代很少 → 极可能是"在 NaN 模型上迭代"。
    # 真实收敛要多次迭代、max_relc 趋近但极少精确到 0。
    pseudo_warn = False
    if last_max_relc == 0.0 or (len(iters) <= 2 and last_max_relc < 1e-10):
        pseudo_warn = True

    # 结论
    real_converged = converged and stop_iter is None and nan9 == 0
    print()
    if real_converged:
        print(f"✅ 数值收敛（末 max_relc {last_max_relc:.4g} < {args.chmax}，无 STOP，无 NaN）")
        if pseudo_warn:
            print("   ⚠⚠ 伪收敛嫌疑：max_relc≈0 且迭代极少，极可能是『在 NaN 模型上迭代』(NaN 上相对变化=0)。")
            print("      务必运行 check_temperature.py 查温度结构——若 .7 全 NaN，则此『收敛』无效。")
        else:
            print("   注意：数值收敛 ≠ 物理收敛。请再用 check_temperature.py 查温度结构（防伪收敛）。")
        return 0
    print("❌ 未数值收敛:" if not real_converged else "")
    reasons = []
    if not converged:
        reasons.append(f"末 max_relc {last_max_relc:.4g} ≥ chmax {args.chmax}")
    if stop_iter is not None:
        reasons.append(f"STOP in SOLVE after ITER {stop_iter}")
    if nan9 > 0:
        reasons.append(f"fort.9 有 {nan9} 行 NaN/坏值")
    if warns:
        reasons.append("迭代轨迹有雪崩")
    for r in reasons:
        print(f"   - {r}")
    return 1


if __name__ == "__main__":
    sys.exit(main())