"""Small numpy mesh kit for the Aethel-9 build: quad lofts with Coons caps, revolves, sweeps, rounded-rect rings.
Meshes are dicts: v (N,3) float, f list of index tuples, uv list of per-corner (u,v) tuples, m list of material names."""
import numpy as np


class Mesh:
    def __init__(self):
        self.v, self.f, self.uv, self.m = [], [], [], []

    def add_verts(self, pts):
        base = len(self.v)
        self.v.extend([tuple(p) for p in np.asarray(pts, float).reshape(-1, 3)])
        return base

    def face(self, idx, uv, mat):
        self.f.append(tuple(int(i) for i in idx)); self.uv.append([tuple(map(float, u)) for u in uv]); self.m.append(mat)

    def merge(self, other):
        base = len(self.v)
        self.v.extend(other.v)
        self.f.extend([tuple(i + base for i in f) for f in other.f])
        self.uv.extend(other.uv); self.m.extend(other.m)
        return self

    def arrays(self):
        return np.array(self.v, float), self.f, self.uv, self.m

    def transform(self, M):
        V = np.array(self.v); V = V @ np.asarray(M)[:3, :3].T + np.asarray(M)[:3, 3]
        self.v = [tuple(p) for p in V]
        return self

    def mirror_x(self):
        """A mirrored copy (X -> -X) with faces reversed so normals stay outward."""
        o = Mesh()
        o.v = [(-x, y, z) for x, y, z in self.v]
        o.f = [tuple(reversed(f)) for f in self.f]
        o.uv = [list(reversed(u)) for u in self.uv]
        o.m = list(self.m)
        return o


def arclen(pts, closed=False):
    p = np.asarray(pts, float)
    if closed:
        p = np.vstack([p, p[:1]])
    return np.concatenate([[0], np.cumsum(np.linalg.norm(np.diff(p, axis=0), axis=1))])


def grid(mesh, P, mat, closed_u=False, uv_scale=1.0, flip=False, uvs=None, row_mats=None):
    """P: (ni, nj, 3) grid of points; quads between rows i and columns j (j wraps if closed_u).
    UVs by arc length in metres * uv_scale unless uvs (ni, nj, 2) is given."""
    P = np.asarray(P, float); ni, nj = P.shape[:2]
    base = mesh.add_verts(P.reshape(-1, 3))
    if uvs is None:
        uu = np.zeros((ni, nj + (1 if closed_u else 0)))
        Pc = np.concatenate([P, P[:, :1]], 1) if closed_u else P
        for i in range(ni):
            uu[i] = arclen(Pc[i])
        um = uu.mean(0)
        vv = np.zeros((ni, nj))
        for j in range(nj):
            vv[:, j] = arclen(P[:, j])
        vm = vv.mean(1)
        uvs = np.zeros((ni, nj + (1 if closed_u else 0), 2))
        uvs[..., 0] = um[None, :] * uv_scale
        uvs[..., 1] = vm[:, None] * uv_scale
    cols = nj if closed_u else nj - 1
    for i in range(ni - 1):
        for j in range(cols):
            jn = (j + 1) % nj
            q = [base + i * nj + j, base + (i + 1) * nj + j, base + (i + 1) * nj + jn, base + i * nj + jn]
            uq = [uvs[i, j], uvs[i + 1, j], uvs[i + 1, j + 1], uvs[i, j + 1]]
            if flip:
                q = q[::-1]; uq = uq[::-1]
            mesh.face(q, uq, row_mats[i] if row_mats else mat)
    return base


def coons_cap(mesh, ring, mat, flip=False, uv_scale=1.0, center=None, ring_idx=None):
    """Quad patch over a closed ring of 4k points (k per side): Coons interior, planar-ish rings.
    ring_idx: the ring's existing vertex indices (weld the cap to a loft so the piece is one closed solid)."""
    R = np.asarray(ring, float); n = len(R); k = n // 4
    assert n % 4 == 0
    s0 = R[0:k + 1]; s1 = R[k:2 * k + 1]; s2 = R[2 * k:3 * k + 1][::-1]; s3 = np.vstack([R[3 * k:], R[:1]])[::-1]
    # s0: bottom (u along), s2: top (u along), s3: left (v along), s1: right (v along)
    G = np.zeros((k + 1, k + 1, 3))
    u = np.linspace(0, 1, k + 1)
    for a in range(k + 1):
        for b in range(k + 1):
            uu, vv = u[a], u[b]
            G[b, a] = ((1 - vv) * s0[a] + vv * s2[a] + (1 - uu) * s3[b] + uu * s1[b]
                       - ((1 - uu) * (1 - vv) * s0[0] + uu * (1 - vv) * s0[-1] + (1 - uu) * vv * s2[0] + uu * vv * s2[-1]))
    # reuse the ring's own vertices on the border: build the full grid and weld by index map
    base = mesh.add_verts(G[1:-1, 1:-1].reshape(-1, 3)) if k > 1 else len(mesh.v)
    idx = np.zeros((k + 1, k + 1), int)
    if ring_idx is None:
        rb = mesh.add_verts(R)
        ring_idx = [rb + a for a in range(n)]
    for a in range(k + 1):
        idx[0, a] = ring_idx[a]
        idx[a, k] = ring_idx[k + a]
        idx[k, k - a] = ring_idx[(2 * k + a) % n]
        idx[k - a, 0] = ring_idx[(3 * k + a) % n]
    for b in range(1, k):
        for a in range(1, k):
            idx[b, a] = base + (b - 1) * (k - 1) + (a - 1)
    # planar UV: project on the patch's best plane
    c = G.reshape(-1, 3).mean(0) if center is None else center
    w = G.reshape(-1, 3) - c
    _, _, vt = np.linalg.svd(w, full_matrices=False)
    e1, e2 = vt[0], vt[1]
    UV = np.stack([(G - c) @ e1, (G - c) @ e2], -1) * uv_scale
    for b in range(k):
        for a in range(k):
            q = [idx[b, a], idx[b, a + 1], idx[b + 1, a + 1], idx[b + 1, a]]
            uq = [UV[b, a], UV[b, a + 1], UV[b + 1, a + 1], UV[b + 1, a]]
            if flip:
                q = q[::-1]; uq = uq[::-1]
            mesh.face(q, uq, mat)


def loft(mesh, rings, mat, cap0=None, cap1=None, uv_scale=1.0, mats=None):
    """Closed rings (all the same count, a multiple of 4 if capped) lofted in order; caps: material name or None.
    Ring orientation: counter-clockwise seen from the loft's start looking along it -> normals point outward."""
    P = np.asarray(rings, float)
    base = grid(mesh, P, mat, closed_u=True, uv_scale=uv_scale)
    ni, nj = P.shape[:2]
    if cap0:
        coons_cap(mesh, P[0], cap0, flip=False, uv_scale=uv_scale, ring_idx=[base + j for j in range(nj)])
    if cap1:
        coons_cap(mesh, P[-1], cap1, flip=True, uv_scale=uv_scale, ring_idx=[base + (ni - 1) * nj + j for j in range(nj)])


def frame(axis):
    """Two unit vectors perpendicular to axis."""
    a = np.asarray(axis, float); a = a / np.linalg.norm(a)
    t = np.array([0, 0, 1.0]) if abs(a[2]) < 0.9 else np.array([1.0, 0, 0])
    e1 = np.cross(a, t); e1 /= np.linalg.norm(e1); e2 = np.cross(a, e1)
    return e1, e2


def circle(center, axis, r, n, phase=0.0, e1=None, e2=None):
    if e1 is None:
        e1, e2 = frame(axis)
    th = phase + np.linspace(0, 2 * np.pi, n, endpoint=False)
    return np.asarray(center, float) + r * (np.cos(th)[:, None] * e1 + np.sin(th)[:, None] * e2)


def revolve(mesh, prof, center, axis, n, mat, cap0=None, cap1=None, uv_scale=1.0, phase=0.0):
    """prof: list of (r, a) along the axis (a in metres from center). Rings counter-clockwise about +axis.
    A first or last point with r <= 1 mm closes that end in a pole (a triangle fan), not a ring of tiny faces."""
    ax = np.asarray(axis, float); ax /= np.linalg.norm(ax)
    e1, e2 = frame(ax)
    prof = list(prof)
    pole0 = prof[0][0] <= 0.001
    pole1 = prof[-1][0] <= 0.001
    p0 = np.asarray(center) + prof[0][1] * ax if pole0 else None
    p1 = np.asarray(center) + prof[-1][1] * ax if pole1 else None
    body = prof[1 if pole0 else 0: len(prof) - 1 if pole1 else len(prof)]
    rings = [circle(np.asarray(center) + a * ax, ax, r, n, phase, e1, e2) for r, a in body]
    P = np.asarray(rings, float)
    base = grid(mesh, P, mat, closed_u=True, uv_scale=uv_scale)
    ni, nj = P.shape[:2]
    for pole, row, fl, capm in ((p0, 0, False, cap0), (p1, ni - 1, True, cap1)):
        if pole is not None:
            pi = mesh.add_verts(pole[None])
            for j in range(nj):
                a_, b_ = base + row * nj + j, base + row * nj + (j + 1) % nj
                f = [pi, a_, b_] if not fl else [pi, b_, a_]          # matches the ring faces' winding
                mesh.face(f, [(0.5, 0.5), (0, 0), (1, 0)] if not fl else [(0.5, 0.5), (1, 0), (0, 0)], mat)
        elif capm:
            idx = [base + row * nj + j for j in range(nj)]
            coons_cap(mesh, P[row], capm, flip=fl, uv_scale=uv_scale, ring_idx=idx)


def rrect(cx, cy, w, h, r, n_side=4, n_corner=4):
    """Rounded rectangle ring in 2D, counter-clockwise, starting at the middle of the bottom-right corner's arc, so each
    quarter of the ring (n_corner + 2*n_side points) is one side: a Coons cap then lines up with the rectangle."""
    r = min(r, 0.45 * w, 0.45 * h)      # keep straight sides (no coincident points)
    nc = n_corner if n_corner % 2 == 0 else n_corner + 1
    cs = [(cx + w / 2 - r, cy - h / 2 + r, -90), (cx + w / 2 - r, cy + h / 2 - r, 0),
          (cx - w / 2 + r, cy + h / 2 - r, 90), (cx - w / 2 + r, cy - h / 2 + r, 180)]
    def arc(k, t0, t1, n):
        ccx, ccy, a0 = cs[k]
        return [(ccx + r * np.cos(np.deg2rad(a0 + 90 * t)), ccy + r * np.sin(np.deg2rad(a0 + 90 * t)))
                for t in np.linspace(t0, t1, n + 1)[:-1]]
    out = []
    for k in range(4):
        out += arc(k, 0.5, 1.0, nc // 2)
        p = arc(k, 1.0, 1.0, 1)[0] if False else (cs[k][0] + r * np.cos(np.deg2rad(cs[k][2] + 90)),
                                                  cs[k][1] + r * np.sin(np.deg2rad(cs[k][2] + 90)))
        k2 = (k + 1) % 4
        q = (cs[k2][0] + r * np.cos(np.deg2rad(cs[k2][2])), cs[k2][1] + r * np.sin(np.deg2rad(cs[k2][2])))
        out += [(p[0] + (q[0] - p[0]) * t, p[1] + (q[1] - p[1]) * t) for t in np.linspace(0, 1, 2 * n_side + 1)[:-1]]
        out += arc(k2, 0.0, 0.5, nc // 2)
    return np.array(out)


def ring3(pts2, origin, ex, ey):
    p = np.asarray(pts2, float)
    return np.asarray(origin, float) + p[:, :1] * np.asarray(ex, float) + p[:, 1:2] * np.asarray(ey, float)


def sweep(mesh, path, normals, prof, mat, closed=False, uv_scale=1.0, side=None, mats=None):
    """Sweep a 2D profile (s, n) along a 3D path: s along the side direction (tangent x normal, or 'side' given per point),
    n along the given normals. Returns the grid (len(path), len(prof), 3)."""
    P = np.asarray(path, float); N = np.asarray(normals, float)
    N = N / np.linalg.norm(N, axis=1, keepdims=True)
    if side is None:
        T = np.gradient(P, axis=0) if not closed else (np.roll(P, -1, 0) - np.roll(P, 1, 0))
        T /= np.linalg.norm(T, axis=1, keepdims=True)
        Sd = np.cross(T, N); Sd /= np.linalg.norm(Sd, axis=1, keepdims=True)
    else:
        Sd = np.asarray(side, float)
    pr = np.asarray(prof, float)
    G = P[:, None, :] + pr[None, :, :1] * Sd[:, None, :] + pr[None, :, 1:2] * N[:, None, :]
    if closed:
        G2 = np.concatenate([G, G[:1]], 0)
        # treat as open grid along the path with the first row repeated (UV seam), then weld by reuse
        _grid_rows(mesh, G, mat, closed_path=True, uv_scale=uv_scale, mats=mats)
    else:
        _grid_rows(mesh, G, mat, closed_path=False, uv_scale=uv_scale, mats=mats)
    return G


def _grid_rows(mesh, G, mat, closed_path, uv_scale, mats=None):
    ni, nj = G.shape[:2]
    base = mesh.add_verts(G.reshape(-1, 3))
    Gc = np.concatenate([G, G[:1]], 0) if closed_path else G
    along = arclen(Gc[:, nj // 2])
    across = arclen(G[0])
    rows = ni if closed_path else ni - 1
    for i in range(rows):
        inx = (i + 1) % ni
        for j in range(nj - 1):
            q = [base + i * nj + j, base + i * nj + j + 1, base + inx * nj + j + 1, base + inx * nj + j]
            uq = [(across[j] * uv_scale, along[i] * uv_scale), (across[j + 1] * uv_scale, along[i] * uv_scale),
                  (across[j + 1] * uv_scale, along[i + 1] * uv_scale), (across[j] * uv_scale, along[i + 1] * uv_scale)]
            mesh.face(q, uq, mats[j] if mats else mat)


def box(mesh, c, size, mat, r=0.0, ex=(1, 0, 0), ey=(0, 1, 0), ez=(0, 0, 1), n_corner=3, uv_scale=1.0):
    """A box (rounded vertical edges if r) as a loft of rounded rectangles along ez, capped."""
    c = np.asarray(c, float); sx, sy, sz = size
    ex, ey, ez = (np.asarray(a, float) for a in (ex, ey, ez))
    if r > 0.25 * min(sx, sy):        # a rounding the plate cannot hold: square corners (no sliver sides)
        r = 0.0
    if r > 0:
        rr = rrect(0, 0, sx, sy, r, n_side=2, n_corner=n_corner)
    else:   # square corners: 8 points (corner, side middle, ...), starting at the bottom-right corner
        rr = np.array([(sx / 2, -sy / 2), (sx / 2, 0), (sx / 2, sy / 2), (0, sy / 2), (-sx / 2, sy / 2), (-sx / 2, 0),
                       (-sx / 2, -sy / 2), (0, -sy / 2)])
    rings = [ring3(rr, c + ez * (-sz / 2), ex, ey), ring3(rr, c + ez * (sz / 2), ex, ey)]
    loft(mesh, rings, mat, cap0=mat, cap1=mat, uv_scale=uv_scale)


def components(F, nv):
    """Connected face groups (sharing vertex indices)."""
    parent = list(range(nv))
    def find(a):
        while parent[a] != a:
            parent[a] = parent[parent[a]]; a = parent[a]
        return a
    for f in F:
        r0 = find(f[0])
        for i in f[1:]:
            r = find(i)
            if r != r0:
                parent[r] = r0
    groups = {}
    for k, f in enumerate(F):
        groups.setdefault(find(f[0]), []).append(k)
    return list(groups.values())


def orient(mesh, ref_point=(0.0, 0.0, 1.4), report=None):
    """Make every connected piece face outward: closed pieces by their signed volume, open pieces by an area-weighted vote
    of the face normals against the direction from the piece's centroid (flat pieces: from ref_point, the vessel's
    middle). Flips whole pieces only, so a piece's own consistent winding is kept."""
    V = np.array(mesh.v, float)
    F = mesh.f
    flipped = 0
    for comp in components(F, len(V)):
        edges = {}
        for k in comp:
            f = F[k]
            for a in range(len(f)):
                e = tuple(sorted((f[a], f[(a + 1) % len(f)])))
                edges[e] = edges.get(e, 0) + 1
        closed = all(c == 2 for c in edges.values())
        nrm, cen, area = [], [], []
        vol = 0.0
        for k in comp:
            p = V[list(F[k])]
            n = np.zeros(3)
            for a in range(len(p)):
                n += np.cross(p[a], p[(a + 1) % len(p)])
            for a in range(1, len(p) - 1):
                vol += np.dot(p[0], np.cross(p[a], p[a + 1])) / 6.0
            ar = np.linalg.norm(n) / 2
            nrm.append(n / (2 * ar + 1e-15)); cen.append(p.mean(0)); area.append(ar)
        nrm, cen, area = np.array(nrm), np.array(cen), np.array(area)
        if closed and abs(vol) > 1e-9:
            flip = vol < 0
        else:
            c0 = (cen * area[:, None]).sum(0) / max(area.sum(), 1e-15)
            d = cen - c0
            vote = np.sum(area * np.sum(nrm * d, 1)) / max(np.sum(area * np.linalg.norm(d, axis=1)), 1e-15)
            if abs(vote) < 0.15:
                d = cen - np.asarray(ref_point)
                vote = np.sum(area * np.sum(nrm * d, 1))
            flip = vote < 0
        if flip:
            flipped += 1
            for k in comp:
                mesh.f[k] = tuple(reversed(mesh.f[k]))
                mesh.uv[k] = list(reversed(mesh.uv[k]))
    if report is not None:
        report.append(flipped)
    return mesh


def winding_report(mesh):
    """Edges shared by two faces that traverse it the same way (inconsistent winding), and open edges."""
    from collections import Counter
    directed = Counter()
    for f in mesh.f:
        for a in range(len(f)):
            directed[(f[a], f[(a + 1) % len(f)])] += 1
    bad = sum(1 for (a, b), c in directed.items() if c > 1 or (a, b) in directed and (b, a) not in directed and False)
    same = sum(1 for (a, b), c in directed.items() if c > 1)
    und = Counter()
    for (a, b), c in directed.items():
        und[tuple(sorted((a, b)))] += c
    open_e = sum(1 for e, c in und.items() if c == 1)
    return dict(same_direction=same, open_edges=open_e)


def quad_ring(corners, k):
    """A ring of 4k points round a quadrilateral (2D corners in order), k points per side, starting at the first corner:
    a Coons cap over it lines up with the quad's sides."""
    C = np.asarray(corners, float)
    out = []
    for a in range(4):
        p, q = C[a], C[(a + 1) % 4]
        for t in np.linspace(0, 1, k + 1)[:-1]:
            out.append(p + (q - p) * t)
    return np.array(out)