# Helpers for the Remora Tender (exec'd after lib2.py, shares its namespace).
import bpy, bmesh, math
from mathutils import Vector, Matrix
from mathutils.bvhtree import BVHTree
from mathutils import kdtree
from mathutils.interpolate import poly_3d_calc

YPOS = Matrix.Rotation(-math.pi / 2, 4, 'X')        # lathe axis (local +Z) -> world +Y
T = lambda *p: Matrix.Translation(p)


def pchip(pts):
    """monotone cubic through [(x, y)]; clamped outside the table"""
    xs = [p[0] for p in pts]; ys = [p[1] for p in pts]
    n = len(xs)
    h = [xs[i + 1] - xs[i] for i in range(n - 1)]
    d = [(ys[i + 1] - ys[i]) / h[i] for i in range(n - 1)]
    m = [d[0]] + [0.0] * (n - 2) + [d[-1]]
    for i in range(1, n - 1):
        if d[i - 1] * d[i] > 0:
            w1, w2 = 2 * h[i] + h[i - 1], h[i] + 2 * h[i - 1]
            m[i] = (w1 + w2) / (w1 / d[i - 1] + w2 / d[i])
    def f(x):
        if x <= xs[0]:
            return ys[0]
        if x >= xs[-1]:
            return ys[-1]
        lo, hi = 0, n - 1
        while hi - lo > 1:
            mid = (lo + hi) // 2
            if xs[mid] <= x:
                lo = mid
            else:
                hi = mid
        i = lo
        t = (x - xs[i]) / h[i]
        t2, t3 = t * t, t * t * t
        return ((2 * t3 - 3 * t2 + 1) * ys[i] + (t3 - 2 * t2 + t) * h[i] * m[i] + (-2 * t3 + 3 * t2) * ys[i + 1]
                + (t3 - t2) * h[i] * m[i + 1])
    return f


def smoothstep(a, b, x):
    t = max(0.0, min(1.0, (x - a) / (b - a)))
    return t * t * (3 - 2 * t)


def lathe(mb, prof, mats, segs=32, xf=None):
    """Revolve a profile [(r, a)] about local +Z (a = height). Starts and ends on the axis (closed shell, fan ends).
    Writes UV_Tile in metres (flat bands seen along the axis, walls unrolled). Same as the docking arm's."""
    bm = bmesh.new()
    uvl = bm.loops.layers.uv.new("UV_Tile")
    cols = []
    for (r, a) in prof:
        if r < 1e-9:
            cols.append([bm.verts.new((0, 0, a))])
        else:
            cols.append([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)])
    ml = list(mats) if isinstance(mats, (list, tuple)) else [mats] * (len(prof) - 1)
    assert len(ml) == len(prof) - 1, (len(ml), len(prof))
    sp = [0.0]
    for i in range(len(prof) - 1):
        sp.append(sp[-1] + math.hypot(prof[i + 1][0] - prof[i][0], prof[i + 1][1] - prof[i][1]))
    for i in range(len(prof) - 1):
        A, B = cols[i], cols[i + 1]
        mi = mb.mi(ml[i])
        flat = abs(prof[i + 1][0] - prof[i][0]) >= abs(prof[i + 1][1] - prof[i][1])
        rm = (prof[i][0] + prof[i + 1][0]) / 2
        for k in range(segs):
            k2 = (k + 1) % segs
            if len(A) == 1 and len(B) == 1:
                continue
            if len(A) == 1:
                f = bm.faces.new((A[0], B[k], B[k2]))
            elif len(B) == 1:
                f = bm.faces.new((A[k], B[0], A[k2]))
            else:
                f = bm.faces.new((A[k], B[k], B[k2], A[k2]))
            f.material_index = mi
            if flat:
                for l in f.loops:
                    l[uvl].uv = (l.vert.co.x, l.vert.co.y)
            else:
                u_a, u_b = rm * 2 * math.pi * k / segs, rm * 2 * math.pi * (k + 1) / segs
                uvs = {}
                for col, s in ((A, sp[i]), (B, sp[i + 1])):
                    if len(col) == 1:
                        uvs[col[0]] = ((u_a + u_b) / 2, s)
                    else:
                        uvs[col[k]] = (u_a, s); uvs[col[k2]] = (u_b, s)
                for l in f.loops:
                    l[uvl].uv = uvs[l.vert]
    bm.normal_update()
    bmesh.ops.recalc_face_normals(bm, faces=bm.faces[:])
    mb.merge(bm, None, xf)


def revolve_loop(mb, loop, mats, segs=32, xf=None):
    """Revolve a CLOSED profile loop [(r, a)] (no point on the axis) about local +Z: a ring torus-like solid."""
    bm = bmesh.new()
    uvl = bm.loops.layers.uv.new("UV_Tile")
    cols = [[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 loop]
    n = len(loop)
    ml = list(mats) if isinstance(mats, (list, tuple)) else [mats] * n
    sp = [0.0]
    for i in range(n):
        j = (i + 1) % n
        sp.append(sp[-1] + math.hypot(loop[j][0] - loop[i][0], loop[j][1] - loop[i][1]))
    for i in range(n):
        j = (i + 1) % n
        A, B = cols[i], cols[j]
        mi = mb.mi(ml[i])
        rm = (loop[i][0] + loop[j][0]) / 2
        for k in range(segs):
            k2 = (k + 1) % segs
            f = bm.faces.new((A[k], B[k], B[k2], A[k2]))
            f.material_index = mi
            u_a, u_b = rm * 2 * math.pi * k / segs, rm * 2 * math.pi * (k + 1) / segs
            uvs = {A[k]: (u_a, sp[i]), A[k2]: (u_b, sp[i]), B[k]: (u_a, sp[i + 1]), B[k2]: (u_b, sp[i + 1])}
            for l in f.loops:
                l[uvl].uv = uvs[l.vert]
    bm.normal_update()
    bmesh.ops.recalc_face_normals(bm, faces=bm.faces[:])
    mb.merge(bm, None, xf)


def grid_rings(mb, rings, mats, closed_u=True, cap0=None, cap1=None, uv=None, skip=None, expect=None, vattr=None):
    """Skin consecutive rings (lists of points, equal length) with quads. mats: one material or one per band.
    cap0/cap1: a point to fan the first/last ring to. uv: rings of (u, v) for UV_Tile (cap points take uv_cap0/1 from
    the ring's first entry). skip(i, k) -> True leaves out the face between ring i and i+1 at k (i = -1: cap0 fan,
    i = len-1: cap1 fan). expect = (i, k, direction): the faces are turned so that face (i, k) faces that way."""
    bm = bmesh.new()
    uvl = bm.loops.layers.uv.new("UV_Tile")
    lays = {name: bm.verts.layers.float.new(name) for name in (vattr or {})}   # layers first: a new layer invalidates verts
    V = [[bm.verts.new(p) for p in r] for r in rings]
    for name, vals in (vattr or {}).items():                  # per-vertex float fields (rings of values), e.g. seam distances
        lay = lays[name]
        for r, vr in zip(V, vals):
            for v, x in zip(r, vr):
                v[lay] = x
    m = len(rings[0])
    nr = len(rings)
    ml = list(mats) if isinstance(mats, (list, tuple)) else [mats] * (nr - 1)
    span = m if closed_u else m - 1
    probe = None
    for i in range(nr - 1):
        mi = mb.mi(ml[i])
        for k in range(span):
            if skip is not None and skip(i, k):
                continue
            k2 = (k + 1) % m
            ids = [(i, k), (i + 1, k), (i + 1, k2), (i, k2)]
            f = bm.faces.new([V[a][b] for a, b in ids])
            f.material_index = mi
            if uv is not None:
                for l, (a, b) in zip(f.loops, ids):
                    l[uvl].uv = uv[a][b]
            if expect is not None and (i, k) == tuple(expect[:2]):
                probe = f
    for cap, ri, ci in ((cap0, 0, -1), (cap1, nr - 1, nr - 1)):
        if cap is None:
            continue
        cv = bm.verts.new(cap)
        mi = mb.mi(ml[0] if ri == 0 else ml[-1])
        for k in range(span):
            if skip is not None and skip(ci, k):
                continue
            k2 = (k + 1) % m
            vs = [V[ri][k], V[ri][k2], cv] if ri == 0 else [V[ri][k2], V[ri][k], cv]
            f = bm.faces.new(vs)
            f.material_index = mi
            if uv is not None:
                for l in f.loops:
                    l[uvl].uv = uv[ri][k]
    loose = [v for v in bm.verts if not v.link_faces]
    if loose:
        bmesh.ops.delete(bm, geom=loose, context='VERTS')
    bm.normal_update()
    if expect is not None:
        assert probe is not None, "orientation probe face was skipped"
        if probe.normal.dot(Vector(expect[2])) < 0:
            bmesh.ops.reverse_faces(bm, faces=bm.faces[:])
            bm.normal_update()
    mb.merge(bm, None)


def loft(mb, secs, mat, xf=None):
    """Beam along +X through chamfered-rectangle sections [(x, half width, z0, z1, chamfer)]: 8 quads per bay,
    3 quads per end (the docking arm's loft without pockets)."""
    bm = bmesh.new()
    rings = []
    for (x, w, z0, z1, c) in secs:
        pts = [(-w + c, z0), (w - c, z0), (w, z0 + c), (w, z1 - c), (w - c, z1), (-w + c, z1), (-w, z1 - c), (-w, z0 + c)]
        rings.append([bm.verts.new((x, y, z)) for y, z in pts])
    for i in range(len(rings) - 1):
        for k in range(8):
            k2 = (k + 1) % 8
            bm.faces.new((rings[i][k], rings[i + 1][k], rings[i + 1][k2], rings[i][k2]))
    for r in (rings[0], rings[-1]):
        bm.faces.new((r[7], r[0], r[1], r[2])); bm.faces.new((r[7], r[2], r[3], r[6])); bm.faces.new((r[6], r[3], r[4], r[5]))
    bm.normal_update()
    bmesh.ops.recalc_face_normals(bm, faces=bm.faces[:])
    mb.merge(bm, mat, xf)


def foil(mb, sections, mat, xf=None):
    """A thin lofted plate (fin, blade, vane) through sections [(centre Vector, chord dir Vector, normal dir Vector,
    chord, thickness)]: an 8-point lens section, quads along, 3 quads on each end."""
    bm = bmesh.new()
    rings = []
    prof = [(-0.50, 0.0), (-0.40, 0.55), (-0.10, 1.0), (0.25, 0.80), (0.50, 0.0), (0.25, -0.80), (-0.10, -1.0), (-0.40, -0.55)]
    for (c, cd, nd, ch, th) in sections:
        rings.append([bm.verts.new(c + cd * (u * ch) + nd * (v * th * 0.5)) for u, v in prof])
    for i in range(len(rings) - 1):
        for k in range(8):
            k2 = (k + 1) % 8
            bm.faces.new((rings[i][k], rings[i + 1][k], rings[i + 1][k2], rings[i][k2]))
    for r in (rings[0], rings[-1]):
        bm.faces.new((r[0], r[1], r[2], r[7])); bm.faces.new((r[7], r[2], r[3], r[6])); bm.faces.new((r[6], r[3], r[4], r[5]))
    bm.normal_update()
    bmesh.ops.recalc_face_normals(bm, faces=bm.faces[:])
    mb.merge(bm, mat, xf)


class Surface:
    """BVH of a closed mesh with smooth (interpolated) normals: ray hits and nearest points."""
    def __init__(self, me):
        self.co = [v.co.copy() for v in me.vertices]
        self.vn = [v.normal.copy() for v in me.vertices]
        self.polys = [tuple(p.vertices) for p in me.polygons]
        self.bvh = BVHTree.FromPolygons(self.co, self.polys)

    def _n(self, p, fi):
        ids = self.polys[fi]
        w = poly_3d_calc([self.co[i] for i in ids], p)
        n = Vector()
        for wi, i in zip(w, ids):
            n += self.vn[i] * wi
        return n.normalized()

    def ray(self, o, d):
        hit, nrm, fi, dist = self.bvh.ray_cast(Vector(o), Vector(d).normalized())
        if hit is None:
            return None, None
        return hit, self._n(hit, fi)

    def near(self, p):
        hit, nrm, fi, dist = self.bvh.find_nearest(Vector(p))
        return hit, self._n(hit, fi)


def prune_inside(mb, inside_depth, keep_depth=0.015):
    """drop faces of an attached part whose every vertex lies deeper than keep_depth inside the hull"""
    bm = mb.bm
    dead = [f for f in bm.faces if all(inside_depth(v.co) > keep_depth for v in f.verts)]
    bmesh.ops.delete(bm, geom=dead, context='FACES')
    loose = [v for v in bm.verts if not v.link_faces]
    if loose:
        bmesh.ops.delete(bm, geom=loose, context='VERTS')
    return len(dead)


def resample_profile(prof, mats, max_step):
    """split every profile segment longer than max_step into equal pieces (materials follow their segment)"""
    ml = list(mats) if isinstance(mats, (list, tuple)) else [mats] * (len(prof) - 1)
    out, om = [prof[0]], []
    for i in range(len(prof) - 1):
        (r0, a0), (r1, a1) = prof[i], prof[i + 1]
        n = max(1, int(math.ceil(math.hypot(r1 - r0, a1 - a0) / max_step)))
        for j in range(1, n + 1):
            out.append((r0 + (r1 - r0) * j / n, a0 + (a1 - a0) * j / n))
            om.append(ml[i])
    return out, om


def bm_tree(bm):
    bm.verts.ensure_lookup_table()
    return BVHTree.FromBMesh(bm)


def overlaps(bm_a, bm_b):
    """intersecting face pairs between two bmeshes (indices into their face sequences)"""
    bm_a.faces.ensure_lookup_table(); bm_b.faces.ensure_lookup_table()
    return BVHTree.FromBMesh(bm_a).overlap(BVHTree.FromBMesh(bm_b))