"""Check the GPU near-field weights against brute-force integration in NumPy.

For a few (target, patch) pairs we compute the reference moments
I_n = int K_m(x_t, y) P_n |J| du dv with independent quadratures: for the self
patch, polar rules on the triangles fanned from the target to the edges, with
a sinh substitution along each edge and very many nodes; for neighbouring patches, a uniform quadtree of depth 6 with an
8x8 Gauss-Legendre rule per leaf. The reference weights are U^T I,
computed with the same geometry model (exact or interpolated) as the GPU code.
"""

import os
import sys
import time

import cupy as cp
import numpy as np

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), ".."))

from fmm3dbie_gpu import cuda_ops as ops  # noqa: E402
from fmm3dbie_gpu import surface as S  # noqa: E402

OVER4PI = 1 / (4 * np.pi)


def kernels6(xt, nt, y, ns, zk):
    """Six kernels (6, n) for target xt (3,), nt (3,), sources y (3, n), ns (3, n)."""
    d = xt[:, None] - y
    r2 = np.sum(d * d, axis=0)
    r = np.sqrt(r2)
    kabs = abs(zk)
    dnt = np.sum(d * nt[:, None], axis=0)
    dns = np.sum(d * ns, axis=0)
    nn = np.sum(ns * nt[:, None], axis=0)
    Ek = np.exp(1j * zk * r) * OVER4PI / r
    Ec = np.exp(-kabs * r) * OVER4PI / r
    fk, fc = 1 - 1j * zk * r, 1 + kabs * r
    gk = 3 - 3j * zk * r - zk**2 * r2
    gc = 3 + 3 * kabs * r + kabs**2 * r2
    K = np.zeros((6, y.shape[1]), dtype=complex)
    K[4] = Ek
    K[1] = Ec
    K[0] = -dnt * fk * Ek / r2
    K[2] = -dnt * fc * Ec / r2
    K[5] = dns * fk * Ek / r2
    K[3] = (nn * fk / r2 - dns * dnt * gk / r2**2) * Ek - (nn * fc / r2 - dns * dnt * gc / r2**2) * Ec
    return K


def moments_at(surf, patch, uv, w, xt, nt, zk):
    """sum_nodes K_m P_n |J| w -> (6, npols)."""
    g = surf.eval_geometry(patch, uv)
    ns = np.cross(g[3:6].T, g[6:9].T).T
    jac = np.linalg.norm(ns, axis=0)
    ns = ns / jac
    K = kernels6(xt, nt, g[0:3], ns, zk)
    P = S.basis_2d(uv, surf.norder)
    return (K * (jac * w)[None, :]) @ P.T


def ref_self(surf, patch, u0, v0, xt, nt, zk, nx=400, nr=80):
    """Polar rules on the triangles (target, edge point) with many nodes. Along each
    edge we split at the foot of the perpendicular and use x = |h| sinh(xi), which
    resolves the peak of width |h| next to the foot."""
    xr, wr = np.polynomial.legendre.leggauss(nr)
    xr, wr = 0.5 * (xr + 1), 0.5 * wr
    xx, wx = np.polynomial.legendre.leggauss(nx)
    I = np.zeros((6, surf.npols), dtype=complex)
    for e in range(4):
        h = [1 - u0, -1 - u0, 1 - v0, -1 - v0][e]
        if abs(h) < 1e-14:
            continue
        if e < 2:
            x1, x2 = -1 - v0, 1 - v0
        else:
            x1, x2 = -1 - u0, 1 - u0
        for xend in (x1, x2):
            if abs(xend) < 1e-14:
                continue
            xi1 = np.arcsinh(abs(xend) / abs(h))
            xi = 0.5 * xi1 * (xx + 1)
            x = np.sign(xend) * abs(h) * np.sinh(xi)
            dx = abs(h) * np.cosh(xi) * 0.5 * xi1 * wx
            R, X = np.meshgrid(xr, x, indexing="ij")
            WW = np.outer(wr, dx) * abs(h) * R
            if e < 2:
                uv = np.vstack([(u0 + R * h).ravel(), (v0 + R * X).ravel()])
            else:
                uv = np.vstack([(u0 + R * X).ravel(), (v0 + R * h).ravel()])
            I += moments_at(surf, patch, uv, WW.ravel(), xt, nt, zk)
    return I


def ref_leaf(surf, patch, xt, nt, zk, depth=6, nq=8):
    gx, gw = np.polynomial.legendre.leggauss(nq)
    nl = 2**depth
    h = 1.0 / nl
    c = -1 + h + 2 * h * np.arange(nl)
    U, V, A, B = np.meshgrid(c, c, gx, gx, indexing="ij")
    uv = np.vstack([(U + h * A).ravel(), (V + h * B).ravel()])
    w = (h * h * np.outer(gw, gw)[None, None]).repeat(nl, 0).repeat(nl, 1).ravel()
    return moments_at(surf, patch, uv, w, xt, nt, zk)


def main():
    """Usage: python3 test_quadrature.py [norder] [k] [key=value ...]

    Keys: nq, nqr, nqx, sub_len, osc_len (quadrature), m (patches per face edge),
    tp (target patch), rfac, rfac_q, kR_max, exact (1 for the exact geometry, 0 for
    the interpolated one)."""
    norder = int(sys.argv[1]) if len(sys.argv) > 1 else 6
    zk = float(sys.argv[2]) if len(sys.argv) > 2 else 10.0
    quad = dict(nq=8, nqr=14, nqx=8, sub_len=1.0, osc_len=2.0, m=3, tp=4, rfac=2.0, rfac_q=2.0,
                kR_max=3.0, exact=1)
    for kv in sys.argv[3:]:
        key, val = kv.split("=")
        quad[key] = type(quad[key])(val)
    m, tp, rfac, rfac_q, kR_max, exact = (quad.pop(x) for x in ["m", "tp", "rfac", "rfac_q", "kR_max", "exact"])
    print("k", zk, "norder", norder, quad, "m", m, "target patch", tp, "exact geometry", bool(exact))
    surf = S.rounded_cube(a=1.0, p=8, m=m, norder=norder, exact=bool(exact))
    npols = surf.npols
    # targets: three nodes of patch tp (corner, centre, next to an edge); all their near pairs
    node_ids = [tp * npols, tp * npols + npols // 2, tp * npols + npols - 2]
    xt_all = surf.srcvals[0:3, node_ids]
    pairs = []
    for q, i in enumerate(node_ids):
        d = np.linalg.norm(surf.centroids - xt_all[:, q:q + 1], axis=0)
        pairs += [(i, j) for j in np.nonzero(d <= rfac * surf.rads)[0]]
    pair_t = cp.asarray([i for i, _ in pairs], dtype=cp.int32)
    pair_p = cp.asarray([j for _, j in pairs], dtype=cp.int32)

    geo = ops.Geometry(surf)
    x = cp.asarray(np.ascontiguousarray(surf.srcvals[0:3].T))
    n = cp.asarray(np.ascontiguousarray(surf.srcvals[9:12].T))
    pid = cp.asarray(surf.patch_id.astype(np.int32))
    uv = cp.asarray(np.ascontiguousarray(surf.uv_nodes.T))
    umat = cp.asarray(np.ascontiguousarray(surf.umat))

    t0 = time.time()
    leaf_ptr, leaves = ops.build_leaves(pair_t, pair_p, x, pid, geo, zk=zk, rfac_q=rfac_q,
                                        nlmax=14, kR_max=kR_max)
    W = ops.near_weights(pair_t, pair_p, x, n, pid, uv, geo, leaf_ptr, leaves, umat, zk, **quad)
    cp.cuda.Device().synchronize()
    print(f"GPU weights in {time.time() - t0:.2f} s for {len(pairs)} pairs")
    W = W.get()
    leaf_ptr = leaf_ptr.get()
    worst = np.zeros(6)
    for q, (i, j) in enumerate(pairs):
        xt, nt = surf.srcvals[0:3, i], surf.srcvals[9:12, i]
        if j == surf.patch_id[i]:
            I = ref_self(surf, j, *surf.uv_nodes[:, i], xt, nt, zk)
            kind = "self"
        else:
            I = ref_leaf(surf, j, xt, nt, zk)
            kind = "near"
        Wref = I @ surf.umat  # w_j = sum_n U[n, j] I_n
        err = np.linalg.norm(W[q] - Wref, axis=1) / np.linalg.norm(Wref, axis=1)
        d = np.linalg.norm(xt - surf.centroids[:, j]) / surf.rads[j]
        worst = np.maximum(worst, err)
        if err.max() > 1e-6 or kind == "self":
            print(f"target {i} patch {j} ({kind}, dist/rad {d:.2f}, {leaf_ptr[q + 1] - leaf_ptr[q]} leaves): "
                  "rel err per kernel " + " ".join(f"{e:.1e}" for e in err))
    print("worst per kernel:", " ".join(f"{e:.1e}" for e in worst))


if __name__ == "__main__":
    main()