"""Quadrilateral patch surfaces in the fmm3dbie convention, without the library.

Each patch is the image of the reference square [-1, 1]^2. It carries
(norder+1)^2 tensor Gauss-Legendre nodes, with the u index varying fastest,
and an unnormalized Legendre basis P_j(u) P_i(v) with the same ordering.
Node data (srcvals) is xyz, dx/du, dx/dv, unit normal, and the geometry is
also stored as basis coefficients (srccoefs) so that it can be evaluated at
any point of the reference square. These are the same conventions as
fmm3dbie's iptype 11, and were checked against the library.
"""

import numpy as np


def gauss_legendre(n):
    x, w = np.polynomial.legendre.leggauss(n)
    return x, w


def legendre_vander(x, norder):
    """P_0(x), ..., P_norder(x) as an array of shape (norder+1, len(x))."""
    x = np.asarray(x, dtype=float)
    P = np.zeros((norder + 1, x.size))
    P[0] = 1.0
    if norder >= 1:
        P[1] = x
    for n in range(1, norder):
        P[n + 1] = ((2 * n + 1) * x * P[n] - n * P[n - 1]) / (n + 1)
    return P


def disc_nodes(norder):
    """Reference nodes (2, npols) and weights (npols) of the tensor GL rule."""
    x, w = gauss_legendre(norder + 1)
    v, u = np.meshgrid(x, x, indexing="ij")  # u fastest
    wv, wu = np.meshgrid(w, w, indexing="ij")
    return np.vstack([u.ravel(), v.ravel()]), (wu * wv).ravel()


def basis_2d(uv, norder):
    """Basis values (npols, npts): index n = i*(norder+1) + j is P_j(u) P_i(v)."""
    Pu = legendre_vander(uv[0], norder)
    Pv = legendre_vander(uv[1], norder)
    return (Pv[:, None, :] * Pu[None, :, :]).reshape((norder + 1) ** 2, -1)


def vals_to_coefs_matrix(norder):
    """U with coefs = U @ vals for values at the tensor GL nodes."""
    x, w = gauss_legendre(norder + 1)
    P = legendre_vander(x, norder)
    U1 = (2 * np.arange(norder + 1)[:, None] + 1) / 2 * P * w[None, :]
    return np.kron(U1, U1)  # (i, j) x (ii, jj) with i, ii slow


def oversample_matrix(norder, nover):
    """M with values at the order-nover nodes = M @ values at the order-norder nodes."""
    uv_o, _ = disc_nodes(nover)
    return basis_2d(uv_o, norder).T @ vals_to_coefs_matrix(norder)


class Surface:
    """exact: optional dict(params=(npatches, 13), a=, p=) giving the superellipsoid face
    chart of every patch (n, e1, e2, s0, s1, t0, t1); when present the integrals use
    the exact geometry instead of the polynomial interpolant."""

    def __init__(self, norder, srcvals, exact=None):
        self.exact = exact
        self.norder = norder
        self.iptype = 11
        self.srcvals = np.ascontiguousarray(srcvals, dtype=float)  # (12, npts)
        self.npts = self.srcvals.shape[1]
        self.npols = (norder + 1) ** 2
        assert self.npts % self.npols == 0
        self.npatches = self.npts // self.npols
        self.uvs, self.wref = disc_nodes(norder)
        self.umat = vals_to_coefs_matrix(norder)
        vals = self.srcvals[:9].reshape(9, self.npatches, self.npols)
        self.srccoefs = np.einsum("nj,cpj->cpn", self.umat, vals).reshape(9, self.npts)
        ru, rv = self.srcvals[3:6], self.srcvals[6:9]
        self.jac = np.linalg.norm(np.cross(ru.T, rv.T).T, axis=0)
        self.wts = self.jac * np.tile(self.wref, self.npatches)
        # centroid and radius of each patch from a 5x5 sample grid (fmm3dbie uses the
        # four corners, which underestimates the extent of strongly curved patches)
        g5 = np.linspace(-1, 1, 5)
        corners = np.array(np.meshgrid(g5, g5, indexing="ij")).reshape(2, -1)
        Pc = basis_2d(corners, norder)  # (npols, 25)
        xc = np.einsum("cpn,nk->cpk", self.srccoefs[:3].reshape(3, self.npatches, self.npols), Pc)
        self.centroids = xc.mean(axis=2)  # (3, npatches)
        self.rads = np.linalg.norm(xc - self.centroids[:, :, None], axis=0).max(axis=1)
        self.patch_id = np.repeat(np.arange(self.npatches), self.npols)
        self.uv_nodes = np.tile(self.uvs, self.npatches)  # (2, npts)

    def eval_geometry(self, patch, uv):
        """xyz, r_u, r_v at reference points uv (2, n) of one patch: (9, n)."""
        if self.exact is not None:
            e = self.exact
            return superellipsoid_chart(e["params"][patch], e["a"], e["p"], uv)
        P = basis_2d(uv, self.norder)
        c = self.srccoefs[:, patch * self.npols:(patch + 1) * self.npols]
        return c @ P

    def oversampled(self, nover):
        """Nodes of the order-nover rule on every patch: xyz, normal, weights."""
        if nover == self.norder:
            return self.srcvals[0:3], self.srcvals[9:12], self.wts
        uv_o, w_o = disc_nodes(nover)
        if self.exact is not None:
            g = np.hstack([self.eval_geometry(j, uv_o) for j in range(self.npatches)])
        else:
            P = basis_2d(uv_o, self.norder)  # (npols, npols_o)
            g = np.einsum("cpn,nk->cpk", self.srccoefs.reshape(9, self.npatches, self.npols), P)
            g = g.reshape(9, -1)
        nrm = np.cross(g[3:6].T, g[6:9].T).T
        jac = np.linalg.norm(nrm, axis=0)
        nrm /= jac
        return g[0:3], nrm, jac * np.tile(w_o, self.npatches)


# ---------------------------------------------------------------------------
# Rounded cube (superellipsoid), the same construction as ../rounded_cube.py

_FACES = [
    ((1, 0, 0), (0, 1, 0), (0, 0, 1)),
    ((-1, 0, 0), (0, 0, 1), (0, 1, 0)),
    ((0, 1, 0), (0, 0, 1), (1, 0, 0)),
    ((0, -1, 0), (1, 0, 0), (0, 0, 1)),
    ((0, 0, 1), (1, 0, 0), (0, 1, 0)),
    ((0, 0, -1), (0, 1, 0), (1, 0, 0)),
]


def graded_breaks(m, grading):
    t = np.linspace(-1, 1, m + 1)
    return (1 - grading) * t + grading * np.sin(0.5 * np.pi * t)


def superellipsoid_chart(par, a, p, uv):
    """xyz, r_u, r_v (9, n) of the chart par = (n, e1, e2, s0, s1, t0, t1) at uv (2, n)."""
    n, e1, e2 = par[0:3], par[3:6], par[6:9]
    s0, s1, t0, t1 = par[9:13]
    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 / 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 rounded_cube(a=1.0, p=8, m=6, norder=8, grading=0.5, exact=True):
    """Superellipsoid |x|^p+|y|^p+|z|^p = a^p, m x m patches per face.

    With exact=True the integrals use the analytic parametrization; otherwise the
    polynomial interpolant of the node data, as fmm3dbie does."""
    uvs, _ = disc_nodes(norder)
    br = graded_breaks(m, grading)
    blocks, params = [], []
    for n, e1, e2 in _FACES:
        for i in range(m):
            for j in range(m):
                par = np.array(list(n) + list(e1) + list(e2) + [br[i], br[i + 1], br[j], br[j + 1]],
                               dtype=float)
                params.append(par)
                g = superellipsoid_chart(par, a, p, uvs)
                nrm = np.cross(g[3:6].T, g[6:9].T).T
                nrm /= np.linalg.norm(nrm, axis=0)
                blocks.append(np.vstack([g, nrm]))
    ex = dict(params=np.array(params), a=a, p=p) if exact else None
    return Surface(norder, np.hstack(blocks), exact=ex)


def superellipsoid_volume(a, p):
    from math import gamma
    return 8 * a**3 * gamma(1 + 1 / p) ** 3 / gamma(1 + 3 / p)


def inside(x, a, p, margin=0.0):
    return np.sum(np.abs(x / a) ** p, axis=0) < (1 - margin) ** p