#!/usr/bin/env python3
"""Compare C++ wave-only A with Python replication of MATLAB PhysicMatrixAssembly."""

from __future__ import annotations

import re
from pathlib import Path

import h5py
import numpy as np
from scipy import sparse

PI = np.pi


def bf_edge(num: int, u: float, v: float, w: float) -> np.ndarray:
    out = np.zeros(3)
    if num == 1:
        out[:] = [-v, u, 0.0]
    elif num == 2:
        out[:] = [-w, 0.0, u]
    elif num == 3:
        out[:] = [-1 + v + w, -u, -u]
    elif num == 4:
        out[:] = [0.0, -w, v]
    elif num == 5:
        out[:] = [-v, -1 + u + w, -v]
    elif num == 6:
        out[:] = [-w, -w, -1 + u + v]
    else:
        raise ValueError(num)
    return out


def bf_curl_edge(num: int, u: float, v: float, w: float) -> np.ndarray:
    out = np.zeros(3)
    if num == 1:
        out[2] = 2.0
    elif num == 2:
        out[1] = -2.0
    elif num == 3:
        out[1] = 2.0
        out[2] = -2.0
    elif num == 4:
        out[0] = 2.0
    elif num == 5:
        out[0] = -2.0
        out[2] = 2.0
    elif num == 6:
        out[0] = 2.0
        out[1] = -2.0
    else:
        raise ValueError(num)
    return out


def gauss_tet_order2() -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    xt = np.array([0.25, 0.166666666667, 0.166666666667, 0.166666666667, 0.5])
    yt = np.array([0.25, 0.166666666667, 0.166666666667, 0.5, 0.166666666667])
    zt = np.array([0.25, 0.166666666667, 0.5, 0.166666666667, 0.166666666667])
    pt = np.array([-0.133333333333, 0.075, 0.075, 0.075, 0.075])
    return xt, yt, zt, pt


def build_edges(elements: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    el2no = elements.T
    n1 = el2no[[0, 0, 0, 1, 1, 2], :]
    n2 = el2no[[1, 2, 3, 2, 3, 3], :]
    el_ed = np.column_stack([n1.reshape(-1, order="F"), n2.reshape(-1, order="F")])
    edge, ic = np.unique(el_ed, axis=0, return_inverse=True)
    eoe = ic.reshape(6, -1, order="F").T
    return edge, eoe


def assemble_wave_matlab(
    nodes: np.ndarray, elements: np.ndarray, domains: np.ndarray
) -> sparse.csr_matrix:
    lam0 = 1.55e-6
    k0 = 2 * PI / lam0
    epsilonr = np.array([4.0, 11.9, 11.9, 11.9], dtype=complex)
    sigma = np.array([0.0, 0.0, 5000.0, 0.0])
    mur = np.ones(4)
    temp = 1 / k0 * 120 * PI
    epsilon = epsilonr - 1j * sigma * temp

    _, eoe = build_edges(elements)
    xt, yt, zt, pt = gauss_tet_order2()
    trips: list[tuple[int, int, complex]] = []

    for n in range(elements.shape[0]):
        idx = elements[n] - 1
        x = nodes[idx, 0]
        y = nodes[idx, 1]
        z = nodes[idx, 2]
        l = np.ones(6)
        pairs = [(0, 1), (0, 2), (0, 3), (1, 2), (1, 3), (2, 3)]
        for i, (a, b) in enumerate(pairs):
            l[i] = np.linalg.norm(nodes[idx[a]] - nodes[idx[b]])
        jac = np.column_stack([x[:3] - x[3], y[:3] - y[3], z[:3] - z[3]])
        inv_jac = np.linalg.inv(jac)
        t_jac = jac.T / np.linalg.det(jac)
        det_jac = abs(np.linalg.det(jac))
        domain = int(domains[n]) - 1
        mu = 1.0 / mur[domain]
        eps = epsilon[domain]

        ne = np.zeros((3, 6, len(xt)), dtype=complex)
        curl_ne = np.zeros((3, 6, len(xt)), dtype=complex)
        for i in range(6):
            for k in range(len(xt)):
                temp_bf = bf_edge(i + 1, xt[k], yt[k], zt[k])
                ne[:, i, k] = inv_jac @ temp_bf * l[i]
                temp_c = bf_curl_edge(i + 1, xt[k], yt[k], zt[k])
                curl_ne[:, i, k] = t_jac @ temp_c * l[i]

        ae = np.zeros((6, 6), dtype=complex)
        for i in range(6):
            for j in range(6):
                for k in range(len(pt)):
                    ae[i, j] += pt[k] * det_jac * (
                        mu * np.vdot(curl_ne[:, i, k], curl_ne[:, j, k])
                        - k0 * k0 * eps * np.vdot(ne[:, i, k], ne[:, j, k])
                    )
        for i in range(6):
            for j in range(6):
                ii = int(eoe[n, i])
                jj = int(eoe[n, j])
                trips.append((ii, jj, ae[i, j]))

    ai = np.array([t[0] for t in trips], dtype=np.int64)
    aj = np.array([t[1] for t in trips], dtype=np.int64)
    av = np.array([t[2] for t in trips])
    n = int(max(ai.max(), aj.max()) + 1)
    return sparse.coo_matrix((av, (ai, aj)), shape=(n, n)).tocsr()


def load_cpp_wave(cpp_asm: Path) -> sparse.csr_matrix:
    ai = np.loadtxt(cpp_asm / "Ai.txt", dtype=np.int64)
    aj = np.loadtxt(cpp_asm / "Aj.txt", dtype=np.int64)
    av = []
    for line in (cpp_asm / "Av.txt").read_text().splitlines():
        line = line.strip()
        if not line or line.startswith("//"):
            continue
        m = re.match(r"\(([^,]+),([^)]+)\)", line)
        av.append(complex(float(m.group(1)), float(m.group(2))))
    av = np.array(av)
    n = int(max(ai.max(), aj.max()) + 1)
    return sparse.coo_matrix((av, (ai, aj)), shape=(n, n)).tocsr()


def main() -> int:
    root = Path(__file__).resolve().parents[2]
    mat_path = next(root.rglob("MeshData2x.mat"))
    cpp_asm = root / "3D opticsfem-master/port/Release/OutFile_asm"

    with h5py.File(mat_path, "r") as f:
        nodes = np.asarray(f["Mesh/Nodes"]).T
        elements = np.asarray(f["Mesh/Elements"]).T.astype(int)
        domains = np.asarray(f["Mesh/Domains"]).reshape(-1).astype(int)

    print("Assembling MATLAB-style wave matrix ...")
    a_py = assemble_wave_matlab(nodes, elements, domains)
    a_cpp = load_cpp_wave(cpp_asm)
    print(f"Python nnz={a_py.nnz}  C++ nnz={a_cpp.nnz}")

    diff = a_py - a_cpp
    print(f"max |A_py - A_cpp| = {max(abs(diff.data).max() if diff.nnz else 0.0, 0.0):.6e}")
    rel = diff.data / (a_cpp.data + 1e-30)
    if diff.nnz:
        print(f"median rel diff on overlapping nnz = {np.median(np.abs(rel)):.6e}")
    return 0


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