# Helpers for the Tenders (copied from tender-manta) (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))