# Spearfish: clean-topology helpers (exec'd after lib2.py and tender_lib.py). Everything here makes quads only:
# caps are Coons-patch grids instead of fans or n-gons, spheres are cube spheres, patches are ladders.
import bmesh, math
from mathutils import Vector, Matrix, Euler


def coons_cap(bm, ring, mat_index=0):
    """fill a closed ring of 4n BMVerts with an n x n quad grid (a Coons patch over its four sides); returns the faces.
    The ring's order sets the faces' winding (same direction as the ring)."""
    m = len(ring)
    assert m % 4 == 0 and m >= 4, m
    n = m // 4
    P = [v.co.copy() for v in ring]
    B = [P[i] for i in range(0, n + 1)]                      # u 0..1 along side 0
    R = [P[n + i] for i in range(0, n + 1)]                  # v 0..1 along side 1
    T = [P[(3 * n - i) % m] for i in range(0, n + 1)]        # u 0..1, side 2 reversed
    L = [P[(4 * n - i) % m] for i in range(0, n + 1)]        # v 0..1, side 3 reversed
    G = [[None] * (n + 1) for _ in range(n + 1)]             # G[i][j], i along u, j along v
    for i in range(n + 1):
        G[i][0] = ring[i]
        G[i][n] = ring[(3 * n - i) % m]
    for j in range(n + 1):
        G[0][j] = ring[(4 * n - j) % m]
        G[n][j] = ring[n + j]
    for i in range(1, n):
        for j in range(1, n):
            u, v = i / n, j / n
            p = ((1 - v) * B[i] + v * T[i] + (1 - u) * L[j] + u * R[j]
                 - ((1 - u) * (1 - v) * B[0] + u * (1 - v) * B[n] + (1 - u) * v * T[0] + u * v * T[n]))
            G[i][j] = bm.verts.new(p)
    faces = []
    for i in range(n):
        for j in range(n):
            f = bm.faces.new((G[i][j], G[i + 1][j], G[i + 1][j + 1], G[i][j + 1]))
            f.material_index = mat_index
            faces.append(f)
    return faces


def ring_rows(bm, rows, closed=True, mat_index=0):
    """quads between consecutive rows of BMVerts (equal length)"""
    out = []
    for a, b in zip(rows, rows[1:]):
        n = len(a)
        for k in range(n if closed else n - 1):
            k2 = (k + 1) % n
            f = bm.faces.new((a[k], b[k], b[k2], a[k2]))
            f.material_index = mat_index
            out.append(f)
    return out


def orient_out(bm, centre):
    """turn every face so it points away from centre (closed, roughly convex parts)"""
    bm.normal_update()
    flip = [f for f in bm.faces if f.normal.dot(f.calc_center_median() - Vector(centre)) < 0]
    if flip:
        bmesh.ops.reverse_faces(bm, faces=flip)
    bm.normal_update()


def qsphere(mb, ctr, radii, mat, cuts=3, rot=(0, 0, 0)):
    """a cube sphere (quads only), scaled to radii"""
    bm = bmesh.new()
    bmesh.ops.create_cube(bm, size=2.0)
    bmesh.ops.subdivide_edges(bm, edges=bm.edges[:], cuts=cuts, use_grid_fill=True)
    for v in bm.verts:
        d = v.co.normalized()
        v.co = Vector((d.x * radii[0], d.y * radii[1], d.z * radii[2]))
    bm.normal_update()
    mb.merge(bm, mat, Matrix.Translation(ctr) @ Matrix.LocRotScale(None, Euler(rot), None))


def revolve(mb, prof, mats, segs=16, xf=None, cap0=True, cap1=True, cap0_mat=None, cap1_mat=None):
    """revolve a profile [(r, a)] (r > 0 everywhere) about local +Z; open ends closed by Coons caps (segs % 4 == 0)"""
    bm = bmesh.new()
    rows = [[bm.verts.new((r * math.cos(2 * math.pi * k / segs), r * math.sin(2 * math.pi * k / segs), a)) for k in range(segs)]
            for r, a in prof]
    ml = list(mats) if isinstance(mats, (list, tuple)) else [mats] * (len(prof) - 1)
    for i in range(len(rows) - 1):
        mi = mb.mi(ml[i])
        ring_rows(bm, [rows[i], rows[i + 1]], mat_index=mi)
    if cap0:
        coons_cap(bm, rows[0][::-1], mb.mi(cap0_mat or ml[0]))
    if cap1:
        coons_cap(bm, rows[-1], mb.mi(cap1_mat or ml[-1]))
    c = Vector((0, 0, 0.5 * (prof[0][1] + prof[-1][1])))
    bm.normal_update()
    # orientation: the side faces point away from the axis
    flip = []
    for f in bm.faces:
        p = f.calc_center_median()
        radial = Vector((p.x, p.y, 0))
        ref = radial if radial.length > 1e-6 and abs(f.normal.z) < 0.9 else Vector((0, 0, p.z - c.z))
        if f.normal.dot(ref) < 0:
            flip.append(f)
    if flip:
        bmesh.ops.reverse_faces(bm, faces=flip)
    bm.normal_update()
    mb.merge(bm, None, xf)


def qcyl(mb, r, h, ctr, mat, rot=(0, 0, 0), segs=12, r2=None):
    """a cylinder along local Z with Coons-cap ends (quads only)"""
    r2 = r if r2 is None else r2
    xf = Matrix.Translation(ctr) @ Matrix.LocRotScale(None, Euler(rot), None)
    revolve(mb, [(r, -h / 2), (r2, h / 2)], mat, segs=segs, xf=xf)


def qbox(mb, size, ctr, mat, rot=(0, 0, 0), bevel=0.0):
    """a box; bevelled with two segments, which keeps the corners quads"""
    mb.box(size, ctr, mat, rot=rot, bevel=bevel, seg=2 if bevel > 0 else 1)


def ladder_patch(mb, lo_hi, t0, t1, to3d, mat, nt=8, nacross=3, lift=0.002, bulge=0.0, normal=None):
    """a thin plate over a region given as t in [t0, t1] -> (s_lo, s_hi) across: a grid of nt x nacross quads.
    to3d(t, s) -> point on the hull; normal(p) -> outward normal. Ends at t0/t1 keep their width (no fan)."""
    bm = bmesh.new()
    rows = []
    for i in range(nt + 1):
        t = t0 + (t1 - t0) * i / nt
        lo, hi = lo_hi(t)
        row = []
        for j in range(nacross + 1):
            s = lo + (hi - lo) * j / nacross
            p = to3d(t, s)
            n = normal(p)
            w = 1.0 - (2.0 * j / nacross - 1.0) ** 2
            e = 1.0 - (2.0 * i / nt - 1.0) ** 2
            row.append(bm.verts.new(p + n * (lift + bulge * max(0.0, w) * max(0.0, e) ** 0.5)))
        rows.append(row)
    ring_rows(bm, rows, closed=False)
    bm.normal_update()
    mid = rows[nt // 2][nacross // 2].co
    nm = normal(mid)
    if sum((f.normal.dot(nm) for f in bm.faces)) < 0:
        bmesh.ops.reverse_faces(bm, faces=bm.faces[:])
    mb.merge(bm, mat)
    return rows


def stadium(t0, t1, half):
    """lo_hi for a stadium (capsule) from t0 to t1 of half width `half`: round ends"""
    r = half
    def f(t):
        a, b = min(t0, t1), max(t0, t1)
        if t < a + r:
            d = (a + r - t) / r
        elif t > b - r:
            d = (t - (b - r)) / r
        else:
            d = 0.0
        w = r * math.sqrt(max(0.0, 1.0 - min(1.0, d) ** 2))
        return -max(w, 0.18 * r), max(w, 0.18 * r)
    return f