"""A small z-buffer rasteriser in numpy: triangles through the picture's camera into a class (or id) image.

raster(cam, tris (n,3,3), ids (n,), scale=1) -> id image (H, W) with -1 for empty, depth image
Used by the fitting overlays (no Blender needed).
"""
import numpy as np
from tg_cam import project, W_IMG, H_IMG


def raster(cam, tris, ids, scale=1.0):
    W, H = int(round(W_IMG * scale)), int(round(H_IMG * scale))
    n = len(tris)
    uv, z = project(cam, tris.reshape(-1, 3))
    uv = uv.reshape(n, 3, 2) * scale
    z = z.reshape(n, 3)
    img = np.full((H, W), -1, np.int32)
    dep = np.full((H, W), np.inf)
    ok = (z > 0.05).all(1)
    for t in np.nonzero(ok)[0]:
        (x0, y0), (x1, y1), (x2, y2) = uv[t]
        xmin, xmax = int(max(0, np.floor(min(x0, x1, x2)))), int(min(W - 1, np.ceil(max(x0, x1, x2))))
        ymin, ymax = int(max(0, np.floor(min(y0, y1, y2)))), int(min(H - 1, np.ceil(max(y0, y1, y2))))
        if xmin > xmax or ymin > ymax:
            continue
        den = (y1 - y2) * (x0 - x2) + (x2 - x1) * (y0 - y2)
        if abs(den) < 1e-12:
            continue
        ys, xs = np.mgrid[ymin:ymax + 1, xmin:xmax + 1]
        px, py = xs + 0.5, ys + 0.5
        a = ((y1 - y2) * (px - x2) + (x2 - x1) * (py - y2)) / den
        b = ((y2 - y0) * (px - x2) + (x0 - x2) * (py - y2)) / den
        c = 1 - a - b
        inside = (a >= -1e-9) & (b >= -1e-9) & (c >= -1e-9)
        if not inside.any():
            continue
        zi = 1.0 / (a / z[t, 0] + b / z[t, 1] + c / z[t, 2])      # perspective-correct depth
        sub = dep[ymin:ymax + 1, xmin:xmax + 1]
        win = inside & (zi < sub)
        sub[win] = zi[win]
        img[ymin:ymax + 1, xmin:xmax + 1][win] = ids[t]
    return img, dep


def edges_of(img):
    """pixels where the id changes (to the right or below)"""
    e = np.zeros(img.shape, bool)
    e[:, :-1] |= img[:, :-1] != img[:, 1:]
    e[:-1, :] |= img[:-1, :] != img[1:, :]
    return e