"""Geometry helpers for the nutrition feeder (exec'd by nf_build.py).

Everything is laid out in the reference picture's own pixels (u right, v down, 1024 x 1536) and a front distance f in
millimetres off the wall; W() turns that into Blender metres (1 px = 0.5 mm, x right, -y toward the viewer, z up).
Surfaces are quads built from loops of equal vertex count (lofts), Coons patches and lathes; nothing is a boolean.
Faces are made with any winding; MB.to_mesh() makes the winding consistent and turns each part to face out.
"""
import bpy, bmesh, math
from mathutils import Vector

S = 0.0005                       # metres per reference pixel (from the hand in the picture)
U0, V0 = 513.0, 1345.0           # the panel's centre line and its bottom edge in the picture


def W(u, v, f):
    return Vector(((u - U0) * S, -f * 0.001, (V0 - v) * S))


def UV_OF(co):
    """picture pixels of a world point (u, v)"""
    return co.x / S + U0, V0 - co.z / S


# ------------------------------------------------------------------ rounded rectangles in picture pixels
class RR:
    """a rounded rectangle; a boundary point is named by a parameter that any other rounded rectangle can take:
    ('c', k, th) on corner k (0 TL, 1 TR, 2 BR, 3 BL) at th degrees, ('e', k, w) on edge k (0 top, 1 right,
    2 bottom, 3 left) at w along its straight part. The loop runs TL corner, top, TR, right, BR, bottom, BL, left."""

    def __init__(s, u0, v0, u1, v1, r):
        s.u0, s.v0, s.u1, s.v1, s.r = u0, v0, u1, v1, r

    def grown(s, d, r=None):
        return RR(s.u0 - d, s.v0 - d, s.u1 + d, s.v1 + d, max(0.5, s.r + d) if r is None else r)

    def at(s, prm):
        u0, v0, u1, v1, r = s.u0, s.v0, s.u1, s.v1, s.r
        kind, k, t = prm
        if kind == 'c':
            a = math.radians(t)
            ca, sa = math.cos(a), math.sin(a)
            return [(u0 + r - r * ca, v0 + r - r * sa), (u1 - r + r * sa, v0 + r - r * ca),
                    (u1 - r + r * ca, v1 - r + r * sa), (u0 + r - r * sa, v1 - r + r * ca)][k]
        if k == 0:
            return (u0 + r + (u1 - u0 - 2 * r) * t, v0)
        if k == 1:
            return (u1, v0 + r + (v1 - v0 - 2 * r) * t)
        if k == 2:
            return (u1 - r - (u1 - u0 - 2 * r) * t, v1)
        return (u0, v1 - r - (v1 - v0 - 2 * r) * t)

    def param_of(s, u, v, eps=1e-4):
        """parameter of a point on the sharp rectangle: corner zones (r from a corner along each side) map onto the
        arcs by arc length through the corner, so the rectangle's corner vertex lands on the 45 degree point"""
        u0, v0, u1, v1, r = s.u0, s.v0, s.u1, s.v1, s.r
        th = lambda sv: 90.0 * sv / (2 * r)
        if abs(v - v0) < eps:
            if u < u0 + r:
                return ('c', 0, th(r + (u - u0)))
            if u > u1 - r:
                return ('c', 1, th(u - (u1 - r)))
            return ('e', 0, (u - (u0 + r)) / (u1 - u0 - 2 * r))
        if abs(u - u1) < eps:
            if v < v0 + r:
                return ('c', 1, th(r + (v - v0)))
            if v > v1 - r:
                return ('c', 2, th(v - (v1 - r)))
            return ('e', 1, (v - (v0 + r)) / (v1 - v0 - 2 * r))
        if abs(v - v1) < eps:
            if u > u1 - r:
                return ('c', 2, th(r + (u1 - u)))
            if u < u0 + r:
                return ('c', 3, th((u0 + r) - u))
            return ('e', 2, ((u1 - r) - u) / (u1 - u0 - 2 * r))
        if abs(u - u0) < eps:
            if v > v1 - r:
                return ('c', 3, th(r + (v1 - v)))
            if v < v0 + r:
                return ('c', 0, th((v0 + r) - v))
            return ('e', 3, ((v1 - r) - v) / (v1 - v0 - 2 * r))
        raise ValueError("point not on the rectangle: %r" % ((u, v),))

    def arc_centres(s):
        r = s.r
        return [(s.u0 + r, s.v0 + r), (s.u1 - r, s.v0 + r), (s.u1 - r, s.v1 - r), (s.u0 + r, s.v1 - r)]


# ------------------------------------------------------------------ squircle (the tunnel), centre and exponent from the picture
SQ_C = (674.0, 972.0)
SQ_N = 2.5


def sq_point(phi, rho, c=SQ_C):
    """the squircle |x|^n + |y|^n = rho^n in the direction phi (radians, picture axes: v down)"""
    ca, sa = math.cos(phi), math.sin(phi)
    R = rho / (abs(ca) ** SQ_N + abs(sa) ** SQ_N) ** (1.0 / SQ_N)
    return (c[0] + R * ca, c[1] + R * sa)


# ------------------------------------------------------------------ mesh builder
class MB:
    def __init__(s):
        s.co, s.faces, s.mats = [], [], []
        s.flat = set()               # faces to shade flat (none so far)

    def v(s, co):
        s.co.append(Vector(co))
        return len(s.co) - 1

    def vs(s, pts):
        return [s.v(p) for p in pts]

    def f(s, idx, mat):
        uq = []
        for i in idx:
            if i not in uq:
                uq.append(i)
        if len(uq) >= 3:
            s.faces.append(uq)
            s.mats.append(mat)

    def bridge(s, la, lb, mat, closed=True):
        """quads between two vertex-index loops of equal length; mat: a name or a function (i) -> name"""
        n = len(la)
        assert n == len(lb), (n, len(lb))
        for i in range(n if closed else n - 1):
            j = (i + 1) % n
            s.f([la[i], la[j], lb[j], lb[i]], mat(i) if callable(mat) else mat)

    def coons(s, loop, corners, mat):
        """fill a closed loop with a quad grid; corners: four positions in the loop (in loop order) that split it
        into sides of equal counts two by two"""
        n = len(loop)
        c = list(corners)
        side = lambda a, b: [loop[(a + k) % n] for k in range((b - a) % n + 1)]
        top, right = side(c[0], c[1]), side(c[1], c[2])
        bot, left = side(c[2], c[3])[::-1], side(c[3], c[0])[::-1]
        nu, nv = len(top) - 1, len(right) - 1
        assert len(bot) - 1 == nu and len(left) - 1 == nv, (nu, nv, len(bot) - 1, len(left) - 1)
        P = lambda i: s.co[i]

        def chord(pts):
            d = [0.0]
            for a, b in zip(pts, pts[1:]):
                d.append(d[-1] + (P(b) - P(a)).length)
            return [x / d[-1] for x in d]
        su = [(a + b) / 2 for a, b in zip(chord(top), chord(bot))]
        tv = [(a + b) / 2 for a, b in zip(chord(left), chord(right))]
        grid = [[None] * (nv + 1) for _ in range(nu + 1)]
        for i in range(nu + 1):
            grid[i][0], grid[i][nv] = top[i], bot[i]
        for j in range(nv + 1):
            grid[0][j], grid[nu][j] = left[j], right[j]
        T00, T10, T01, T11 = P(top[0]), P(top[nu]), P(bot[0]), P(bot[nu])
        for i in range(1, nu):
            for j in range(1, nv):
                a, b = su[i], tv[j]
                p = ((1 - b) * P(top[i]) + b * P(bot[i]) + (1 - a) * P(left[j]) + a * P(right[j])
                     - ((1 - a) * (1 - b) * T00 + a * (1 - b) * T10 + (1 - a) * b * T01 + a * b * T11))
                grid[i][j] = s.v(p)
        for i in range(nu):
            for j in range(nv):
                s.f([grid[i][j], grid[i + 1][j], grid[i + 1][j + 1], grid[i][j + 1]], mat)
        return grid

    def to_mesh(s, name, mat_names, front_hint=True):
        """make the mesh; winding made consistent per part, each part turned so its frontmost face looks at -y"""
        me = bpy.data.meshes.get(name) or bpy.data.meshes.new(name)
        bm = bmesh.new()
        bv = [bm.verts.new(c) for c in s.co]
        bm.verts.ensure_lookup_table()
        idx = {m: i for i, m in enumerate(mat_names)}
        bad = 0
        for fi, (f, m) in enumerate(zip(s.faces, s.mats)):
            try:
                face = bm.faces.new([bv[i] for i in f])
                face.material_index = idx[m]
            except ValueError:
                bad += 1
        bmesh.ops.remove_doubles(bm, verts=bm.verts, dist=1e-7)
        bmesh.ops.recalc_face_normals(bm, faces=bm.faces)
        if front_hint:
            seen = set()
            for f0 in bm.faces:
                if f0.index in seen:
                    continue
                comp, stack = [], [f0]
                seen.add(f0.index)
                while stack:
                    f = stack.pop()
                    comp.append(f)
                    for e in f.edges:
                        for g in e.link_faces:
                            if g.index not in seen:
                                seen.add(g.index)
                                stack.append(g)
                front = min(comp, key=lambda f: (f.calc_center_median().y, -abs(f.normal.y)))
                if front.normal.y > 0:
                    for f in comp:
                        f.normal_flip()
        bm.normal_update()
        me.clear_geometry() if len(me.vertices) else None
        me.materials.clear()                 # before the faces: clearing the slots also zeroes the faces' indices
        for m in mat_names:
            me.materials.append(bpy.data.materials.get(m) or bpy.data.materials.new(m))
        bm.to_mesh(me)
        bm.free()
        return me, bad


# ------------------------------------------------------------------ lathe (vertical axis)
def lathe(mb, prof, uc, fc, nseg, mats, cap_start=False, cap_end=False):
    """prof: [(r_px, v_px)] top to bottom; mats: one name per segment (len(prof)-1) or a single name.
    A point with r 0 is a single vertex (a cone tip). Caps close a ring with a Coons patch (quads)."""
    rings = []
    for (r, v) in prof:
        if r <= 1e-6:
            rings.append([mb.v(W(uc, v, fc))])
        else:
            rings.append([mb.v(W(uc + r * math.cos(a), v, fc + 0.5 * r * math.sin(a)))
                          for a in (2 * math.pi * k / nseg for k in range(nseg))])
    for k in range(len(rings) - 1):
        m = mats if isinstance(mats, str) else mats[k]
        a, b = rings[k], rings[k + 1]
        if len(a) == 1 and len(b) == 1:
            continue
        if len(a) == 1 or len(b) == 1:
            ring, tip = (b, a[0]) if len(a) == 1 else (a, b[0])
            for i in range(nseg):
                mb.f([ring[i], ring[(i + 1) % nseg], tip], m)
        else:
            mb.bridge(a, b, m)
    q = nseg // 4
    if cap_start and len(rings[0]) > 1:
        mb.coons(rings[0], [0, q, 2 * q, 3 * q], mats if isinstance(mats, str) else mats[0])
    if cap_end and len(rings[-1]) > 1:
        mb.coons(rings[-1], [0, q, 2 * q, 3 * q], mats if isinstance(mats, str) else mats[-1])
    return rings


# ------------------------------------------------------------------ prisms on a grid of cells (the injector mount)
def cell_solid(mb, cols, rows, cells, moved=None):
    """cells: {(i, j): (f_back, f_front, mat, front_mat)} on columns/rows in picture px; a solid of axis-aligned
    prisms whose shared parts are left out. Faces at f 0 (against the back plate) are left out. moved: {(i, j):
    (u, v)} moves a grid corner (for chamfered corners) at every depth."""
    moved = moved or {}
    levels = sorted({x for c in cells.values() for x in c[:2]})
    vid = {}

    def V_(i, j, f):
        key = (i, j, f)
        if key not in vid:
            u, v = moved.get((i, j), (cols[i], rows[j]))
            vid[key] = mb.v(W(u, v, f))
        return vid[key]

    def sub(iv, ivs):
        """interval iv minus the union of ivs"""
        out = [iv]
        for (a, b) in ivs:
            nxt = []
            for (x, y) in out:
                if b <= x or a >= y:
                    nxt.append((x, y))
                else:
                    if a > x:
                        nxt.append((x, a))
                    if b < y:
                        nxt.append((b, y))
            out = nxt
        return [o for o in out if o[1] - o[0] > 1e-9]

    for (i, j), (fb, ff, mat, fmat) in cells.items():
        mb.f([V_(i, j, ff), V_(i + 1, j, ff), V_(i + 1, j + 1, ff), V_(i, j + 1, ff)], fmat or mat)
        if fb > 1e-9:
            mb.f([V_(i, j, fb), V_(i, j + 1, fb), V_(i + 1, j + 1, fb), V_(i + 1, j, fb)], mat)
        # four sides: (neighbour, the side's two corners)
        for (di, dj), (a, b) in (((0, -1), ((i, j), (i + 1, j))), ((1, 0), ((i + 1, j), (i + 1, j + 1))),
                                 ((0, 1), ((i + 1, j + 1), (i, j + 1))), ((-1, 0), ((i, j + 1), (i, j)))):
            nb = cells.get((i + di, j + dj))
            for (lo, hi) in sub((fb, ff), [nb[:2]] if nb else []):
                cuts = [lo] + [x for x in levels if lo < x < hi] + [hi]
                for f0, f1 in zip(cuts, cuts[1:]):
                    mb.f([V_(*a, f1), V_(*b, f1), V_(*b, f0), V_(*a, f0)], mat)


# ------------------------------------------------------------------ audit
def audit(ob):
    me = ob.data
    bm = bmesh.new()
    bm.from_mesh(me)
    q = sum(1 for f in bm.faces if len(f.verts) == 4)
    t = sum(1 for f in bm.faces if len(f.verts) == 3)
    ng = sum(1 for f in bm.faces if len(f.verts) > 4)
    open_e = sum(1 for e in bm.edges if len(e.link_faces) == 1)
    nonman = sum(1 for e in bm.edges if len(e.link_faces) > 2)
    tiny = sum(1 for f in bm.faces if f.calc_area() < 1e-9)
    tris = sum(len(f.verts) - 2 for f in bm.faces)
    bm.free()
    return dict(quads=q, tris=t, ngons=ng, open_edges=open_e, nonmanifold_edges=nonman, zero_area=tiny, triangles=tris,
                verts=len(me.vertices))