"""Rounded cube (superellipsoid) surface built from quadrilateral patches.

The surface is |x|^p + |y|^p + |z|^p = a^p with p even. Each of the six cube
faces (s, t) in [-1, 1]^2 is mapped radially onto the surface,
x = a c / ||c||_p with c = n + s e1 + t e2. Larger p gives sharper edges.
Patch breakpoints in s and t are graded toward the face edges, where the
curvature is concentrated.
"""

import numpy as np

import fmm3dbie_ctypes as fb

# (n, e1, e2) with e1 x e2 = n, so that d/ds x d/dt points outward
_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):
    """m+1 breakpoints on [-1, 1]; grading in [0, 1] blends uniform and sine-clustered."""
    t = np.linspace(-1, 1, m + 1)
    return (1 - grading) * t + grading * np.sin(0.5 * np.pi * t)


def rounded_cube(a=1.0, p=8, m=6, norder=8, grading=0.5):
    """Return an fmm3dbie Surface for the superellipsoid of half-width a.

    m x m patches per face, each discretized with (norder+1)^2 Gauss-Legendre nodes.
    """
    iptype = 11
    uvs = fb.get_disc_nodes(norder, iptype)  # (2, npols) on [-1, 1]^2
    br = graded_breaks(m, grading)
    blocks = []
    for n, e1, e2 in _FACES:
        n, e1, e2 = (np.array(v, dtype=float) for v in (n, e1, e2))
        for i in range(m):
            for j in range(m):
                s0, s1 = br[i], br[i + 1]
                t0, t1 = br[j], br[j + 1]
                s = s0 + 0.5 * (uvs[0] + 1) * (s1 - s0)
                t = t0 + 0.5 * (uvs[1] + 1) * (t1 - t0)
                c = n[:, None] + s * e1[:, None] + t * e2[:, None]  # (3, npols)
                N = np.sum(c**p, axis=0) ** (1.0 / p)
                gradN = c ** (p - 1) / N ** (p - 1)  # dN/dc
                xyz = a * c / N

                def deriv(dc, scale):
                    dN = np.sum(gradN * dc[:, None], axis=0)
                    return a * scale * (dc[:, None] / N - c * dN / N**2)

                du = deriv(e1, 0.5 * (s1 - s0))
                dv = deriv(e2, 0.5 * (t1 - t0))
                nrm = np.cross(du.T, dv.T).T
                nrm /= np.linalg.norm(nrm, axis=0)
                blocks.append(np.vstack([xyz, du, dv, nrm]))
    srcvals = np.hstack(blocks)
    return fb.Surface(norder, iptype, srcvals)


def superellipsoid_volume(a, p):
    """Exact volume of |x|^p+|y|^p+|z|^p <= a^p (used to check the discretization)."""
    from math import gamma
    return 8 * a**3 * gamma(1 + 1 / p) ** 3 / gamma(1 + 3 / p)