#!/usr/bin/env python3
"""Replicate MATLAB FemMatrixAssembly (wave+port+PEC) and compare with reference."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import h5py
import numpy as np
from scipy import sparse
from scipy.sparse.linalg import spsolve

from compare_wave_asm import assemble_wave_matlab, bf_edge, build_edges
from export_ab_coo import export_ab_coo


def gauss_tri_order2() -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    xft = np.array([1.0 / 3, 0.6, 0.2, 0.2])
    yft = np.array([1.0 / 3, 0.2, 0.6, 0.2])
    pft = np.array([-0.28125, 25.0 / 96, 25.0 / 96, 25.0 / 96])
    return xft, yft, pft


def bf_et(i: int, u: float, v: float) -> np.ndarray:
    if i == 1:
        return np.array([1.0 - v, u, 0.0])
    if i == 2:
        return np.array([v, 1.0 - u, 0.0])
    if i == 3:
        return np.array([-v, u, 0.0])
    raise ValueError(i)


def bf_ez(i: int, u: float, v: float) -> np.ndarray:
    return np.array([0.0, 0.0, [1 - u - v, u, v][i - 1]])


def bf_curl_et(i: int, u: float, v: float) -> np.ndarray:
    return np.array([0.0, 0.0, [2.0, -2.0, 2.0][i - 1]])


def bf_curl_ez(i: int, u: float, v: float) -> np.ndarray:
    if i == 1:
        return np.array([-1.0, 1.0, 0.0])
    if i == 2:
        return np.array([0.0, -1.0, 0.0])
    if i == 3:
        return np.array([1.0, 0.0, 0.0])
    raise ValueError(i)


def get_tet_face(num_face: int) -> tuple[list[int], list[int]]:
    mapping = {
        1: ([0, 1, 2], [0, 1, 3]),
        2: ([0, 1, 3], [0, 2, 4]),
        3: ([0, 2, 3], [1, 2, 5]),
        4: ([1, 2, 3], [3, 4, 5]),
    }
    return mapping[num_face]


def assemble_port_face(
    nodes: np.ndarray,
    elements: np.ndarray,
    eoe: np.ndarray,
    mode: dict,
    face_idx: int,
    is_input: bool,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, complex, complex]:
    n_edges = int(eoe.max()) + 1
    fc = mode["facesConn"][face_idx]
    num_ele = int(fc[0]) - 1
    num_face = int(fc[1])
    face_node, face_bf = get_tet_face(num_face)
    idx = elements[num_ele] - 1
    x, y, z = nodes[idx, 0], nodes[idx, 1], nodes[idx, 2]
    pairs = [(0, 1), (0, 2), (0, 3), (1, 2), (1, 3), (2, 3)]
    l = np.array([np.linalg.norm(nodes[idx[a]] - nodes[idx[b]]) for a, b in pairs])

    mu = 1.0
    gamma = complex(mode["gamma"][0], mode["gamma"][1])
    power = mode["powerCoef"]
    normal = np.array(mode["normal"], dtype=float)
    et_coef = np.array(mode["Et_re"]) + 1j * np.array(mode["Et_im"])
    ez_coef = np.array(mode["Ez_re"]) + 1j * np.array(mode["Ez_im"])
    port_edge = np.array(mode["portEdgeOfFace"][face_idx], dtype=int)
    port_ez = np.array(mode["portNewFaces"][face_idx], dtype=int)

    jac = np.column_stack([x[:3] - x[3], y[:3] - y[3], z[:3] - z[3]])
    inv_jac = np.linalg.inv(jac)
    xx, yy, zz = x[face_node], y[face_node], z[face_node]

    f_jac = np.zeros((3, 3))
    f_jac[0, 0] = -xx[0] + xx[1]
    f_jac[0, 1] = -yy[0] + yy[1]
    f_jac[1, 0] = -xx[0] + xx[2]
    f_jac[1, 1] = -yy[0] + yy[2]
    f_jac[2, 2] = 1.0
    inv_f = np.linalg.inv(f_jac)
    f_jac_s = np.zeros((3, 3))
    inv22 = np.linalg.inv(f_jac[:2, :2])
    f_jac_s[0, 0] = inv22[1, 1]
    f_jac_s[0, 1] = -inv22[1, 0]
    f_jac_s[1, 0] = -inv22[0, 1]
    f_jac_s[1, 1] = inv22[0, 0]
    f_jac_s[2, 2] = 1.0
    f_t_jac = f_jac.T / np.linalg.det(f_jac)
    f_det = abs(np.linalg.det(f_jac))

    f_jac2 = np.zeros((2, 3))
    f_jac2[0, 0] = -xx[0] + xx[1]
    f_jac2[0, 1] = -yy[0] + yy[1]
    f_jac2[1, 0] = -xx[0] + xx[2]
    f_jac2[1, 1] = -yy[0] + yy[2]
    v_jac = np.linalg.inv(jac)

    xft, yft, pft = gauss_tri_order2()
    g = np.zeros(n_edges, dtype=complex)
    st = np.zeros(n_edges, dtype=complex)
    b_edge = np.zeros(n_edges, dtype=complex)
    p = 0.0 + 0.0j
    b1 = 0.0 + 0.0j

    for gp in range(len(pft)):
        f_e = np.zeros(3, dtype=complex)
        f_curl_e = np.zeros(3, dtype=complex)
        for i in range(3):
            et = inv_f @ bf_et(i + 1, xft[gp], yft[gp]) * l[face_bf[i]]
            ez = bf_ez(i + 1, xft[gp], yft[gp])
            curl_et = f_t_jac @ bf_curl_et(i + 1, xft[gp], yft[gp]) * l[face_bf[i]]
            curl_ez = f_jac_s @ bf_curl_ez(i + 1, xft[gp], yft[gp])
            f_e += et * et_coef[port_edge[i]] + ez * ez_coef[port_ez[i]]
            f_curl_e += curl_et * et_coef[port_edge[i]] + curl_ez * ez_coef[port_ez[i]]

        rot = np.array([f_e[1], -f_e[0], 0.0])
        if is_input:
            f_n0 = f_e * power
            f_curl_n0 = (f_curl_e + gamma * rot) * power
            f_n1 = f_e
            f_curl_n1 = f_curl_e - gamma * rot
        else:
            f_n2 = f_e
            f_curl_n2 = f_curl_e + gamma * rot
        f_w = f_e

        ref = np.array([xft[gp], yft[gp]]) @ f_jac2 + np.array([xx[0], yy[0], zz[0]])
        ref3 = ref @ v_jac - np.array([x[3], y[3], z[3]]) @ v_jac

        if is_input:
            for i in range(3):
                bf = inv_jac @ bf_edge(face_bf[i] + 1, ref3[0], ref3[1], ref3[2]) * l[face_bf[i]]
                edge_id = int(eoe[num_ele, face_bf[i]])
                g[edge_id] += pft[gp] * f_det * (f_w[0] * bf[0] + f_w[1] * bf[1])
                temp = mu * np.cross(normal, f_curl_n1)
                st[edge_id] += pft[gp] * f_det * np.vdot(bf, temp)
                temp0 = mu * np.cross(normal, f_curl_n0)
                b_edge[edge_id] += pft[gp] * f_det * np.vdot(bf, temp0)
            p += pft[gp] * f_det * (f_w[0] * f_n1[0] + f_w[1] * f_n1[1])
            b1 += pft[gp] * f_det * (f_w[0] * f_n0[0] + f_w[1] * f_n0[1])
        else:
            for i in range(3):
                bf = inv_jac @ bf_edge(face_bf[i] + 1, ref3[0], ref3[1], ref3[2]) * l[face_bf[i]]
                edge_id = int(eoe[num_ele, face_bf[i]])
                g[edge_id] += pft[gp] * f_det * (f_w[0] * bf[0] + f_w[1] * bf[1])
                temp = mu * np.cross(normal, f_curl_n2)
                st[edge_id] += pft[gp] * f_det * np.vdot(bf, temp)
            p += pft[gp] * f_det * (f_w[0] * f_n2[0] + f_w[1] * f_n2[1])

    return st, g, b_edge, p, b1


from scipy.sparse import bmat


def block_append_input(a: sparse.csr_matrix, b: np.ndarray, s: np.ndarray, g: np.ndarray, p: complex, b1: complex):
    n = b.size
    s_col = sparse.csr_matrix(s.reshape(-1, 1))
    g_row = sparse.csr_matrix(g.reshape(1, -1))
    p_corner = sparse.csr_matrix(np.array([[-p]]))
    a2 = bmat([[a, s_col], [g_row, p_corner]], format="csr")
    b2 = np.concatenate([-b, [b1]])
    return a2, b2


def block_append_output(a: sparse.csr_matrix, b: np.ndarray, t: np.ndarray, g: np.ndarray, p: complex):
    n = b.size
    t_col = sparse.csr_matrix(t.reshape(-1, 1))
    g_row = sparse.csr_matrix(g.reshape(1, -1))
    p_corner = sparse.csr_matrix(np.array([[-p]]))
    a2 = bmat([[a, t_col], [g_row, p_corner]], format="csr")
    b2 = np.concatenate([b, [0.0]])
    return a2, b2


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument(
        "--export",
        type=Path,
        default=None,
        help="export reduced A/b (COO) to this directory",
    )
    args = ap.parse_args()

    root = Path(__file__).resolve().parents[2]
    rel = root / "3D opticsfem-master/port/Release"
    with h5py.File(next(root.rglob("MeshData2x.mat")), "r") as f:
        nodes = np.asarray(f["Mesh/Nodes"]).T
        elements = np.asarray(f["Mesh/Elements"]).T.astype(int)
        faces = np.asarray(f["Mesh/Faces"]).T.astype(int)
        faces_index = np.asarray(f["Mesh/FacesIndex"]).reshape(-1).astype(int)
        domains = np.asarray(f["Mesh/Domains"]).reshape(-1).astype(int)

    modes = json.loads(rel.joinpath("port_modes_fem4.json").read_text())
    _, eoe = build_edges(elements)

    a = assemble_wave_matlab(nodes, elements, domains)
    b = np.zeros(a.shape[0], dtype=complex)

    n_edge = int(eoe.max()) + 1

    for mode, is_in in [(modes["input"], True), (modes["output"], False)]:
        st = np.zeros(a.shape[0], dtype=complex)
        g = np.zeros(a.shape[0], dtype=complex)
        bedge = np.zeros(n_edge, dtype=complex)
        p = 0.0 + 0.0j
        b1 = 0.0 + 0.0j
        for fi in range(len(mode["facesConn"])):
            s1, g1, be1, p1, b1e = assemble_port_face(nodes, elements, eoe, mode, fi, is_in)
            st[:n_edge] += s1
            g[:n_edge] += g1
            if is_in:
                bedge += be1
                b1 += b1e
            p += p1
        if is_in:
            b[:n_edge] += bedge
            a, b = block_append_input(a, b, st[: a.shape[0]], g[: a.shape[0]], p, b1)
            print(f"after input |b|={np.linalg.norm(b):.6g} b1={b1:.6g}")
        else:
            a, b = block_append_output(a, b, st[: a.shape[0]], g[: a.shape[0]], p)
            print(f"after output |b|={np.linalg.norm(b):.6g}")

    pec_faces = faces[np.isin(faces_index, [1, 2, 5, 24])]
    pec_pairs = set()
    for tri in pec_faces:
        for a0, b0 in [(tri[0], tri[1]), (tri[1], tri[2]), (tri[2], tri[0])]:
            pec_pairs.add(tuple(sorted((int(a0), int(b0)))))
    edge_u, _ = build_edges(elements)
    pec_idx = {i for i, e in enumerate(edge_u) if tuple(e) in pec_pairs}
    free = [i for i in range(a.shape[0]) if i not in pec_idx]
    a = a.tocsr()[free, :][:, free]
    b = b[free]

    export_dir = args.export
    if export_dir is None:
        export_dir = root / "三维matlab代码/2023-2-端口激励问题(四面体网格)/OutFile_fem4_ab"
    export_ab_coo(a, b, export_dir)

    x = spsolve(a, b)
    pcoef = modes["input"]["powerCoef"]
    e1, e2 = x[-2], x[-1]
    print(f"Python S11={e1/pcoef} S21={e2/pcoef}")
    ref = next(root.rglob("OutFile_fem4/S_params.txt"))
    vals = ref.read_text().strip().splitlines()[1].split()
    print(f"MATLAB S11={complex(float(vals[0]), float(vals[1]))} S21={complex(float(vals[2]), float(vals[3]))}")
    cpp = rel / "OutFile_fem4" / "S_params.txt"
    print("C++", cpp.read_text().strip())
    return 0


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