#!/usr/bin/env python3
"""Diagnose why ref A + cpp b works but cpp A + cpp b fails despite tiny A diff."""

from __future__ import annotations

from pathlib import Path

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

ROOT = Path(__file__).resolve().parents[2]
REF = ROOT / "三维matlab代码/2023-2-端口激励问题(四面体网格)/OutFile_fem4_ab"
CPP = ROOT / "3D opticsfem-master/port/Release/OutFile_fem4"
PCOEF = 2.69467866496258e15


def parse_av_line(s: str) -> complex:
    s = s.strip()
    if s.startswith("("):
        body = s[1:-1]
        a, b = body.split(",", 1)
        return complex(float(a), float(b))
    return complex(float(s))


def load_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 = [parse_av_line(line) for line in (prefix / "Av.txt").read_text(encoding="utf-8", errors="ignore").splitlines() if line.strip()]
    return sparse.coo_matrix((av, (ai, aj)), shape=(79616, 79616)).tocsr()


def load_b(prefix: Path) -> np.ndarray:
    return np.loadtxt(prefix / "Bv_real.txt") + 1j * np.loadtxt(prefix / "Bv_imag.txt")


def main() -> int:
    print("=== Load matrices ===")
    A_ref = load_coo(REF)
    A_cpp = load_coo(CPP)
    b_ref = load_b(REF)
    b_cpp = load_b(CPP)

    print(f"A_ref nnz={A_ref.nnz}, A_cpp nnz={A_cpp.nnz}")
    print(f"indices equal: {np.array_equal(A_ref.indices, A_cpp.indices) and np.array_equal(A_ref.indptr, A_cpp.indptr)}")
    ddata = np.abs(A_ref.data - A_cpp.data)
    print(f"A.data max|diff|={ddata.max():.6g}, mean|diff|={ddata.mean():.6g}")
    print(f"A.data rel max|diff|={ddata.max()/max(np.abs(A_ref.data).max(),1e-30):.6g}")
    print(f"b max|diff|={np.abs(b_ref-b_cpp).max():.6g}, count>1e-6={(np.abs(b_ref-b_cpp)>1e-6).sum()}")

    # duplicate (i,j) in raw COO before CSR merge?
    ai = np.loadtxt(REF / "Ai.txt", dtype=np.int64)
    aj = np.loadtxt(REF / "Aj.txt", dtype=np.int64)
    keys = np.column_stack([ai, aj])
    uniq, counts = np.unique(keys, axis=0, return_counts=True)
    ndup = int((counts > 1).sum())
    print(f"raw COO duplicate (i,j) groups: {ndup}")

    print("\n=== CRITICAL: same b_cpp, two A ===")
    x_refA = spsolve(A_ref, b_cpp)
    x_cppA = spsolve(A_cpp, b_cpp)
    rel = np.linalg.norm(x_refA - x_cppA) / np.linalg.norm(x_refA)
    print(f"||x_refA - x_cppA|| / ||x_refA|| = {rel:.6g}")
    print(f"x_refA  S11={x_refA[-2]/PCOEF}  S21={x_refA[-1]/PCOEF}")
    print(f"x_cppA  S11={x_cppA[-2]/PCOEF}  S21={x_cppA[-1]/PCOEF}")

    print("\n=== Same A_ref, two b ===")
    x1 = spsolve(A_ref, b_ref)
    x2 = spsolve(A_ref, b_cpp)
    print(f"||x(b_ref)-x(b_cpp)||/||x(b_ref)|| = {np.linalg.norm(x1-x2)/np.linalg.norm(x1):.6g}")
    print(f"x(b_ref) S11={x1[-2]/PCOEF}")
    print(f"x(b_cpp) S11={x2[-2]/PCOEF}")

    print("\n=== C++ exported X ===")
    x_exp = np.loadtxt(CPP / "X_real.txt") + 1j * np.loadtxt(CPP / "X_imag.txt")
    free_len = 79616
    x_red = x_exp  # reduced? export is full 83350
    # map: C++ exports full _mX; reduced solve is 79616
    print(f"X size={x_exp.size}, e1={x_exp[83348]/PCOEF}, e2={x_exp[83349]/PCOEF}")
    print(f"||x_cppA - x_exp[?]|| : reduced x vs full X needs free map")

    # Where do A values differ most (relative)?
    mask = ddata > 0
    if mask.any():
        rel_d = ddata[mask] / np.maximum(np.abs(A_ref.data[mask]), 1e-30)
        top = np.argsort(-rel_d)[:10]
        print("\nTop relative A diffs (CSR data index):")
        for i in top:
            print(f"  idx={i} rel={rel_d[i]:.3g} abs={ddata[i]:.3g} ref={A_ref.data[i]:.6g} cpp={A_cpp.data[i]:.6g}")

    # Port columns 79614,79615
    for col in [79614, 79615]:
        dcol = (A_ref[:, col] - A_cpp[:, col]).toarray().ravel()
        nz = np.abs(dcol) > 0
        print(f"\nPort col {col}: nnz diff={nz.sum()}, max|diff|={np.abs(dcol).max():.3g}")

    return 0


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