#!/usr/bin/env python3
"""Diagnose double-PBC A matrix diff: MATLAB vs C++ projection."""
from __future__ import annotations

import math
import re
from collections import Counter, defaultdict
from pathlib import Path

import numpy as np
import scipy.io as sio
from scipy import sparse

ROOT = Path(__file__).resolve().parents[1]
MAT_ROOT = ROOT.parent / "三维matlab代码" / "matlab 3D一阶基+散射边界条件+周期边界"
CPP_OUT = ROOT / "build" / "Release" / "OutFile_double"
MAT_OUT = MAT_ROOT / "OutFile_double"
MAT_ASM = MAT_ROOT / "OutFile_double_asm"
MESH_MAT = MAT_ROOT / "doublePBC_mesh.mat"


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 = []
    with open(prefix / "Av.txt", 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)
            av.append(complex(float(m.group(1)), float(m.group(2))))
    n = int(max(ai.max(), aj.max()) + 1)
    return sparse.csr_matrix((av, (ai, aj)), shape=(n, n))


def find_tri(domains, mesh) -> np.ndarray:
    domains = np.atleast_1d(domains)
    dom_tri = np.asarray(mesh.DomainOfTri).flatten()
    conn = np.asarray(mesh.ConnOfTri, dtype=int)
    out = []
    for j, d in enumerate(dom_tri):
        if d in domains:
            out.append(j)
    return np.array(out, dtype=int)


def face_edges(mesh, tri_idx: int) -> list[int]:
    num_tet, num_face = np.asarray(mesh.ConnOfTri[tri_idx], dtype=int) - 1
    e = np.asarray(mesh.EdgeOfTet[num_tet], dtype=int).flatten() - 1
    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, dis, mesh) -> np.ndarray:
    """MATLAB findPBCIndex (1-based edge ids in output)."""
    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)

    dis = np.asarray(dis, dtype=float).reshape(3)
    dl = np.linalg.norm(dis)
    err = dl * 0.00005
    vertex = np.asarray(mesh.Vertex, dtype=float)
    edge = np.asarray(mesh.Edge, dtype=int) - 1

    pairs = []
    for si in src_edges:
        v1 = vertex[edge[si, 0]]
        v2 = vertex[edge[si, 1]]
        matched = False
        for dj in dst_edges:
            v3 = vertex[edge[dj, 0]]
            v4 = vertex[edge[dj, 1]]
            l1 = abs(np.linalg.norm(v1 - v3) - dl)
            l2 = abs(np.linalg.norm(v2 - v4) - dl)
            l3 = abs(np.linalg.norm(v1 - v4) - dl)
            l4 = abs(np.linalg.norm(v2 - v3) - dl)
            if l1 + l2 < err:
                pairs.append((si + 1, dj + 1, 1))
                matched = True
                break
            if l3 + l4 < err:
                pairs.append((si + 1, dj + 1, -1))
                matched = True
                break
        if not matched:
            pairs.append((si + 1, si + 1, 1))
    return np.array(pairs, dtype=int)


def tabulate_duplicate_rows(values: np.ndarray) -> np.ndarray:
    """Return row indices where tabulate count > 1 (MATLAB assembly_pbc_double bug)."""
    vals = np.sort(values)
    uniq, counts = np.unique(vals, return_counts=True)
    rows = np.where(counts > 1)[0] + 1  # 1-based row index in tabulate table
    return rows


def duplicate_values(values: np.ndarray) -> np.ndarray:
    c = Counter(values.tolist())
    return np.array(sorted([v for v, n in c.items() if n > 1]), dtype=int)


def merge_pbc_double_matlab_fixed(pbc1: np.ndarray, pbc2: np.ndarray, phi1: float, phi2: float):
    """Intended MATLAB merge: ovDstIndex = tbl2(find(tbl2(:,2)>1), 1)."""
    return merge_pbc_double_correct(pbc1, pbc2, phi1, phi2)


def merge_pbc_double_correct(pbc1: np.ndarray, pbc2: np.ndarray, phi1: float, phi2: float):
    """C++ mergeDoublePbcPairs logic (duplicate dst values)."""
    ov = duplicate_values(np.concatenate([pbc1[:, 1], pbc2[:, 1]]))
    remove1 = np.zeros(len(pbc1), dtype=bool)
    remove2 = np.zeros(len(pbc2), dtype=bool)
    merged = []

    dst1 = pbc1[:, 1]
    src1 = pbc1[:, 0]
    dst2 = pbc2[:, 1]
    src2 = pbc2[:, 0]

    for ov_dst in ov:
        i1 = np.where(dst2 == ov_dst)[0]
        if len(i1) == 0:
            continue
        i1 = i1[0]
        s2 = src2[i1]
        i3 = np.where(dst1 == s2)[0]
        if len(i3) == 0:
            continue
        i3 = i3[0]
        s1 = src1[i3]
        sign = pbc2[i1, 2] * pbc1[i3, 2]
        merged.append((s1, ov_dst, sign, phi1 * phi2 * sign))
        remove1[i3] = True
        remove2[i1] = True
        remove1[dst1 == ov_dst] = True
        remove2[dst2 == ov_dst] = True

    for i in range(len(pbc1)):
        if not remove1[i]:
            merged.append((pbc1[i, 0], pbc1[i, 1], pbc1[i, 2], phi1 * pbc1[i, 2]))
    for i in range(len(pbc2)):
        if not remove2[i]:
            merged.append((pbc2[i, 0], pbc2[i, 1], pbc2[i, 2], phi2 * pbc2[i, 2]))

    arr = np.array([(m[0], m[1], m[2]) for m in merged], dtype=int)
    phi = np.array([m[3] for m in merged], dtype=float)
    return arr, phi


def build_p(dof: int, pbc: np.ndarray, phi: np.ndarray) -> sparse.csr_matrix:
    rows, cols, data = [], [], []
    for i in range(dof):
        rows.append(i)
        cols.append(i)
        data.append(1.0)
    for k in range(len(pbc)):
        src = int(pbc[k, 0]) - 1
        dst = int(pbc[k, 1]) - 1
        rows.append(dst)
        cols.append(src)
        data.append(phi[k])
    p_full = sparse.csr_matrix((data, (rows, cols)), shape=(dof, dof))
    dst_cols = sorted({int(pbc[k, 1]) - 1 for k in range(len(pbc))})
    keep = np.ones(dof, dtype=bool)
    keep[dst_cols] = False
    return p_full[:, keep]


def diff_stats(a: sparse.csr_matrix, b: sparse.csr_matrix, label: str) -> None:
    if a.shape != b.shape:
        print(f"{label}: shape mismatch {a.shape} vs {b.shape}")
        return
    d = a - b
    mx = float(np.max(np.abs(d.data))) if d.nnz else 0.0
    rel = mx / max(float(np.max(np.abs(b.data))), 1e-30)
    print(f"{label}: max|diff|={mx:.6g}, rel_max={rel:.6g}, diff_nnz={d.nnz}")


def main() -> None:
    mesh = sio.loadmat(MESH_MAT, squeeze_me=True, struct_as_record=False)["mesh"]
    dof = int(mesh.NbrEdge)

    pbc1 = find_pbc_index([1, 4, 7, 10, 13], [54, 55, 56, 57, 58], [2e-6, 0, 0], mesh)
    pbc2 = find_pbc_index([2, 5, 8, 11, 14], [17, 18, 19, 20, 21], [0, 2e-6, 0], mesh)
    print(f"PBC1 pairs: {len(pbc1)}, PBC2 pairs: {len(pbc2)}")

    ov_vals = duplicate_values(np.concatenate([pbc1[:, 1], pbc2[:, 1]]))
    print(f"Corner overlap dst count (correct): {len(ov_vals)}")

    pbc_ok, phi_ok = merge_pbc_double_correct(pbc1, pbc2, -1.0, 1.0)
    pbc_fix, phi_fix = merge_pbc_double_matlab_fixed(pbc1, pbc2, -1.0, 1.0)
    print(f"Merged constraints: correct={len(pbc_ok)}")

    a_asm = load_coo(MAT_ASM)
    a_mat = load_coo(MAT_OUT)
    a_cpp = load_coo(CPP_OUT)

    p_ok = build_p(dof, pbc_ok, phi_ok)
    a_proj_ok = p_ok.conj().T @ a_asm @ p_ok

    diff_stats(a_mat, a_cpp, "MAT out vs C++ out")
    diff_stats(a_mat, a_proj_ok, "MAT out vs asm+correct P")
    diff_stats(a_cpp, a_proj_ok, "C++ out vs asm+correct P")

    # normE quick
    ne_cpp = np.loadtxt(CPP_OUT / "normE", comments="//")
    ne_mat = np.loadtxt(MAT_OUT / "normE", comments="//")
    rel = np.linalg.norm(ne_cpp - ne_mat) / np.linalg.norm(ne_mat)
    print(f"normE L2 rel diff: {rel:.6g}, corr={np.corrcoef(ne_cpp, ne_mat)[0,1]:.6f}")


if __name__ == "__main__":
    main()