"""
Tiny dependency-light renderer (numpy + Pillow) for STL meshes.

Point-splat rasterisation: every triangle is sampled in proportion to its
projected area, the nearest sample wins each pixel, then the image is shaded
with a key light (with a shadow map), a fill light, a rim light, specular
highlights and a soft contact shadow on an optional ground plane.

World frame: X right, Y up, Z towards the default viewer (board frame of the
Coral Eye files when the device stands upright).
"""
import struct
import numpy as np
np.seterr(all="ignore")
from PIL import Image, ImageFilter


# --------------------------------------------------------------------------
def load_stl(path):
    with open(path, "rb") as f:
        head = f.read(80)
        n = struct.unpack("<I", f.read(4))[0]
        dt = np.dtype([("n", "<f4", 3), ("v", "<f4", (3, 3)), ("a", "<u2")])
        data = np.fromfile(f, dtype=dt, count=n)
    if len(data) != n:
        raise ValueError("not a binary STL: " + path)
    return data["v"].astype(np.float64)


def rot_x(deg):
    a = np.radians(deg); c, s = np.cos(a), np.sin(a)
    return np.array([[1, 0, 0], [0, c, -s], [0, s, c]])


def rot_y(deg):
    a = np.radians(deg); c, s = np.cos(a), np.sin(a)
    return np.array([[c, 0, s], [0, 1, 0], [-s, 0, c]])


def rot_z(deg):
    a = np.radians(deg); c, s = np.cos(a), np.sin(a)
    return np.array([[c, -s, 0], [s, c, 0], [0, 0, 1]])


class Mesh:
    def __init__(self, tris, color, spec=0.25, gloss=40.0, name=""):
        self.tris = np.asarray(tris, float)          # (n,3,3)
        self.color = np.array(color, float) / 255.0
        self.spec, self.gloss, self.name = spec, gloss, name

    def transformed(self, R=np.eye(3), t=(0, 0, 0)):
        T = self.tris @ np.asarray(R).T + np.asarray(t, float)
        return Mesh(T, self.color * 255, self.spec, self.gloss, self.name)


# --------------------------------------------------------------------------
class Camera:
    def __init__(self, eye, target, up=(0, 1, 0), fov=30.0, width=1200, height=900, ortho_scale=None):
        self.eye, self.target = np.array(eye, float), np.array(target, float)
        f = self.target - self.eye; f /= np.linalg.norm(f)
        r = np.cross(f, up); r /= np.linalg.norm(r)
        u = np.cross(r, f)
        self.f, self.r, self.u = f, r, u
        self.fov, self.W, self.H = fov, width, height
        self.ortho = ortho_scale

    def project(self, P):
        d = P - self.eye
        x, y, z = d @ self.r, d @ self.u, d @ self.f
        if self.ortho:
            s = self.H / self.ortho
            sx = self.W / 2 + x * s
            sy = self.H / 2 - y * s
        else:
            fpx = (self.H / 2) / np.tan(np.radians(self.fov) / 2)
            zz = np.maximum(z, 1e-6)
            sx = self.W / 2 + fpx * x / zz
            sy = self.H / 2 - fpx * y / zz
        return sx, sy, z


def _raster(T, cam, W, H, id_offset=0):
    """Exact coverage rasteriser. Returns (pixel index, depth, triangle id) of the nearest hit per pixel."""
    sx, sy, z = cam.project(T.reshape(-1, 3))
    sx, sy, z = sx.reshape(-1, 3), sy.reshape(-1, 3), z.reshape(-1, 3)
    ok = (z > 1e-3).all(1) & np.isfinite(sx).all(1) & np.isfinite(sy).all(1)
    x0 = np.floor(sx.min(1) - 0.5).astype(np.int64) + 1
    y0 = np.floor(sy.min(1) - 0.5).astype(np.int64) + 1
    x1 = np.floor(sx.max(1) - 0.5).astype(np.int64)
    y1 = np.floor(sy.max(1) - 0.5).astype(np.int64)
    ok &= (x1 >= 0) & (y1 >= 0) & (x0 < W) & (y0 < H)
    x0, y0 = np.clip(x0, 0, W - 1), np.clip(y0, 0, H - 1)
    x1, y1 = np.clip(x1, 0, W - 1), np.clip(y1, 0, H - 1)
    bw, bh = x1 - x0 + 1, y1 - y0 + 1
    ok &= (bw > 0) & (bh > 0)
    size = np.maximum(bw, bh)
    out_lin, out_z, out_t = [], [], []

    def emit(ids, k):
        if len(ids) == 0:
            return
        gy, gx = np.mgrid[0:k, 0:k]
        gx, gy = gx.ravel(), gy.ravel()
        px = x0[ids, None] + gx[None, :]
        py = y0[ids, None] + gy[None, :]
        cx, cy = px + 0.5, py + 0.5
        X, Y = sx[ids], sy[ids]
        d = (X[:, 1] - X[:, 0]) * (Y[:, 2] - Y[:, 0]) - (X[:, 2] - X[:, 0]) * (Y[:, 1] - Y[:, 0])
        good = np.abs(d) > 1e-12
        ids, px, py, cx, cy, X, Y, d = ids[good], px[good], py[good], cx[good], cy[good], X[good], Y[good], d[good]
        w0 = ((X[:, 1, None] - cx) * (Y[:, 2, None] - cy) - (X[:, 2, None] - cx) * (Y[:, 1, None] - cy)) / d[:, None]
        w1 = ((X[:, 2, None] - cx) * (Y[:, 0, None] - cy) - (X[:, 0, None] - cx) * (Y[:, 2, None] - cy)) / d[:, None]
        w2 = 1 - w0 - w1
        e = -1e-9
        inside = (w0 >= e) & (w1 >= e) & (w2 >= e) & (px <= x1[ids, None]) & (py <= y1[ids, None])
        Z = z[ids]
        if cam.ortho:
            dep = w0 * Z[:, 0, None] + w1 * Z[:, 1, None] + w2 * Z[:, 2, None]
        else:
            dep = 1.0 / (w0 / Z[:, 0, None] + w1 / Z[:, 1, None] + w2 / Z[:, 2, None])
        r, c = np.nonzero(inside)
        out_lin.append(py[r, c] * W + px[r, c])
        out_z.append(dep[r, c])
        out_t.append(ids[r] + id_offset)

    idx_all = np.nonzero(ok)[0]
    for k in (1, 2, 4, 8, 16, 32, 64):
        sel = idx_all[size[idx_all] <= k]
        idx_all = idx_all[size[idx_all] > k]
        step = max(1, 4_000_000 // (k * k))
        for s in range(0, len(sel), step):
            emit(sel[s:s + step], k)
    for i in idx_all:                         # the few big ones, one by one in tiles
        kx, ky = bw[i], bh[i]
        k = int(max(kx, ky))
        emit(np.array([i]), k)
    if not out_lin:
        return np.zeros(0, np.int64), np.zeros(0), np.zeros(0, np.int64)
    lin, dep, tid = np.concatenate(out_lin), np.concatenate(out_z), np.concatenate(out_t)
    order = np.lexsort((dep, lin))
    lin_s = lin[order]
    first = np.ones(len(lin_s), bool)
    first[1:] = lin_s[1:] != lin_s[:-1]
    sel = order[first]
    return lin[sel], dep[sel], tid[sel]


def _pixel_world(cam, lin, depth, W, H):
    ys, xs = np.divmod(lin, W)
    if cam.ortho:
        s = H / cam.ortho
        dx = (xs + 0.5 - W / 2) / s
        dy = -(ys + 0.5 - H / 2) / s
        return cam.eye + np.outer(dx, cam.r) + np.outer(dy, cam.u) + np.outer(depth, cam.f)
    fpx = (H / 2) / np.tan(np.radians(cam.fov) / 2)
    dx = (xs + 0.5 - W / 2) / fpx
    dy = -(ys + 0.5 - H / 2) / fpx
    return cam.eye + np.outer(depth, cam.f) + np.outer(depth * dx, cam.r) + np.outer(depth * dy, cam.u)


def _ssao(wpos, N, vis, H, W, cam, world_r=3.5):
    """Screen-space ambient occlusion from world positions and normals."""
    P_img = wpos.reshape(H, W, 3)
    N_img = np.zeros((H * W, 3)); N_img[vis] = N; N_img = N_img.reshape(H, W, 3)
    valid = vis.reshape(H, W)
    dist = np.linalg.norm(cam.target - cam.eye)
    if cam.ortho:
        px_per_mm = H / cam.ortho
    else:
        px_per_mm = (H / 2) / np.tan(np.radians(cam.fov) / 2) / dist
    occ = np.zeros((H, W)); cnt = 0
    for frac in (0.25, 0.6, 1.0):
        r = max(2, int(round(world_r * frac * px_per_mm)))
        for k in range(8):
            a = 2 * np.pi * k / 8 + frac * 1.7
            dx, dy = int(round(r * np.cos(a))), int(round(r * np.sin(a)))
            Pn = np.roll(np.roll(P_img, dy, 0), dx, 1)
            vn = np.roll(np.roll(valid, dy, 0), dx, 1)
            D = Pn - P_img
            d = np.linalg.norm(D, axis=2) + 1e-6
            c = (D * N_img).sum(2) / d
            occ += np.clip(c - 0.12, 0, 1) * np.clip(1 - d / world_r, 0, 1) * vn
            cnt += 1
    ao = np.clip(1 - 2.2 * occ / cnt, 0.3, 1.0)
    return ao.reshape(-1)[vis]


def render(meshes, cam, ss=2, density=1.6, ground_y=None, bg_top=(246, 247, 249), bg_bottom=(222, 225, 230),
           key_dir=(-0.45, 0.8, 0.55), shadow=True, outline=0.0, out=None, ssao=True):
    W, H = cam.W * ss, cam.H * ss
    camS = Camera(cam.eye, cam.target, cam.u, cam.fov, W, H, cam.ortho)
    tris = np.concatenate([m.tris for m in meshes])
    owner = np.concatenate([np.full(len(m.tris), i) for i, m in enumerate(meshes)])
    nrm = np.cross(tris[:, 1] - tris[:, 0], tris[:, 2] - tris[:, 0])
    nrm /= np.linalg.norm(nrm, axis=1, keepdims=True) + 1e-12

    zb = np.full(W * H, np.inf); tid = np.full(W * H, -1, np.int64); wpos = np.zeros((W * H, 3))
    chunk = 200_000
    for s in range(0, len(tris), chunk):
        lin, z, t = _raster(tris[s:s + chunk], camS, W, H, id_offset=s)
        better = z < zb[lin]
        zb[lin[better]], tid[lin[better]] = z[better], t[better]
    vis_lin = np.nonzero(tid >= 0)[0]
    wpos[vis_lin] = _pixel_world(camS, vis_lin, zb[vis_lin], W, H)

    # ---- ground plane ----
    rays = None
    if ground_y is not None:
        ys, xs = np.mgrid[0:H, 0:W]
        if camS.ortho:
            s = H / camS.ortho
            dx = (xs.ravel() + 0.5 - W / 2) / s
            dy = -(ys.ravel() + 0.5 - H / 2) / s
            origins = camS.eye + np.outer(dx, camS.r) + np.outer(dy, camS.u)
            dirs = np.tile(camS.f, (W * H, 1))
        else:
            fpx = (H / 2) / np.tan(np.radians(camS.fov) / 2)
            dx = (xs.ravel() + 0.5 - W / 2) / fpx
            dy = -(ys.ravel() + 0.5 - H / 2) / fpx
            dirs = camS.f + np.outer(dx, camS.r) + np.outer(dy, camS.u)
            dirs /= np.linalg.norm(dirs, axis=1, keepdims=True)
            origins = np.tile(camS.eye, (W * H, 1))
        tg = (ground_y - origins[:, 1]) / np.where(np.abs(dirs[:, 1]) < 1e-9, -1e-9, dirs[:, 1])
        gpos = origins + dirs * tg[:, None]
        gz = (gpos - camS.eye) @ camS.f
        is_ground = (tg > 0) & (gz < zb)
        rays = (is_ground, gpos)

    # ---- background ----
    t = np.linspace(0, 1, H)[:, None, None]
    img = (np.array(bg_top) * (1 - t) + np.array(bg_bottom) * t) / 255.0
    img = np.broadcast_to(img, (H, W, 3)).copy().reshape(-1, 3)

    key = np.array(key_dir, float); key /= np.linalg.norm(key)
    fill = np.array([0.7, 0.25, 0.6]); fill /= np.linalg.norm(fill)
    rim = np.array([0.1, 0.5, -0.9]); rim /= np.linalg.norm(rim)

    # ---- shadow map from the key light (orthographic) ----
    def shadow_factor(P):
        if not shadow:
            return np.ones(len(P))
        center = tris.reshape(-1, 3).mean(0)
        ext = np.ptp(tris.reshape(-1, 3), axis=0).max() * 1.6
        if ground_y is not None:
            ext *= 1.6
        SM = 1600
        lcam = Camera(center + key * 500, center, up=(0, 0, 1) if abs(key[1]) > 0.95 else (0, 1, 0),
                      width=SM, height=SM, ortho_scale=ext)
        sz = np.full(SM * SM, np.inf)
        for s in range(0, len(tris), chunk):
            lin, z, _ = _raster(tris[s:s + chunk], lcam, SM, SM)
            np.minimum.at(sz, lin, z)
        sx, sy, z = lcam.project(P)
        lit = np.zeros(len(P))
        taps = 0
        for ox in (-2.0, 0, 2.0):
            for oy in (-2.0, 0, 2.0):
                ix = np.clip(np.floor(sx + ox).astype(int), 0, SM - 1)
                iy = np.clip(np.floor(sy + oy).astype(int), 0, SM - 1)
                lit += (z <= sz[iy * SM + ix] + 0.9)
                taps += 1
        return lit / taps

    vis = tid >= 0
    if vis.any():
        t_v = tid[vis]
        N = nrm[t_v].copy()
        P = wpos[vis]
        V = camS.eye - P; V /= np.linalg.norm(V, axis=1, keepdims=True)
        flipn = (N * V).sum(1) < 0
        N[flipn] *= -1
        own = owner[t_v]
        base = np.array([meshes[i].color for i in range(len(meshes))])[own]
        spec = np.array([meshes[i].spec for i in range(len(meshes))])[own]
        gloss = np.array([meshes[i].gloss for i in range(len(meshes))])[own]
        sh = shadow_factor(P + N * 0.4)
        ao = _ssao(wpos, N, vis, H, W, camS) if ssao else np.ones(len(P))
        dk = np.clip((N * key).sum(1), 0, 1)
        df = np.clip((N * fill).sum(1), 0, 1)
        dr = np.clip((N * rim).sum(1), 0, 1)
        hemi = 0.5 + 0.5 * N[:, 1]
        Hk = key + V; Hk /= np.linalg.norm(Hk, axis=1, keepdims=True)
        sp = spec * np.clip((N * Hk).sum(1), 0, 1) ** gloss * sh
        light = (0.30 * hemi + 0.18) * ao + (0.78 * dk * sh + 0.22 * df + 0.18 * dr) * (0.55 + 0.45 * ao)
        col = base * light[:, None] + sp[:, None]
        fres = (1 - np.clip((N * V).sum(1), 0, 1)) ** 4 * 0.12
        col = col + fres[:, None]
        img[vis] = np.clip(col, 0, 1)

    if rays is not None:
        is_ground, gpos = rays
        gp = gpos[is_ground]
        sh = shadow_factor(gp + np.array([0, 0.05, 0]))
        # contact darkening: distance to the object's footprint
        base_col = img[is_ground]
        shade = 0.55 + 0.45 * sh
        img[is_ground] = base_col * shade[:, None]

    img = img.reshape(H, W, 3)
    if outline > 0:
        z2 = zb.reshape(H, W).copy()
        z2[~np.isfinite(z2)] = 1e6
        gy, gx = np.gradient(np.log(z2 + 1))
        edge = (np.hypot(gx, gy) > outline).astype(float)
        idm = tid.reshape(H, W)
        e2 = np.zeros_like(edge)
        e2[:, 1:] += (idm[:, 1:] != idm[:, :-1]) & (np.abs(nrm[idm[:, 1:]] - nrm[idm[:, :-1]]).sum(-1) > 0.6)
        e2[1:, :] += (idm[1:, :] != idm[:-1, :]) & (np.abs(nrm[idm[1:, :]] - nrm[idm[:-1, :]]).sum(-1) > 0.6)
        edge = np.clip(edge + (e2 > 0), 0, 1)
        img = img * (1 - 0.55 * edge[..., None])

    im = Image.fromarray((np.clip(img, 0, 1) * 255).astype(np.uint8))
    if ss > 1:
        im = im.resize((cam.W, cam.H), Image.LANCZOS)
    if out:
        im.save(out)
    return im