"""Door robot arm: small shape helpers shared by the part modules (metres)."""
import math
import bmesh
from mathutils import Vector, Matrix
from dra_lib import (lathe, slab, rbox, tube, knurl, star_slab, holed_surface, fill_cap, face, merge, xform, basis,
                     frame, fillet, ccw, W, G, S, R, O, GL, LA, TAU, MKEY, circle2, decal_strip, decal_poly, decal_dot)

M_XZ = basis((0, 0, 0), (0, -1, 0), (1, 0, 0))     # local x -> X, local y -> Z, local z -> -Y
M_YZ = basis((0, 0, 0), (1, 0, 0), (0, 1, 0))      # local x -> Y, local y -> Z, local z -> X


def rrect(w, h, r, cx=0.0, cy=0.0, r2=None):
    """rounded rectangle outline (x, y, round), counter-clockwise; r2: round of the two top corners if different"""
    x0, x1, y0, y1 = cx - w / 2, cx + w / 2, cy - h / 2, cy + h / 2
    rt = r if r2 is None else r2
    return [(x0, y0, r), (x1, y0, r), (x1, y1, rt), (x0, y1, rt)]


def stadium(L, wd, cx=0.0, cy=0.0, ang=0.0, n=6):
    """a slot outline (rounded ends) of length L and width wd, turned by ang"""
    r = wd / 2
    pts = []
    for k in range(n + 1):
        a = -math.pi / 2 + math.pi * k / n
        pts.append((L / 2 - r + r * math.cos(a), r * math.sin(a)))
    for k in range(n + 1):
        a = math.pi / 2 + math.pi * k / n
        pts.append((-L / 2 + r + r * math.cos(a), r * math.sin(a)))
    ca, sa = math.cos(ang), math.sin(ang)
    return [(cx + x * ca - y * sa, cy + x * sa + y * ca) for x, y in pts]


def xz_slab(outline, y0, y1, bev_lo=0.003, bev_hi=0.003, **kw):
    """an outline in the X-Z plane pushed across Y from y0 to y1 (bev_lo rounds the rim at y0, bev_hi at y1)"""
    bm = slab(outline, -y1, -y0, bev=bev_lo, bev0=bev_hi, **kw)
    return xform(bm, M_XZ)


def yz_slab(outline, x0, x1, bev_lo=0.003, bev_hi=0.003, **kw):
    """an outline in the Y-Z plane pushed along X from x0 to x1"""
    bm = slab(outline, x0, x1, bev=bev_hi, bev0=bev_lo, **kw)
    return xform(bm, M_YZ)


def xy_slab(outline, z0, z1, bev_lo=0.003, bev_hi=0.003, **kw):
    return slab(outline, z0, z1, bev=bev_hi, bev0=bev_lo, **kw)


def screw(r=0.006, h=None, mat=S, sock=G, segs=None):
    """socket-head cap screw standing on z = 0 (its foot sinks 1 mm)"""
    h = h or r * 0.8
    segs = segs or (8 if r < 0.0065 else 12)
    prof = [(0, -0.001, 0, mat), (r, -0.001, 0, mat), (r, h, r * 0.4, mat), (r * 0.5, h, 0, sock),
            (r * 0.5, h * 0.45, 0, sock), (0, h * 0.45, 0, sock)]
    return lathe(prof, segs, n=1)


def dome(r=0.008, h=None, mat=W, segs=10):
    """a round dome cap (screw cover) on z = 0"""
    h = h or r * 0.6
    prof = [(0, -0.001, 0, mat), (r, -0.001, 0, mat), (r, h * 0.3, 0, mat)]
    for k in range(1, 4):
        a = (math.pi / 2) * k / 4
        prof.append((r * math.cos(a), h * 0.3 + (h * 0.7) * math.sin(a), 0, mat))
    prof.append((0, h, 0, mat))
    return lathe(prof, segs)


def hexbolt(s=0.024, h=0.012, washer=True, mat=S):
    """hex head bolt on z = 0 (s = across flats), with a washer"""
    Rh = s / 2 / math.cos(math.pi / 6)
    bm = lathe([(0, -0.001, 0, mat), (Rh, -0.001, 0, mat), (Rh, h, 0.0012, mat), (0, h, 0, mat)], 6, phase=math.pi / 6)
    if washer:
        xform(bm, Matrix.Translation((0, 0, 0.0025)))
        w = lathe([(0, -0.001, 0, mat), (s * 0.95, -0.001, 0, mat), (s * 0.95, 0.0025, 0.0006, mat), (0, 0.0025, 0, mat)], 16)
        merge(bm, w)
        w.free()
    # the boss of the thread end
    return bm


def ring(r_in, r_out, h, segs=32, f=0.0015, mat=G, z0=0.0):
    return lathe([(r_in, z0, f, mat), (r_out, z0, f, mat), (r_out, z0 + h, f, mat), (r_in, z0 + h, f, mat)], segs, closed=True)


def cyl(r, h, segs=32, f=0.002, mat=G, z0=0.0, f0=None):
    f0 = f if f0 is None else f0
    return lathe([(0, z0, 0, mat), (r, z0, f0, mat), (r, z0 + h, f, mat), (0, z0 + h, 0, mat)], segs)


def grooved_cyl(r, h, grooves, depth=0.003, width=0.004, segs=32, f=0.003, mat=W, groove_mat=None, z0=0.0):
    """a cylinder along +Z with ring grooves at the given heights (from z0)"""
    gm = mat if groove_mat is None else groove_mat
    prof = [(0, z0, 0, mat), (r, z0, f, mat)]
    for g in sorted(grooves):
        prof += [(r, z0 + g - width / 2, 0, mat), (r - depth, z0 + g - width / 2 + 0.0008, 0, gm),
                 (r - depth, z0 + g + width / 2 - 0.0008, 0, mat), (r, z0 + g + width / 2, 0, mat)]
    prof += [(r, z0 + h, f, mat), (0, z0 + h, 0, mat)]
    return lathe(prof, segs)


def holed_plate(w, h, t, cols, rows, pitch, hole_r, depth, mat=G, hole_mat=None, ch=0.0012):
    """a plate (w x h, thickness t, front face at z = t, back at 0) with a grid of round blind holes in its face"""
    bm = bmesh.new()
    us = [-w / 2] + [-(cols * pitch) / 2 + i * pitch for i in range(cols + 1)] + [w / 2]
    vs = [-h / 2] + [-(rows * pitch) / 2 + j * pitch for j in range(rows + 1)] + [h / 2]
    holes = {(i, j) for i in range(1, cols + 1) for j in range(1, rows + 1)}
    lat = holed_surface(bm, lambda u, v: Vector((u, v, t)), us, vs, holes, hole_r, depth,
                        lambda u, v: Vector((0, 0, 1)), mat=mat, hole_mat=hole_mat, ret_lattice=True)
    per = lat_perimeter(lat)
    ring1 = []
    ring2 = []
    for v in per:
        p = Vector(v.co)
        o = Vector((0 if abs(abs(p.x) - w / 2) > 1e-7 else math.copysign(1, p.x),
                    0 if abs(abs(p.y) - h / 2) > 1e-7 else math.copysign(1, p.y), 0))
        ring1.append(bm.verts.new((p.x + o.x * ch, p.y + o.y * ch, t - ch)))
        ring2.append(bm.verts.new((p.x + o.x * ch, p.y + o.y * ch, 0.0)))
    n = len(per)
    for k in range(n):
        k2 = (k + 1) % n
        mid = (per[k].co + per[k2].co) / 2
        o = Vector((mid.x, mid.y, 0))
        face(bm, [per[k], per[k2], ring1[k2], ring1[k]], mat, out=o + Vector((0, 0, 0.5 * o.length)))
        face(bm, [ring1[k], ring1[k2], ring2[k2], ring2[k]], mat, out=o)
    fill_cap(bm, ring2, (0, 0, -1), mat)
    bmesh.ops.remove_doubles(bm, verts=bm.verts[:], dist=1e-8)
    return bm


def lat_perimeter(lat):
    """the boundary of a lattice (rows of BMVerts) counter-clockwise"""
    nv, nu = len(lat), len(lat[0])
    per = [lat[0][i] for i in range(nu)]
    per += [lat[j][nu - 1] for j in range(1, nv)]
    per += [lat[nv - 1][i] for i in range(nu - 2, -1, -1)]
    per += [lat[j][0] for j in range(nv - 2, 0, -1)]
    return per


def gear_outline(teeth, r_root, r_tip, w_root=0.36, w_tip=0.14, phase=0.0):
    """spur gear outline, points in angular order; w_* = half widths as fractions of the tooth pitch angle"""
    pts = []
    pa = TAU / teeth
    for k in range(teeth):
        c = phase + k * pa
        for a, r in ((c - w_root * pa, r_root), (c - w_tip * pa, r_tip), (c + w_tip * pa, r_tip), (c + w_root * pa, r_root)):
            pts.append((r * math.cos(a), r * math.sin(a)))
    return pts


def notched_disc(r, notches, nw, nd, n_arc=4, phase=0.0, open_rim=True):
    """a disc outline with rectangular notches cut into its rim (radial slots open at the rim)"""
    pts = []
    pa = TAU / notches
    half = math.asin(min(0.95, (nw / 2) / r))
    for k in range(notches):
        c = phase + k * pa
        a0, a1 = c + half, c + pa - half
        for j in range(n_arc + 1):
            a = a0 + (a1 - a0) * j / n_arc
            pts.append((r * math.cos(a), r * math.sin(a)))
        # the notch at angle c + pa: walls straight in toward the centre, width nw
        cn = c + pa
        dn = Vector((math.cos(cn), math.sin(cn)))
        tn = Vector((-dn.y, dn.x))
        rin = r - nd
        for p in ((dn * math.sqrt(max(r * r - (nw / 2) ** 2, 0)) - tn * nw / 2), (dn * rin - tn * nw / 2),
                  (dn * rin + tn * nw / 2), (dn * math.sqrt(max(r * r - (nw / 2) ** 2, 0)) + tn * nw / 2)):
            pts.append((p.x, p.y))
    return pts


def half_discs(r, gap, h, mat, segs=24, bev=0.001, z0=0.0):
    """two half discs with a straight slot of width 'gap' between them (a 'minus' screw head), slot along local X"""
    g = gap / 2
    a_lo = math.asin(min(0.99, g / r))
    up = [(r * math.cos(a_lo + (math.pi - 2 * a_lo) * k / segs), r * math.sin(a_lo + (math.pi - 2 * a_lo) * k / segs))
          for k in range(segs + 1)]
    acc = slab(up, z0, z0 + h, bev=bev, bev0=0.0, mat=mat)
    lo = [(x, -y) for x, y in reversed(up)]
    b2 = slab(lo, z0, z0 + h, bev=bev, bev0=0.0, mat=mat)
    merge(acc, b2)
    b2.free()
    return acc


def quarter_discs(r, gap, h, mat, segs=10, bev=0.001, z0=0.0, phase=0.0):
    """four quarter discs with a '+' cross groove of width gap"""
    g = gap / 2
    a_lo = math.asin(min(0.99, g / r))
    base = [(g, g)] + [(r * math.cos(a_lo + (math.pi / 2 - 2 * a_lo) * k / segs),
                         r * math.sin(a_lo + (math.pi / 2 - 2 * a_lo) * k / segs)) for k in range(segs + 1)]
    acc = None
    for q in range(4):
        a = phase + q * math.pi / 2
        ca, sa = math.cos(a), math.sin(a)
        poly = [(x * ca - y * sa, x * sa + y * ca) for x, y in base]
        bm = slab(poly, z0, z0 + h, bev=bev, bev0=0.0, mat=mat)
        if acc is None:
            acc = bm
        else:
            merge(acc, bm)
            bm.free()
    return acc


def u_handle(width, height, r_tube, rad=None, sleeve=None, blocks=True, mat=W, block_mat=None):
    """a bent-tube handle in the local XZ plane: feet at x = 0 (z = +-height/2), standing out along +X by width.
    sleeve: (n_rings, ring_r, length) on the outer leg."""
    rad = rad or min(width, height) * 0.25
    pts = [(0.0, 0, -height / 2), (width, 0, -height / 2), (width, 0, height / 2), (0.0, 0, height / 2)]
    bm = tube(pts, r_tube, 12, rad, mat=mat)
    if sleeve:
        n, rr, L = sleeve
        for k in range(n):
            z = -L / 2 + (k + 0.5) * L / n
            sl = lathe([(r_tube * 0.95, -L / n * 0.42, 0.0015, mat), (rr, -L / n * 0.42, 0.0015, mat),
                        (rr, L / n * 0.42, 0.0015, mat), (r_tube * 0.95, L / n * 0.42, 0.0015, mat)], 12, closed=True)
            xform(sl, Matrix.Translation((width, 0, z)))
            merge(bm, sl)
            sl.free()
    if blocks:
        bmat = mat if block_mat is None else block_mat
        for z in (-height / 2, height / 2):
            b = rbox((0.04, r_tube * 3.2, r_tube * 3.2), 0.003)
            xform(b, Matrix.Translation((0.004, 0, z)))
            for f in b.faces:
                f.material_index = bmat
            merge(bm, b)
            b.free()
    return bm


def ribbed_boot(r_in, r_out, length, ribs, mat=R, segs=16):
    """a ribbed rubber boot (bellows) along +Z"""
    prof = [(0.0, 0.0, 0, mat), (r_in, 0.0, 0, mat)]
    p = length / ribs
    for k in range(ribs):
        z = k * p
        prof += [(r_in, z + 0.15 * p, 0, mat), (r_out, z + 0.35 * p, 0, mat), (r_out, z + 0.65 * p, 0, mat), (r_in, z + 0.85 * p, 0, mat)]
    prof += [(r_in, length, 0, mat), (0.0, length, 0, mat)]
    return lathe(prof, segs)


def thread_rod(r, length, pitch, mat=S, segs=10):
    """a threaded rod along +Z (ring ribs standing in for the thread)"""
    prof = [(0.0, 0.0, 0, mat), (r * 0.8, 0.0, 0, mat)]
    n = max(2, int(length / pitch))
    p = length / n
    for k in range(n):
        z = k * p
        prof += [(r * 0.8, z + 0.1 * p, 0, mat), (r, z + 0.4 * p, 0, mat), (r, z + 0.6 * p, 0, mat), (r * 0.8, z + 0.9 * p, 0, mat)]
    prof += [(r * 0.8, length, 0, mat), (0.0, length, 0, mat)]
    return lathe(prof, segs)