#!/usr/bin/env python3
"""Compare PBC-projected A matrix: MATLAB vs C++, and verify PBC projection pipeline."""
import re
import math
from pathlib import Path
import numpy as np
from scipy import sparse

ROOT = Path(__file__).resolve().parents[1]
MATLAB_DIR = ROOT.parent / "三维matlab代码" / "matlab 3D一阶基+散射边界条件+单周期边界"
CPP_OUT = ROOT / "build" / "Release" / "OutFile"
MAT_OUT = MATLAB_DIR / "OutFile"
MAT_ASM = MATLAB_DIR / "OutFile_asm"
MESH_FILE = MATLAB_DIR / "PBCmesh.dat"
PBC_ANGLE = math.pi / 3
PBC_PHI = complex(0.5, 0.8660254037844386)


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

    def read_val(p):
        vals = []
        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, dtype=np.complex128)

    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 read_mesh_dat(path: Path):
    """Minimal PBCmesh.dat reader (0-based indices)."""
    lines = path.read_text(encoding="utf-8", errors="ignore").splitlines()
    i = 0

    def read_block(name):
        nonlocal i
        while i < len(lines) and lines[i].strip() != name:
            i += 1
        if i >= len(lines):
            raise KeyError(name)
        i += 1
        return i

    def read_int():
        nonlocal i
        v = int(lines[i].strip())
        i += 1
        return v

    def read_floats(n):
        nonlocal i
        vals = list(map(float, lines[i].split()))
        i += 1
        return np.array(vals[:n])

    read_block("NbrVertex")
    n_vertex = read_int()
    read_block("Vertex")
    vertex = np.zeros((n_vertex, 3))
    for r in range(n_vertex):
        vertex[r] = read_floats(3)

    read_block("NbrEdge")
    n_edge = read_int()
    read_block("Edge")
    edge = np.zeros((n_edge, 2), dtype=np.int64)
    for r in range(n_edge):
        a, b = map(int, lines[i].split())
        edge[r] = (a - 1, b - 1)
        i += 1

    read_block("NbrTet")
    n_tet = read_int()
    read_block("Tet")
    tet = np.zeros((n_tet, 4), dtype=np.int64)
    for r in range(n_tet):
        tet[r] = np.array(list(map(int, lines[i].split()))) - 1
        i += 1

    read_block("EdgeOfTet")
    edge_of_tet = np.zeros((n_tet, 6), dtype=np.int64)
    for r in range(n_tet):
        edge_of_tet[r] = np.array(list(map(int, lines[i].split()))) - 1
        i += 1

    read_block("NbrTri")
    n_tri = read_int()
    read_block("ConnOfTri")
    conn_of_tri = np.zeros((n_tri, 2), dtype=np.int64)
    for r in range(n_tri):
        a, b = map(int, lines[i].split())
        conn_of_tri[r] = (a - 1, b - 1)
        i += 1

    read_block("DomainOfTri")
    domain_of_tri = np.zeros(n_tri, dtype=np.int64)
    for r in range(n_tri):
        domain_of_tri[r] = int(lines[i].strip())  # keep 1-based like MATLAB mesh
        i += 1

    return {
        "vertex": vertex,
        "edge": edge,
        "tet": tet,
        "edge_of_tet": edge_of_tet,
        "conn_of_tri": conn_of_tri,
        "domain_of_tri": domain_of_tri,
        "n_edge": n_edge,
    }


def find_tri(domains, mesh):
    domains = np.atleast_1d(domains)
    out = []
    for j in range(mesh["conn_of_tri"].shape[0]):
        num_tet, _ = mesh["conn_of_tri"][j]
        dom = mesh["domain_of_tri"][j]
        # DomainOfTri in dat is 1-based after +1 in C++ loader; MATLAB c6 may differ.
        # PBCmesh.dat stores 0-based domain ids before C++ +1.
        if dom in domains or (dom + 1) in domains:
            out.append(j)
    return np.array(out, dtype=np.int64)


def face_edges(mesh, tri_idx):
    num_tet, num_face = mesh["conn_of_tri"][tri_idx]
    e = mesh["edge_of_tet"][num_tet]
    face = num_face + 1
    if face == 1:
        return [e[0], e[1], e[3]]
    if face == 2:
        return [e[0], e[2], e[4]]
    if face == 3:
        return [e[1], e[2], e[5]]
    if face == 4:
        return [e[3], e[4], e[5]]
    raise ValueError(face)


def find_pbc_index(src, dst, theta, mesh):
    src_tris = find_tri(src, mesh)
    dst_tris = find_tri(dst, mesh)
    src_edges, dst_edges = [], []
    for tri in src_tris:
        src_edges.extend(face_edges(mesh, tri))
    for tri in dst_tris:
        dst_edges.extend(face_edges(mesh, tri))
    src_edges = np.unique(src_edges)
    dst_edges = np.unique(dst_edges)

    tra = np.array(
        [[math.cos(theta), -math.sin(theta), 0.0],
         [math.sin(theta), math.cos(theta), 0.0],
         [0.0, 0.0, 1.0]]
    )
    tol = 0.01 * 0.00005
    pairs = []
    for si in src_edges:
        v1 = mesh["vertex"][mesh["edge"][si, 0]]
        v2 = mesh["vertex"][mesh["edge"][si, 1]]
        matched = False
        for dj, di in enumerate(dst_edges):
            v3 = tra @ mesh["vertex"][mesh["edge"][di, 0]]
            v4 = tra @ mesh["vertex"][mesh["edge"][di, 1]]
            l1 = np.linalg.norm(v1 - v3)
            l2 = np.linalg.norm(v2 - v4)
            l3 = np.linalg.norm(v1 - v4)
            l4 = np.linalg.norm(v2 - v3)
            if l1 + l2 < tol:
                pairs.append((si, di, 1))
                matched = True
                break
            if l3 + l4 < tol:
                pairs.append((si, di, -1))
                matched = True
                break
        if not matched:
            pairs.append((si, -1, 0))
    pairs = [(s, d, sg) for s, d, sg in pairs if s != d and d >= 0]
    return pairs


def build_P(dof, pairs):
    P = sparse.lil_matrix((dof, dof), dtype=np.complex128)
    for i in range(dof):
        P[i, i] = 1.0
    dst_cols = []
    for src, dst, sign in pairs:
        P[dst, src] = sign * PBC_PHI
        dst_cols.append(dst)
    keep = [c for c in range(dof) if c not in dst_cols]
    return P[:, keep].tocsr()


def diff_stats(A, B, label):
    D = A - B
    print(f"\n=== {label} ===")
    print(f"  max |diff|      = {np.max(np.abs(D.data)) if D.nnz else 0:.6g}")
    print(f"  max |real diff| = {np.max(np.abs(D.real.data)) if D.nnz else 0:.6g}")
    print(f"  max |imag diff| = {np.max(np.abs(D.imag.data)) if D.nnz else 0:.6g}")
    if D.nnz and A.shape[0] == B.shape[0]:
        di = np.abs(A.diagonal().imag - B.diagonal().imag)
        print(f"  diag imag max diff = {di.max():.6g}, >10 count = {(di > 10).sum()}")


def main():
    print("Loading matrices...")
    A_cpp = load_coo(CPP_OUT)
    A_mat = load_coo(MAT_OUT)
    A_asm = load_coo(MAT_ASM)
    diff_stats(A_cpp, A_mat, "C++ OutFile vs MATLAB OutFile (after PBC)")

    print("\nLoading mesh and building PBC pairs...")
    mesh = read_mesh_dat(MESH_FILE)
    # domains in case_config: src=[1,4], dst=[2,5] (1-based in MATLAB)
    pairs = find_pbc_index(np.array([1, 4]), np.array([2, 5]), PBC_ANGLE, mesh)
    print(f"  PBC pairs: {len(pairs)}")
    signs = [p[2] for p in pairs]
    print(f"  sign +1: {signs.count(1)}, sign -1: {signs.count(-1)}")

    P = build_P(mesh["n_edge"], pairs)
    A_proj = P.conj().T @ A_asm @ P
    diff_stats(A_proj, A_mat, "MATLAB asm projected with Python P vs MATLAB OutFile")
    diff_stats(A_proj, A_cpp, "MATLAB asm projected with Python P vs C++ OutFile")

    # Test sign-flip hypothesis on all -1 pairs
    pairs_flip = [(s, d, -sg if sg < 0 else sg) for s, d, sg in pairs]
    Pf = build_P(mesh["n_edge"], pairs_flip)
    Af = Pf.conj().T @ A_asm @ Pf
    diff_stats(Af, A_cpp, "MATLAB asm with flipped -1 signs vs C++ OutFile")


if __name__ == "__main__":
    main()