#!/usr/bin/env python3
"""Compare two OpticsFEM COO exports (Ai/Aj/Av, Bv_real/Bv_imag)."""

from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np
from scipy import sparse


def load_coo(prefix: Path) -> tuple[sparse.csr_matrix, np.ndarray]:
    ai = np.loadtxt(prefix / "Ai.txt", dtype=np.int64)
    aj = np.loadtxt(prefix / "Aj.txt", dtype=np.int64)
    av: list[complex] = []
    for line in (prefix / "Av.txt").read_text(encoding="utf-8", errors="ignore").splitlines():
        line = line.strip()
        if not line:
            continue
        if line.startswith("(") and line.endswith(")"):
            re_s, im_s = line[1:-1].split(",")
            av.append(complex(float(re_s), float(im_s)))
        else:
            av.append(complex(float(line)))
    a = sparse.coo_matrix((np.asarray(av), (ai, aj))).tocsr()
    br = np.loadtxt(prefix / "Bv_real.txt")
    bi = np.loadtxt(prefix / "Bv_imag.txt")
    b = br + 1j * bi
    return a, b


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("ref", type=Path, help="reference directory (e.g. MATLAB OutFile_fem4_ab)")
    ap.add_argument("other", type=Path, help="other directory (e.g. C++ OutFile)")
    args = ap.parse_args()

    a_ref, b_ref = load_coo(args.ref)
    a_other, b_other = load_coo(args.other)

    print(f"ref:   n={a_ref.shape[0]} nnz={a_ref.nnz} |b|={np.linalg.norm(b_ref):.6g}")
    print(f"other: n={a_other.shape[0]} nnz={a_other.nnz} |b|={np.linalg.norm(b_other):.6g}")

    if a_ref.shape != a_other.shape:
        print("[FAIL] matrix shape mismatch")
        return 1
    if b_ref.size != b_other.size:
        print("[FAIL] b length mismatch")
        return 1

    d = a_ref - a_other
    max_abs = abs(d).max()
    nnz_diff = d.nnz
    rel = max_abs / max(abs(a_ref).max(), 1e-30)
    print(f"A max|diff|={max_abs:.6g}  rel={rel:.6g}  nnz(diff)={nnz_diff}")

    b_diff = np.max(np.abs(b_ref - b_other))
    print(f"b max|diff|={b_diff:.6g}")

    if max_abs < 1e-9 and b_diff < 1e-9:
        print("[OK] matrices match within 1e-9")
        return 0
    if max_abs < 1e-6 and b_diff < 1e-6:
        print("[WARN] close but not identical")
        return 0
    print("[FAIL] significant difference")
    return 2


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