#!/usr/bin/env python3
"""Verify FemType=5 (3D eigen frequency) outputs without requiring a pre-made reference.

Checks (in order):
  1. Self-consistency: normE == sqrt(|Ex|^2+|Ey|^2+|Ez|^2) per mode
  2. Independent freq: SciPy eigs(A,B,sigma) on exported A/B vs OutFile/freq
  3. Optional MATLAB: compare freq/normE if mat_outdir contains those files

Usage:
  python tools/verify_eigen3d.py build/Release/OutFile
  python tools/verify_eigen3d.py build/Release/OutFile --json eigen3d.json
  python tools/verify_eigen3d.py build/Release/OutFile --mat-outdir path/to/matlab/OutFile
"""
from __future__ import annotations

import argparse
import json
import re
from pathlib import Path

import numpy as np
from scipy import sparse
from scipy.sparse.linalg import eigs

C_LIGHT = 2.9979e8


def load_real_coo(prefix: Path) -> sparse.csr_matrix:
    ai = np.loadtxt(prefix / "Ai.txt", dtype=np.int64)
    aj = np.loadtxt(prefix / "Aj.txt", dtype=np.int64)
    av = np.loadtxt(prefix / "Av.txt", dtype=np.float64)
    n = int(max(ai.max(), aj.max()) + 1)
    return sparse.csr_matrix((av, (ai, aj)), shape=(n, n))


def load_complex_vector_modes(path: Path) -> list[np.ndarray]:
    """Ex/Ey/Ez: each mode is nv complex values."""
    modes: list[list[complex]] = []
    cur: list[complex] = []
    for line in path.read_text().splitlines():
        line = line.strip()
        if not line:
            continue
        if line == "//" or line.startswith("//"):
            if line.startswith("//") and len(line) > 2:
                parts = line[2:].strip().split()
                if parts:
                    cur.append(complex(float(parts[0]), float(parts[1]) if len(parts) > 1 else 0.0))
            if cur:
                modes.append(cur)
                cur = []
            continue
        parts = line.split()
        cur.append(complex(float(parts[0]), float(parts[1])))
    if cur:
        modes.append(cur)
    return [np.asarray(m, dtype=np.complex128) for m in modes]


def load_real_modes(path: Path) -> list[np.ndarray]:
    modes: list[list[float]] = []
    cur: list[float] = []
    for line in path.read_text().splitlines():
        line = line.strip()
        if not line:
            continue
        if line == "//" or line.startswith("//"):
            if line.startswith("//") and len(line) > 2:
                rest = line[2:].strip().split()
                if rest:
                    cur.append(float(rest[0]))
            if cur:
                modes.append(cur)
                cur = []
            continue
        cur.append(float(line.split()[0]))
    if cur:
        modes.append(cur)
    return [np.asarray(m, dtype=np.float64) for m in modes]


def load_freq(path: Path) -> np.ndarray:
    rows = []
    for line in path.read_text().splitlines():
        line = line.strip()
        if not line:
            continue
        parts = line.split()
        rows.append(complex(float(parts[0]), float(parts[1]) if len(parts) > 1 else 0.0))
    return np.asarray(rows, dtype=np.complex128)


def self_check(outdir: Path) -> dict:
    ex = load_complex_vector_modes(outdir / "Ex")
    ey = load_complex_vector_modes(outdir / "Ey")
    ez = load_complex_vector_modes(outdir / "Ez")
    ne = load_real_modes(outdir / "normE")
    n_mode = len(ne)
    if not (len(ex) == len(ey) == len(ez) == n_mode):
        raise SystemExit(
            f"mode count mismatch: Ex={len(ex)} Ey={len(ey)} Ez={len(ez)} normE={n_mode}"
        )
    rels, corrs = [], []
    for i in range(n_mode):
        n_calc = np.sqrt(np.abs(ex[i]) ** 2 + np.abs(ey[i]) ** 2 + np.abs(ez[i]) ** 2)
        n_ref = ne[i]
        if n_ref.size != n_calc.size:
            raise SystemExit(f"mode {i}: size {n_ref.size} vs {n_calc.size}")
        denom = np.linalg.norm(n_ref)
        rel = np.linalg.norm(n_calc - n_ref) / (denom if denom > 0 else 1.0)
        corr = np.corrcoef(n_calc, n_ref)[0, 1] if denom > 0 else 1.0
        rels.append(rel)
        corrs.append(corr)
    return {
        "n_modes": n_mode,
        "n_vertex": int(ne[0].size) if ne else 0,
        "normE_rel_err": rels,
        "normE_corr": corrs,
    }


def scipy_freq_check(outdir: Path, json_path: Path | None) -> dict:
    A = load_real_coo_with_names(outdir, "A")
    B = load_real_coo_with_names(outdir, "B")
    freq_cpp = load_freq(outdir / "freq")
    k = freq_cpp.size

    ff0 = None
    if json_path and json_path.exists():
        js = json.loads(json_path.read_text(encoding="utf-8"))
        ff0 = float(js.get("searchValue", 0))
        lam0 = float(js.get("lambda0", 0.8))
        if ff0 <= 0:
            ff0 = C_LIGHT / lam0

    if ff0 is None or ff0 <= 0:
        ff0 = C_LIGHT / 0.8

    k0 = 2.0 * np.pi / (C_LIGHT / ff0)
    sigma = k0 * k0

    print(f"  SciPy eigs: n={A.shape[0]}, k={k}, sigma={sigma:.6g} (ff0={ff0:.6g} Hz) ...")
    vals, _ = eigs(A, k=k, M=B, sigma=sigma, which="LM")
    vals = np.sort(vals.real)[::-1]
    freq_scipy = C_LIGHT / (2.0 * np.pi / np.sqrt(vals))

    freq_cpp_hz = freq_cpp.real
    freq_scipy_hz = np.sort(freq_scipy)[::-1]

    # Match modes by nearest frequency (ordering may differ)
    used = set()
    pairs = []
    for f in freq_cpp_hz:
        j = int(np.argmin([abs(f - s) if idx not in used else np.inf for idx, s in enumerate(freq_scipy_hz)]))
        used.add(j)
        pairs.append((f, freq_scipy_hz[j], abs(f - freq_scipy_hz[j]) / max(abs(f), 1e-30)))

    rels = [p[2] for p in pairs]
    return {
        "freq_cpp_Hz": freq_cpp_hz.tolist(),
        "freq_scipy_Hz": freq_scipy_hz.tolist(),
        "freq_rel_err_matched": rels,
        "freq_max_rel_err": float(max(rels) if rels else 0.0),
    }


def load_real_coo_with_names(outdir: Path, which: str) -> sparse.csr_matrix:
    ai = np.loadtxt(outdir / f"{which}i.txt", dtype=np.int64)
    aj = np.loadtxt(outdir / f"{which}j.txt", dtype=np.int64)
    av = np.loadtxt(outdir / f"{which}v.txt", dtype=np.float64)
    n = int(max(ai.max(), aj.max()) + 1)
    return sparse.csr_matrix((av, (ai, aj)), shape=(n, n))


def compare_matlab(outdir: Path, mat_outdir: Path) -> dict:
    freq_cpp = load_freq(outdir / "freq").real
    if not (mat_outdir / "freq").exists():
        return {"skipped": "matlab OutFile/freq not found"}
    freq_mat = load_freq(mat_outdir / "freq").real
    k = min(freq_cpp.size, freq_mat.size)
    freq_cpp = np.sort(freq_cpp[:k])[::-1]
    freq_mat = np.sort(freq_mat[:k])[::-1]
    rel_f = np.linalg.norm(freq_cpp - freq_mat) / np.linalg.norm(freq_mat)

    out = {"freq_rel_err": float(rel_f)}
    if (mat_outdir / "normE").exists():
        ne_cpp_modes = load_real_modes(outdir / "normE")
        ne_mat_modes = load_real_modes(mat_outdir / "normE")
        m = min(len(ne_cpp_modes), len(ne_mat_modes))
        rels, corrs = [], []
        for i in range(m):
            a, b = ne_cpp_modes[i], ne_mat_modes[i]
            if a.size != b.size:
                continue
            rels.append(np.linalg.norm(a - b) / np.linalg.norm(b))
            corrs.append(np.corrcoef(a, b)[0, 1])
        out["normE_rel_err_per_mode"] = rels
        out["normE_corr_per_mode"] = corrs
    return out


def main() -> None:
    parser = argparse.ArgumentParser(description="Verify 3D eigenfrequency OutFile")
    parser.add_argument("outdir", type=Path, help="C++ OutFile directory")
    parser.add_argument("--json", type=Path, default=None, help="eigen3d.json for sigma")
    parser.add_argument("--mat-outdir", type=Path, default=None, help="MATLAB OutFile for comparison")
    parser.add_argument("--skip-scipy", action="store_true", help="Skip SciPy eigs (slow)")
    args = parser.parse_args()

    outdir = args.outdir.resolve()
    if not (outdir / "normE").exists():
        raise SystemExit(f"missing {outdir / 'normE'}")

    print("=== 1. Self-check: normE vs Ex/Ey/Ez ===")
    sc = self_check(outdir)
    print(f"  modes={sc['n_modes']}, vertices/mode={sc['n_vertex']}")
    for i, (rel, corr) in enumerate(zip(sc["normE_rel_err"], sc["normE_corr"])):
        ok = "OK" if rel < 1e-10 else "WARN"
        print(f"  mode {i}: rel_err={rel:.3e}, corr={corr:.6f}  [{ok}]")

    if not args.skip_scipy and (outdir / "Bi.txt").exists():
        print("\n=== 2. Independent freq: SciPy eigs(A,B) vs freq ===")
        js = args.json
        if js is None:
            cand = outdir.parent / "eigen3d.json"
            js = cand if cand.exists() else None
        fc = scipy_freq_check(outdir, js)
        print(f"  C++ freq (Hz):   {[f'{x:.6e}' for x in fc['freq_cpp_Hz']]}")
        print(f"  SciPy freq (Hz): {[f'{x:.6e}' for x in fc['freq_scipy_Hz']]}")
        print(f"  matched max rel err: {fc['freq_max_rel_err']:.3e}")
        if fc["freq_max_rel_err"] < 1e-4:
            print("  => freq solver consistent with exported A/B [OK]")
        else:
            print("  => freq mismatch — check solver or matrix export [WARN]")

    if args.mat_outdir:
        print(f"\n=== 3. Compare MATLAB: {args.mat_outdir} ===")
        mc = compare_matlab(outdir, args.mat_outdir.resolve())
        if mc.get("skipped"):
            print(f"  skipped: {mc['skipped']}")
        else:
            print(f"  freq rel err: {mc['freq_rel_err']:.3e}")
            if "normE_rel_err_per_mode" in mc:
                for i, (rel, corr) in enumerate(
                    zip(mc["normE_rel_err_per_mode"], mc["normE_corr_per_mode"])
                ):
                    print(f"  normE mode {i}: rel={rel:.3e}, corr={corr:.6f}")

    print("\n--- How to get a MATLAB reference ---")
    print("  cd 三维matlab代码/matlab 3D一阶本征问题/3D一阶本征问题2")
    print("  % need SBCmesh.mat (from COMSOL/export); then run main.m")
    print("  % add after main.m loop:")
    print("  %   writematrix([solverff0.real, zeros(num,1)], 'OutFile/freq')")
    print("  %   for i=1:numberSolve, export one normE block with // separator")
    print("  Then: python tools/verify_eigen3d.py ... --mat-outdir path/to/matlab/OutFile")


if __name__ == "__main__":
    main()