"""Green's identity S_k[du/dn] - D_k[u] = u/2 at boundary nodes, brute force over all patches,
with three models of the patch geometry used inside the integrals:
  interp:     xyz, r_u, r_v each interpolated from node values (fmm3dbie's model)
  consistent: xyz interpolated, r_u and r_v the derivatives of that interpolant
  exact:      the analytic superellipsoid parametrization
The density (u and du/dn) is always the polynomial interpolant of node values."""
import os
import sys

import numpy as np

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

from fmm3dbie_gpu import surface as S  # noqa: E402
from problems.common import green, green_dn  # noqa: E402
from test_quadrature import kernels6  # noqa: E402

a_, p_, m_, norder = 1.0, 8, 4, 6
surf = S.rounded_cube(a=a_, p=p_, m=m_, norder=norder)
br = S.graded_breaks(m_, 0.5)
npols = surf.npols
k = 10.0; xe = np.array([1.7, 0.6, -0.5])
u = green(surf.srcvals[0:3], xe, k); un = green_dn(surf.srcvals[0:3], surf.srcvals[9:12], xe, k)

def exact_geom(patch, uv):
    f, rem = divmod(patch, m_ * m_); i, j = divmod(rem, m_)
    n, e1, e2 = (np.array(v, float) for v in S._FACES[f])
    s0, s1, t0, t1 = br[i], br[i + 1], br[j], br[j + 1]
    s = s0 + 0.5 * (uv[0] + 1) * (s1 - s0); t = t0 + 0.5 * (uv[1] + 1) * (t1 - t0)
    c = n[:, None] + s * e1[:, None] + t * e2[:, None]
    N = np.sum(c ** p_, axis=0) ** (1.0 / p_); gradN = c ** (p_ - 1) / N ** (p_ - 1)
    def deriv(dc, scale):
        dN = np.sum(gradN * dc[:, None], axis=0)
        return a_ * scale * (dc[:, None] / N - c * dN / N ** 2)
    return np.vstack([a_ * c / N, deriv(e1, 0.5 * (s1 - s0)), deriv(e2, 0.5 * (t1 - t0))])

def dlegendre_vander(x, n):
    P = S.legendre_vander(x, n); D = np.zeros_like(P)
    for k_ in range(1, n + 1):  # P_k' = sum over j = k-1, k-3, ... of (2j+1) P_j
        for j in range(k_ - 1, -1, -2):
            D[k_] += (2 * j + 1) * P[j]
    return D

def consistent_geom(patch, uv):
    c = surf.srccoefs[0:3, patch * npols:(patch + 1) * npols]
    Pu, Pv = S.legendre_vander(uv[0], norder), S.legendre_vander(uv[1], norder)
    Du, Dv = dlegendre_vander(uv[0], norder), dlegendre_vander(uv[1], norder)
    B = (Pv[:, None] * Pu[None]).reshape(npols, -1)
    Bu = (Pv[:, None] * Du[None]).reshape(npols, -1)
    Bv = (Dv[:, None] * Pu[None]).reshape(npols, -1)
    return np.vstack([c @ B, c @ Bu, c @ Bv])

models = {"interp": lambda P, uv: surf.eval_geometry(P, uv), "consistent": consistent_geom, "exact": exact_geom}

def moments(geom, patch, uv, w, xt, nt):
    g = geom(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, k)
    return (K * (jac * w)[None]) @ S.basis_2d(uv, norder).T

def self_rule(u0, v0, nx=200, nr=40):
    xr, wr = np.polynomial.legendre.leggauss(nr); xr, wr = 0.5 * (xr + 1), 0.5 * wr
    xx, wx = np.polynomial.legendre.leggauss(nx)
    U, W = [], []
    for e in range(4):
        h = [1 - u0, -1 - u0, 1 - v0, -1 - v0][e]
        x1, x2 = (-1 - v0, 1 - v0) if e < 2 else (-1 - u0, 1 - u0)
        for xa, xb in [(x1, 0.0), (0.0, x2)]:
            x = xa + 0.5 * (xb - xa) * (xx + 1); dx = 0.5 * (xb - xa) * wx
            R, X = np.meshgrid(xr, x, indexing="ij"); WW = np.outer(wr, dx) * abs(h) * R
            uv = np.vstack([(u0 + R * h).ravel(), (v0 + R * X).ravel()]) if e < 2 else \
                 np.vstack([(u0 + R * X).ravel(), (v0 + R * h).ravel()])
            U.append(uv); W.append(WW.ravel())
    return np.hstack(U), np.hstack(W)

def leaf_rule(depth, nq=8):
    gx, gw = np.polynomial.legendre.leggauss(nq); nl = 2 ** depth; h = 1.0 / nl
    c = -1 + h + 2 * h * np.arange(nl)
    Uc, Vc, A, B = np.meshgrid(c, c, gx, gx, indexing="ij")
    uv = np.vstack([(Uc + h * A).ravel(), (Vc + h * B).ravel()])
    w = (h * h * np.outer(gw, gw)[None, None]).repeat(nl, 0).repeat(nl, 1).ravel()
    return uv, w

leaf6, leaf3 = leaf_rule(6), leaf_rule(3)
P = 13
for loc in [0, 3, 21, 24]:
    i = P * npols + loc
    xt, nt = surf.srcvals[0:3, i], surf.srcvals[9:12, i]
    out = []
    for name, geom in models.items():
        Sv = Dv = 0
        for j in range(surf.npatches):
            if j == P:
                uv, w = self_rule(*surf.uv_nodes[:, i])
            else:
                d = np.linalg.norm(xt - surf.centroids[:, j]) / surf.rads[j]
                uv, w = leaf6 if d < 3 else leaf3
            Wt = moments(geom, j, uv, w, xt, nt) @ surf.umat
            sl = slice(j * npols, (j + 1) * npols)
            Sv += Wt[4] @ un[sl]; Dv += Wt[5] @ u[sl]
        out.append(f"{name} {abs(Sv - Dv - u[i] / 2) / abs(u[i]):.1e}")
    print(f"node position {loc} (u,v)=({surf.uv_nodes[0, i]:+.2f},{surf.uv_nodes[1, i]:+.2f}): relative residual  " + "   ".join(out), flush=True)