"""MATLAB assembly_equ/out/inc 忠实复现，用于与 C++ 对比定位问题。"""
import numpy as np
from pathlib import Path
from scipy import sparse
from scipy.sparse.linalg import spsolve

PI = np.pi

def get_bf(typ, num, u, v, w):
    if typ == 1:
        table = {
            1: np.array([-v, u, 0.0]),
            2: np.array([-w, 0.0, u]),
            3: np.array([-1 + v + w, -u, -u]),
            4: np.array([0.0, -w, v]),
            5: np.array([-v, -1 + u + w, -v]),
            6: np.array([-w, -w, -1 + u + v]),
        }
        return table[num]
    table = {
        1: np.array([0.0, 0.0, 2.0]),
        2: np.array([0.0, -2.0, 0.0]),
        3: np.array([0.0, 2.0, -2.0]),
        4: np.array([2.0, 0.0, 0.0]),
        5: np.array([-2.0, 0.0, 2.0]),
        6: np.array([2.0, -2.0, 0.0]),
    }
    return table[num]

def gp_tet():
    u = np.array([0.25, 0.166666666667, 0.166666666667, 0.166666666667, 0.5])
    v = np.array([0.25, 0.166666666667, 0.166666666667, 0.5, 0.166666666667])
    w = np.array([0.25, 0.166666666667, 0.5, 0.166666666667, 0.166666666667])
    wg = np.array([-0.133333333333, 0.075, 0.075, 0.075, 0.075])
    return u, v, w, wg

def gp_tri():
    u = np.array([0.333333333333333, 0.6, 0.2, 0.2])
    v = np.array([0.3333333333333333, 0.2, 0.6, 0.2])
    wg = np.array([-0.28125, 0.260416666666, 0.260416666666, 0.260416666666])
    return u, v, wg

def cross_n(n, a):
    return np.cross(n, np.cross(a, n))

def load_mesh(path):
    lines = Path(path).read_text().splitlines()
    i = 0
    def need(t):
        nonlocal i
        assert lines[i] == t, (lines[i], t)
        i += 1
    need("NbrVertex"); nv = int(lines[i]); i += 1
    need("Vertex")
    V = np.array([list(map(float, lines[i + k].split())) for k in range(nv)]); i += nv
    need("NbrTet"); nt = int(lines[i]); i += 1
    need("Tet")
    T = np.array([list(map(int, lines[i + k].split())) for k in range(nt)], int); i += nt
    need("DomainOfTet")
    domT = np.array([int(lines[i + k]) for k in range(nt)]); i += nt
    need("NbrEdge"); ne = int(lines[i]); i += 1
    need("Edge"); i += ne
    need("EdgeOfTet")
    EOT = np.array([list(map(int, lines[i + k].split())) for k in range(nt)], int); i += nt
    need("NbrTri"); ntri = int(lines[i]); i += 1
    need("Tri"); i += ntri
    need("DomainOfTri")
    domTri = np.array([int(lines[i + k]) for k in range(ntri)]); i += ntri
    need("ConnOfTri")
    conn = np.array([list(map(int, lines[i + k].split())) for k in range(ntri)], int); i += ntri
    need("NormOfFace"); nn = int(lines[i]); i += 1
    norm = {}
    for _ in range(nn):
        parts = list(map(float, lines[i].split())); i += 1
        d = int(parts[0]); norm[d] = np.array(parts[1:4])
    norm.setdefault(4, np.array([0.0, 0.0, 1.0]))
    return dict(V=V, T=T, domT=domT, EOT=EOT, domTri=domTri, conn=conn, norm=norm, ne=ne, nt=nt, ntri=ntri)

def find_tri(domains, domTri):
    out = []
    for d in domains:
        out.extend(np.where(domTri == d)[0].tolist())
    return np.array(sorted(set(out)))

def assemble_equ(mesh, lda0, eps, use_inv_jac=True):
    k0 = 2 * PI / lda0
    u, v, w, wg = gp_tet()
    rows, cols, vals = [], [], []
    V, T, domT, EOT = mesh["V"], mesh["T"], mesh["domT"], mesh["EOT"]
    for n in range(mesh["nt"]):
        x = V[T[n] - 1, 0]; y = V[T[n] - 1, 1]; z = V[T[n] - 1, 2]
        Jac = np.column_stack([x[:3] - x[3], y[:3] - y[3], z[:3] - z[3]])
        detJ = np.linalg.det(Jac)
        DetJac = abs(detJ)
        InvJac = np.linalg.inv(Jac)
        TJac = Jac.T / detJ  # curl map unchanged
        ep = eps[domT[n] - 1]
        Ae = np.zeros((6, 6), complex)
        for k in range(len(u)):
            E = np.zeros((3, 6)); curlE = np.zeros((3, 6))
            for j in range(6):
                ref = get_bf(1, j + 1, u[k], v[k], w[k])
                E[:, j] = (InvJac if use_inv_jac else Jac) @ ref
                curlE[:, j] = TJac @ get_bf(2, j + 1, u[k], v[k], w[k])
            for i in range(6):
                for j in range(6):
                    Ae[i, j] += wg[k] * DetJac * np.dot(curlE[:, i], curlE[:, j]) - wg[k] * DetJac * k0 * k0 * ep * np.dot(E[:, i], E[:, j])
        for i in range(6):
            for j in range(6):
                rows.append(EOT[n, i] - 1); cols.append(EOT[n, j] - 1); vals.append(Ae[i, j])
    return rows, cols, vals

def face_map(num_face):
    if num_face == 1: return np.array([1,0,0]), np.array([0,1,0]), np.array([0,0,1]), [1,2,4], [0,1,3]
    if num_face == 2: return np.array([1,0,0]), np.array([0,1,0]), np.array([0,0,0]), [1,3,5], [0,2,4]
    if num_face == 3: return np.array([1,0,0]), np.array([0,0,0]), np.array([0,1,0]), [2,3,6], [1,2,5]
    return np.array([0,0,0]), np.array([1,0,0]), np.array([0,1,0]), [4,5,6], [3,4,5]

def assemble_sbc(mesh, tri_idx, eps, lda0, Einc, is_inc, use_inv_jac=True):
    k0 = 2 * PI / lda0
    u, v, wg = gp_tri()
    rows, cols, vals, b = [], [], [], {}
    V, T, EOT, conn, norm = mesh["V"], mesh["T"], mesh["EOT"], mesh["conn"], mesh["norm"]
    for tri in tri_idx:
        domain = mesh["domTri"][tri]
        num_tet = conn[tri, 0] - 1
        num_face = conn[tri, 1]
        x = V[T[num_tet] - 1, 0]; y = V[T[num_tet] - 1, 1]; z = V[T[num_tet] - 1, 2]
        x2, y2, z2, bf_idx, edge_slot = face_map(num_face)
        if num_face == 1: x3, y3, z3 = x[:3], y[:3], z[:3]
        elif num_face == 2: x3, y3, z3 = x[[0,1,3]], y[[0,1,3]], z[[0,1,3]]
        elif num_face == 3: x3, y3, z3 = x[[0,2,3]], y[[0,2,3]], z[[0,2,3]]
        else: x3, y3, z3 = x[1:], y[1:], z[1:]
        Jac = np.column_stack([x[:3] - x[3], y[:3] - y[3], z[:3] - z[3]])
        InvJac = np.linalg.inv(Jac)
        a = np.linalg.norm(np.array([x3[0]-x3[1], y3[0]-y3[1], z3[0]-z3[1]]))
        b1 = np.linalg.norm(np.array([x3[0]-x3[2], y3[0]-y3[2], z3[0]-z3[2]]))
        c = np.linalg.norm(np.array([x3[1]-x3[2], y3[1]-y3[2], z3[1]-z3[2]]))
        integ = 0.25 * np.sqrt((a+b1+c)*(a+b1-c)*(a-b1+c)*(b1+c-a))
        normal = norm[domain]
        ep = eps[mesh["domT"][num_tet] - 1]
        nn = np.sqrt(ep)
        map_idx = [EOT[num_tet, s] - 1 for s in edge_slot]
        Ae = np.zeros((3, 3), complex)
        Be = np.zeros(3, complex)
        for k in range(len(u)):
            wgp = 1 - u[k] - v[k]
            u2 = x2 @ [u[k], v[k], wgp]; v2 = y2 @ [u[k], v[k], wgp]; w2 = z2 @ [u[k], v[k], wgp]
            Egp = np.zeros((3, 3))
            for j in range(3):
                ref = get_bf(1, bf_idx[j], u2, v2, w2)
                Egp[:, j] = (InvJac if use_inv_jac else Jac) @ ref
            for i in range(3):
                for j in range(3):
                    tj = cross_n(normal, Egp[:, j])
                    Ae[i, j] += 1j * k0 * nn * integ * wg[k] * np.dot(Egp[:, i], tj) * 2
            if is_inc:
                t_inc = cross_n(normal, Einc)
                for i in range(3):
                    Be[i] -= 1j * k0 * nn * 2 * integ * wg[k] * np.dot(Egp[:, i], t_inc) * 2
        for i in range(3):
            for j in range(3):
                rows.append(map_idx[i]); cols.append(map_idx[j]); vals.append(Ae[i, j])
            if is_inc:
                b[map_idx[i]] = b.get(map_idx[i], 0) + Be[i]
    return rows, cols, vals, (b if is_inc else {})

def solve_case(mesh, lda0, eps, out_dom, inc_dom, Einc, use_inv_jac, label):
    r, c, v = assemble_equ(mesh, lda0, eps, use_inv_jac)
    ro, co, vo, _ = assemble_sbc(mesh, find_tri(out_dom, mesh["domTri"]), eps, lda0, Einc, False, use_inv_jac)
    ri, ci, vi, bi = assemble_sbc(mesh, find_tri(inc_dom, mesh["domTri"]), eps, lda0, Einc, True, use_inv_jac)
    rows = r + ro + ri; cols = c + co + ci; vals = v + vo + vi
    n = mesh["ne"]
    A = sparse.coo_matrix((vals, (rows, cols)), shape=(n, n)).tocsr()
    b = np.zeros(n, complex)
    for k, val in bi.items(): b[k] = val
    x = spsolve(A, b)
    print(f"=== {label} ===")
    print(f"nnz={A.nnz} |b|={np.linalg.norm(b):.6f} |x|={np.linalg.norm(x):.6f}")
    return A, b, x

def compare_cpp(cpp_dir, A_ref, b_ref):
    coo = Path(cpp_dir) / "coo.txt"
    if not coo.exists():
        print("cpp dump not found:", coo)
        return
    hdr = coo.read_text().splitlines()[0].split()
    nnz, n = int(hdr[0]), int(hdr[1])
    rows, cols, re, im = [], [], [], []
    for line in coo.read_text().splitlines()[1:]:
        p = line.split(); rows.append(int(p[0])); cols.append(int(p[1])); re.append(float(p[2])); im.append(float(p[3]))
    A_cpp = sparse.coo_matrix((np.array(re)+1j*np.array(im), (rows, cols)), shape=(n, n)).tocsr()
    b_cpp = np.zeros(n, complex)
    for i, line in enumerate((Path(cpp_dir)/"b.txt").read_text().splitlines()):
        p = line.split(); b_cpp[i] = float(p[0])+1j*float(p[1])
    dA = A_cpp - A_ref
    print("=== C++ vs ref (InvJac) ===")
    print(f"max|dA|={np.max(np.abs(dA.data)) if dA.nnz else 0:.6e} rel|db|={np.linalg.norm(b_cpp-b_ref)/(np.linalg.norm(b_ref)+1e-30):.6e}")
    print(f"cpp |b|={np.linalg.norm(b_cpp):.6f} solve|x|={np.linalg.norm(spsolve(A_cpp,b_cpp)):.6f}")

def main():
    mesh_path = Path(r"C:\Users\Administrator\Desktop\西电-合作\三维matlab代码\matlab 3D一阶散射问题\SBCmesh.dat")
    mesh = load_mesh(mesh_path)
    lda0, eps = 0.8, np.array([1.0, 1.5])
    out_dom = [1, 2, 3, 5, 14]; inc_dom = [4]
    Einc = np.array([1.0, 0.0, 0.0])
    A1, b1, x1 = solve_case(mesh, lda0, eps, out_dom, inc_dom, Einc, True, "InvJac (MATLAB Jac\\)")
    A2, b2, x2 = solve_case(mesh, lda0, eps, out_dom, inc_dom, Einc, False, "Jac@ (Piola forward)")
    compare_cpp(Path(r"C:\Users\Administrator\Desktop\西电-合作\3D opticsfem-master\build\Release\cpp_dump"), A1, b1)

if __name__ == "__main__":
    main()