#!/usr/bin/env python3
"""Project MATLAB OutFile_asm with PBC and compare to OutFile / C++."""

from __future__ import annotations

import math
import re
from pathlib import Path

import numpy as np
from scipy import sparse

from compare_pbc_pairs import collect_edges, load_mesh, pair_matlab, remove_self


def load_coo(prefix: Path) -> sparse.csr_matrix:
    def read_col(p: Path) -> np.ndarray:
        with open(p, encoding="utf-8", errors="ignore") as f:
            return np.array([int(x.strip()) for x in f if x.strip()])

    def read_val(p: Path) -> np.ndarray:
        vals: list[complex] = []
        with open(p, encoding="utf-8", errors="ignore") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                m = re.match(r"\(([-+0-9.eE]+),([-+0-9.eE]+)\)", line)
                if m:
                    vals.append(complex(float(m.group(1)), float(m.group(2))))
                else:
                    vals.append(complex(float(line), 0.0))
        return np.array(vals)

    ai = read_col(prefix / "Ai.txt")
    aj = read_col(prefix / "Aj.txt")
    av = read_val(prefix / "Av.txt")
    n = int(max(ai.max(), aj.max()) + 1)
    return sparse.csr_matrix((av, (ai, aj)), shape=(n, n))


def build_p_matrix(dof: int, pairs: list[tuple[int, int, int]], phi: complex) -> sparse.csr_matrix:
    # 1-based edge indices in mesh -> use as matrix indices directly (MATLAB convention)
    rows: list[int] = []
    cols: list[int] = []
    data: list[complex] = []
    for i in range(1, dof + 1):
        rows.append(i)
        cols.append(i)
        data.append(1.0)
    for src, dst, sign in pairs:
        rows.append(dst)
        cols.append(src)
        data.append(sign * phi)
    p_full = sparse.csr_matrix((data, (np.array(rows) - 1, np.array(cols) - 1)), shape=(dof, dof))
    dst_cols = sorted({dst - 1 for _, dst, _ in pairs})
    keep = np.ones(dof, dtype=bool)
    keep[dst_cols] = False
    return p_full[:, keep]


def main() -> None:
    root = Path(__file__).resolve().parents[2]
    mat_dir = root / "三维matlab代码" / "matlab 3D一阶基+散射边界条件+单周期边界"
    cpp_dir = root / "3D opticsfem-master" / "build" / "Release" / "OutFile"
    mesh = load_mesh(root / "3D opticsfem-master" / "build" / "Release" / "PBCmesh.dat")

    phi = complex(0.5, math.sqrt(3) / 2)
    theta = math.pi / 3
    src_edges = collect_edges(mesh, [1, 4])
    dst_edges = collect_edges(mesh, [2, 5])
    pairs = remove_self(pair_matlab(mesh, src_edges, dst_edges, theta))

    a_asm = load_coo(mat_dir / "OutFile_asm")
    a_out = load_coo(mat_dir / "OutFile")
    a_cpp = load_coo(cpp_dir)

    dof = a_asm.shape[0]
    p = build_p_matrix(dof, pairs, phi)
    a_proj = p.conj().T @ a_asm @ p

    d_self = a_proj - a_out
    d_cpp = a_cpp - a_out
    print("Projected asm vs mat OutFile:")
    print("  max |diff|", np.max(np.abs(d_self.data)))
    print("  max |real diff|", np.max(np.abs(d_self.real.data)))
    print("  max |imag diff|", np.max(np.abs(d_self.imag.data)))

    print("\nC++ vs mat OutFile:")
    print("  max |imag diff|", np.max(np.abs(d_cpp.imag.data)))

    i = 20552 - 1  # 0-based for scipy
    print(f"\nDOF 20552:")
    print("  mat out ", a_out[i, i])
    print("  proj asm", a_proj[i, i])
    print("  cpp out ", a_cpp[i, i])
    print("  asm pre ", a_asm[i, i])

    # If projected asm matches mat out, pre-PBC asm from C++ is wrong
    # Compare asm sizes / check if cpp pre-projection differs
    r, c, v = sparse.find(d_self)
    order = np.argsort(-np.abs(v.imag))
    print("\nTop projected-asm vs mat-out imag diffs:")
    for k in order[:5]:
        ri, ci = int(r[k]), int(c[k])
        print(f"  ({ri+1},{ci+1}) proj={a_proj[ri,ci]:.6g} mat={a_out[ri,ci]:.6g}")


if __name__ == "__main__":
    main()