"""A cymbal as a thin closed body built from quadrilateral patches (pure NumPy).

fmm3dbie's integral equations need a closed surface, so the cymbal is not a
zero-thickness shell but a thin body around an axisymmetric midsurface
z = F(rho), rho = x^2 + y^2, with thickness T(rho). The midsurface has a raised
bell at the center and a shallow bow outside it; the thickness decreases from
the center to the rim, with zero slope at the rim. The top and bottom faces
are the offset surfaces M +- (T/2) N, with M the midsurface and N its upward
unit normal, and a half-torus of diameter T closes the body at the rim. Every
patch lies entirely on one face or on the rim, so every patch is analytic even
though the curvature jumps where the faces meet the rim.

The faces are meshed with an O-grid in the (x, y) plane: an equiangular central
square of m_c x m_c patches, then rings of patches between the square and a
circle and between successive circles. Rings are graded toward the rim, where
the density varies on the scale of the thickness. Neighboring rings need not
have the same number of patches around, so the mesh is not conforming; neither
fmm3dbie nor the GPU code requires that it be.

Node derivatives dx/du, dx/dv are computed by the complex step method, which is
exact to rounding for this analytic parametrization.
"""

import numpy as np

from . import surface

CSTEP = 1e-30


class CymbalShape:
    """Axisymmetric cymbal shape; lengths in meters.

    radius: outer radius, including the rim.
    t_rim: thickness at the rim; taper: thickness at the center / t_rim.
    bell_radius, bell_height, bell_width: the bell is a logistic step in rho,
        of height bell_height, centered at r = bell_radius, with width about
        bell_width in r.
    bow_height: fall of the midsurface from the center to the rim, apart from the bell.
    """

    def __init__(self, radius=0.2, t_rim=0.01, taper=2.0, bell_radius=0.06, bell_height=0.025,
                 bell_width=0.01, bow_height=0.015):
        self.radius, self.t_rim, self.taper = radius, t_rim, taper
        self.bell_radius, self.bell_height, self.bell_width = bell_radius, bell_height, bell_width
        self.bow_height = bow_height
        self.r_mid = radius - t_rim / 2  # radius of the midsurface, where the rim is attached

    # midsurface height, its rho-derivative, and the thickness, all functions of rho
    def _bell_arg(self, rho):
        return (rho - self.bell_radius ** 2) / (2 * self.bell_radius * self.bell_width)

    def F(self, rho):
        return (self.bow_height * (1 - rho / self.r_mid ** 2)
                + self.bell_height / (1 + np.exp(self._bell_arg(rho))))

    def dF(self, rho):
        e = np.exp(self._bell_arg(rho))
        return (-self.bow_height / self.r_mid ** 2
                - self.bell_height * e / (1 + e) ** 2 / (2 * self.bell_radius * self.bell_width))

    def T(self, rho):
        return self.t_rim * (1 + (self.taper - 1) * (1 - rho / self.r_mid ** 2) ** 2)

    def face_point(self, x, y, side):
        """Point of the top (side = +1) or bottom (side = -1) face above/below (x, y) of
        the midsurface. Complex-safe."""
        rho = x * x + y * y
        g = 2 * self.dF(rho)  # grad F = g (x, y)
        q = np.sqrt(1 + g * g * rho)
        s = side * self.T(rho) / 2 / q
        return np.array([x - s * g * x, y - s * g * y, self.F(rho) + s])

    def rim_point(self, theta, phi):
        """Point of the rim at azimuth theta; phi from -pi/2 (bottom face) to pi/2 (top face)."""
        R = self.r_mid
        rho = R * R
        fp = 2 * R * self.dF(rho)  # dz/dr of the midsurface at the rim
        q = np.sqrt(1 + fp * fp)
        c, s = np.cos(theta), np.sin(theta)
        h = self.T(rho) / 2
        a, b = h * np.cos(phi) / q, h * np.sin(phi) / q  # along the tangent (1, fp) and normal (-fp, 1)
        rr = R + a - b * fp
        return np.array([rr * c, rr * s, self.F(rho) + a * fp + b])

    def cross_section(self, n=400):
        """Closed outline of the body in the (r, z) half plane, (2, m), counterclockwise."""
        r = np.linspace(0, self.r_mid, n)
        top = self.face_point(r, 0 * r, 1)
        bot = self.face_point(r, 0 * r, -1)
        phi = np.linspace(-np.pi / 2, np.pi / 2, n // 4)
        rim = self.rim_point(0 * phi, phi)
        rz = np.hstack([bot[[0, 2]], rim[[0, 2]], top[[0, 2]][:, ::-1]])
        return rz

    def inside(self, pts, margin=0.0):
        """True for points (3, n) inside the body or within margin of its surface.

        Every vertical line meets the cross section in one interval, so we offset the
        outline outward by margin, split it at its outermost point into a lower and an
        upper chain, and compare z with both chains at the point's radius."""
        rz = self.cross_section(4000)
        if margin > 0:
            t = np.roll(rz, -1, axis=1) - np.roll(rz, 1, axis=1)  # counterclockwise tangent
            t /= np.linalg.norm(t, axis=0)
            rz = rz + margin * np.array([t[1], -t[0]])  # outward normal
            rz[0] = np.maximum(rz[0], 0)
        i = int(np.argmax(rz[0]))
        lo, up = rz[:, :i + 1], rz[:, i:][:, ::-1]
        r = np.hypot(pts[0], pts[1])
        z_lo = np.interp(r, np.maximum.accumulate(lo[0]), lo[1])
        z_up = np.interp(r, np.maximum.accumulate(up[0]), up[1])
        return (r <= rz[0, i]) & (pts[2] > z_lo) & (pts[2] < z_up)

    def volume(self, n=20000):
        """Volume of the body by Pappus' theorem on the cross section (a check)."""
        rz = self.cross_section(n)
        r, z = rz
        r1, z1 = np.roll(r, -1), np.roll(z, -1)
        # int r dA over the cross section, by Green's theorem: (1/2) oint r^2 dz
        return 2 * np.pi * np.sum((r * r + r * r1 + r1 * r1) / 6 * (z1 - z))


def ring_layout(shape, h, grade=1.5, c_frac=0.12):
    """Radii and patch counts of the O-grid for a target patch size h.

    Returns (m_c, c, radii, counts): the central square has half-width c and
    m_c x m_c patches (m_c >= 2); ring i lies between radii[i] and radii[i + 1]
    (radii[0] = c means the square) and has counts[i] patches around (a multiple
    of 4). The ring around the square reaches r = 2c, so that it is 0.59c wide at
    the square's corners. Radial widths are at most h; they grow geometrically
    from t_rim at the rim by the factor grade, and they are at most twice the
    bell width across the bell.
    """
    R = shape.r_mid
    c = c_frac * R
    m_c = max(2, int(np.ceil(2 * c / h)))
    first = 2 * c  # outer radius of the ring around the square
    b0 = shape.bell_radius - 2.5 * shape.bell_width
    b1 = shape.bell_radius + 2.5 * shape.bell_width

    def wmax(edge):
        w = min(h, shape.t_rim + (grade - 1) * (R - edge))
        if edge > b0 and edge - w < b1:
            w = min(w, 2 * shape.bell_width)
        return w

    radii = [R]
    while True:
        w = wmax(radii[-1])
        if radii[-1] - w < first + 0.5 * w:
            break
        radii.append(radii[-1] - w)
    radii = np.array([c, first] + radii[::-1])
    counts = []
    for i in range(len(radii) - 1):
        width = radii[i + 1] - radii[i]
        arc = min(h, 2.5 * width)
        n = 4 * max(m_c, int(np.ceil(2 * np.pi * radii[i + 1] / (4 * arc))))
        counts.append(n)
    counts[0] = 4 * m_c  # the first ring matches the square
    # rings never have fewer patches around than the ring inside them
    counts = list(np.maximum.accumulate(counts))
    return m_c, c, radii, counts


def _face_charts(m_c, c, radii, counts):
    """Planar charts (u, v) in [-1, 1]^2 -> (x, y) for all face patches, as closures.
    Each chart is oriented so that d/du x d/dv points up."""
    charts = []
    q = np.pi / 4
    # central square, equiangular so that its edges match the first ring
    br = np.linspace(-1, 1, m_c + 1)
    for i in range(m_c):
        for j in range(m_c):
            def ch(u, v, a0=br[i], a1=br[i + 1], b0=br[j], b1=br[j + 1]):
                s = a0 + 0.5 * (u + 1) * (a1 - a0)
                t = b0 + 0.5 * (v + 1) * (b1 - b0)
                return c * np.tan(q * s), c * np.tan(q * t)
            charts.append(ch)
    # rings; u is radial, v is the angle (counterclockwise)
    for ir in range(len(radii) - 1):
        r0, r1, n = radii[ir], radii[ir + 1], counts[ir]
        dth = 2 * np.pi / n
        for k in range(n):
            th0 = -q + k * dth  # patch angles start at the square corner (-pi/4)
            if ir == 0:
                def ch(u, v, th0=th0, r1=r1, dth=dth):
                    th = th0 + 0.5 * (v + 1) * dth
                    sec = np.floor((th.real + q) / (2 * q))  # sector index 0..3
                    loc = th - sec * 2 * q  # angle relative to the sector's center
                    rin = c / np.cos(loc)
                    a = 0.5 * (u + 1)
                    rad = (1 - a) * rin + a * r1
                    return rad * np.cos(th), rad * np.sin(th)
            else:
                def ch(u, v, th0=th0, r0=r0, r1=r1, dth=dth):
                    th = th0 + 0.5 * (v + 1) * dth
                    rad = r0 + 0.5 * (u + 1) * (r1 - r0)
                    return rad * np.cos(th), rad * np.sin(th)
            charts.append(ch)
    return charts


def _node_data(fun, uvs):
    """xyz, r_u, r_v, unit normal (12, n) of the map fun(u, v) -> (3, n), with complex step derivatives."""
    u, v = uvs[0].astype(complex), uvs[1].astype(complex)
    x = np.real(fun(u, v))
    du = np.imag(fun(u + 1j * CSTEP, v)) / CSTEP
    dv = np.imag(fun(u, v + 1j * CSTEP)) / CSTEP
    nrm = np.cross(du.T, dv.T).T
    nrm /= np.linalg.norm(nrm, axis=0)
    return np.vstack([x, du, dv, nrm])


def patch_maps(shape, h, grade=1.5):
    """All patch maps (u, v) -> xyz (3, n) with outward orientation, and the ring layout."""
    m_c, c, radii, counts = ring_layout(shape, h, grade)
    charts = _face_charts(m_c, c, radii, counts)
    maps = []
    for ch in charts:  # top faces
        maps.append(lambda u, v, ch=ch: shape.face_point(*ch(u, v), 1))
    for ch in charts:  # bottom faces: swapping u and v flips the normal
        maps.append(lambda u, v, ch=ch: shape.face_point(*ch(v, u), -1))
    nrim = counts[-1]
    for k in range(nrim):  # rim: u is the azimuth, v the angle over the rim
        th0 = -np.pi / 4 + 2 * np.pi * k / nrim
        maps.append(lambda u, v, th0=th0: shape.rim_point(th0 + 0.5 * (u + 1) * 2 * np.pi / nrim,
                                                         0.5 * np.pi * v))
    return maps, dict(m_c=m_c, c=c, radii=radii, counts=counts, nrim=nrim)


def cymbal_srcvals(shape, h, norder, grade=1.5):
    """Node data (12, npts) in the fmm3dbie iptype 11 convention, and the layout."""
    uvs, _ = surface.disc_nodes(norder)
    maps, layout = patch_maps(shape, h, grade)
    return np.hstack([_node_data(f, uvs) for f in maps]), layout


def cymbal(shape, h, norder, grade=1.5):
    """A surface.Surface for the cymbal, with the interpolated geometry inside the
    integrals, as in fmm3dbie."""
    srcvals, layout = cymbal_srcvals(shape, h, norder, grade)
    s = surface.Surface(norder, srcvals)
    s.layout = layout
    return s