"""Small isometric mesh renderer, so the assembly diagram is drawn from the
real geometry rather than an approximation of it."""
import numpy as np
from matplotlib.collections import PolyCollection


def _rot(azim_deg, elev_deg):
    a, e = np.radians(azim_deg), np.radians(elev_deg)
    rz = np.array([[np.cos(a), -np.sin(a), 0], [np.sin(a), np.cos(a), 0], [0, 0, 1]])
    # negative pitch so world +Z ends up pointing UP on screen; with +e the view
    # is correct but vertically mirrored, which silently flips exploded views
    rx = np.array([[1, 0, 0], [0, np.cos(-e), -np.sin(-e)], [0, np.sin(-e), np.cos(-e)]])
    return rx @ rz


def project(pts, azim=32, elev=24):
    """World -> screen. Returns (xy, depth); larger depth is nearer the camera."""
    p = np.asarray(pts, dtype=float) @ _rot(azim, elev).T
    return p[..., :2], p[..., 2]


def draw_mesh(ax, mesh, colour, offset=(0, 0, 0), azim=32, elev=24,
              light=(-0.35, -0.5, 0.78), zorder=1, alpha=1.0, edge=None):
    """Flat-shaded painter's-algorithm render of one mesh."""
    R = _rot(azim, elev)
    tri = mesh.triangles

    # Recompute normals from the vertices. Boolean output carries sliver
    # triangles whose stored normal is arbitrary, and those slivers shade at
    # random -- which is what put the cobweb streaks across the flat faces.
    e1 = tri[:, 1] - tri[:, 0]
    e2 = tri[:, 2] - tri[:, 0]
    n = np.cross(e1, e2)
    ln = np.linalg.norm(n, axis=1)
    good = ln > 1e-9
    n = np.divide(n, np.where(ln[:, None] == 0, 1, ln[:, None]))

    # backface culling -- without it you see straight through the part
    n_cam = n @ R.T
    keep = good & (n_cam[:, 2] > 1e-6)
    tris = (tri + np.asarray(offset, dtype=float))[keep]
    n = n[keep]
    if len(tris) == 0:
        return np.zeros((0, 2))
    L = np.asarray(light, dtype=float); L /= np.linalg.norm(L)

    # face-on surfaces read lightest; near-vertical walls get a soft falloff
    shade = 0.55 + 0.45 * np.clip(n @ L, 0, 1)
    base = np.array(_hex(colour), dtype=float)
    cols = np.clip(base[None, :] * shade[:, None], 0, 1)

    flat = tris.reshape(-1, 3)
    xy, d = project(flat, azim, elev)
    xy = xy.reshape(-1, 3, 2)
    depth = d.reshape(-1, 3).mean(axis=1)

    order = np.argsort(depth)
    # Stroke each triangle in its own fill colour. With edgecolors='none' the
    # antialiased seams between adjacent triangles read as a faint cobweb across
    # what should be one flat face.
    pc = PolyCollection(xy[order], facecolors=cols[order], alpha=alpha,
                        edgecolors=cols[order] if edge is None else edge,
                        linewidths=0.5 if edge is None else 0.15, zorder=zorder)
    ax.add_collection(pc)
    return xy.reshape(-1, 2)


def _hex(h):
    h = h.lstrip('#')
    return tuple(int(h[i:i + 2], 16) / 255 for i in (0, 2, 4))


def draw_group(ax, items, azim=32, elev=24, light=(-0.35, -0.5, 0.78),
               zorder=1, alpha=1.0):
    """Render several meshes as ONE depth-sorted pass.

    draw_mesh per mesh sorts only within that mesh, so a second collection lands
    entirely on top of the first -- which made the grey core paint over the
    yellow face plate that is physically in front of it.
    items: [(mesh, colour_hex, offset)]  (offset optional)
    """
    R = _rot(azim, elev)
    L = np.asarray(light, dtype=float); L /= np.linalg.norm(L)
    XY, COL, D = [], [], []
    for it in items:
        mesh, colour = it[0], it[1]
        offset = it[2] if len(it) > 2 else (0, 0, 0)
        tri = mesh.triangles + np.asarray(offset, dtype=float)
        e1 = tri[:, 1] - tri[:, 0]
        e2 = tri[:, 2] - tri[:, 0]
        n = np.cross(e1, e2)
        ln = np.linalg.norm(n, axis=1)
        good = ln > 1e-9
        n = np.divide(n, np.where(ln[:, None] == 0, 1, ln[:, None]))
        n_cam = n @ R.T
        keep = good & (n_cam[:, 2] > 1e-6)
        tri, n = tri[keep], n[keep]
        if not len(tri):
            continue
        shade = 0.55 + 0.45 * np.clip(n @ L, 0, 1)
        base = np.array(_hex(colour), dtype=float)
        COL.append(np.clip(base[None, :] * shade[:, None], 0, 1))
        xy, d = project(tri.reshape(-1, 3), azim, elev)
        XY.append(xy.reshape(-1, 3, 2))
        D.append(d.reshape(-1, 3).mean(axis=1))
    if not XY:
        return
    xy = np.concatenate(XY); cols = np.concatenate(COL); depth = np.concatenate(D)
    order = np.argsort(depth)
    ax.add_collection(PolyCollection(xy[order], facecolors=cols[order], alpha=alpha,
                                     edgecolors=cols[order], linewidths=0.5,
                                     zorder=zorder))


def raster(items, extent, px=900, azim=32, elev=24, light=(-0.35, -0.5, 0.78),
           ss=2, ambient=0.55):
    """True z-buffer render of flat-shaded triangles -> RGBA array.

    The painter's algorithm cannot cope with this geometry: the core slab's top
    face is a handful of very large triangles whose centroid depth beats the
    small face-plate triangles sitting 0.6 mm in front of them, so the grey core
    paints straight over the yellow plate. A depth buffer is the only honest fix,
    and it also kills the antialiasing seams for free.

    items: [(mesh, colour_hex, offset)]; extent: (x0, x1, y0, y1) in world units.
    """
    x0, x1, y0, y1 = extent
    pw = int(px)
    ph = max(1, int(round(px * (y1 - y0) / (x1 - x0))))
    W, H = pw * ss, ph * ss
    zbuf = np.full((H, W), -np.inf)
    img = np.zeros((H, W, 3))
    hit = np.zeros((H, W), dtype=bool)
    R = _rot(azim, elev)
    L = np.asarray(light, dtype=float); L /= np.linalg.norm(L)

    for it in items:
        mesh, colour = it[0], it[1]
        offset = it[2] if len(it) > 2 else (0, 0, 0)
        tri = mesh.triangles + np.asarray(offset, dtype=float)
        e1, e2 = tri[:, 1] - tri[:, 0], tri[:, 2] - tri[:, 0]
        n = np.cross(e1, e2)
        ln = np.linalg.norm(n, axis=1)
        keep = ln > 1e-9
        n = n[keep] / ln[keep][:, None]
        tri = tri[keep]
        n_cam = n @ R.T
        front = n_cam[:, 2] > 1e-6
        tri, n = tri[front], n[front]
        if not len(tri):
            continue
        shade = ambient + (1 - ambient) * np.clip(n @ L, 0, 1)
        base = np.array(_hex(colour), dtype=float)
        cols = np.clip(base[None, :] * shade[:, None], 0, 1)

        p = tri.reshape(-1, 3) @ R.T
        sx = (p[:, 0] - x0) / (x1 - x0) * (W - 1)
        sy = (y1 - p[:, 1]) / (y1 - y0) * (H - 1)     # screen y grows downward
        sz = p[:, 2]
        sx, sy, sz = sx.reshape(-1, 3), sy.reshape(-1, 3), sz.reshape(-1, 3)

        for k in range(len(sx)):
            ax_, ay_ = sx[k], sy[k]
            i0 = max(int(np.floor(ax_.min())), 0); i1 = min(int(np.ceil(ax_.max())), W - 1)
            j0 = max(int(np.floor(ay_.min())), 0); j1 = min(int(np.ceil(ay_.max())), H - 1)
            if i1 < i0 or j1 < j0:
                continue
            d = ((ay_[1] - ay_[2]) * (ax_[0] - ax_[2]) +
                 (ax_[2] - ax_[1]) * (ay_[0] - ay_[2]))
            if abs(d) < 1e-12:
                continue
            X, Y = np.meshgrid(np.arange(i0, i1 + 1), np.arange(j0, j1 + 1))
            l0 = ((ay_[1] - ay_[2]) * (X - ax_[2]) + (ax_[2] - ax_[1]) * (Y - ay_[2])) / d
            l1 = ((ay_[2] - ay_[0]) * (X - ax_[2]) + (ax_[0] - ax_[2]) * (Y - ay_[2])) / d
            l2 = 1.0 - l0 - l1
            m = (l0 >= -1e-9) & (l1 >= -1e-9) & (l2 >= -1e-9)
            if not m.any():
                continue
            z = l0 * sz[k, 0] + l1 * sz[k, 1] + l2 * sz[k, 2]
            sub = zbuf[j0:j1 + 1, i0:i1 + 1]
            w = m & (z > sub)
            if not w.any():
                continue
            sub[w] = z[w]
            img[j0:j1 + 1, i0:i1 + 1][w] = cols[k]
            hit[j0:j1 + 1, i0:i1 + 1][w] = True

    rgba = np.concatenate([img, hit[..., None].astype(float)], axis=2)
    if ss > 1:
        rgba = rgba.reshape(ph, ss, pw, ss, 4).mean(axis=(1, 3))
        a = rgba[..., 3:4]
        rgba[..., :3] = np.divide(rgba[..., :3], np.where(a == 0, 1, a))
    return rgba
