# Minimal software renderer for checking an OBJ + MTL:  perspective camera, flat shading with one key
# light and a fill light, z-buffer, translucent materials blended in a second pass.
# Usage:  python3 util/previewObj.py model.obj out.png [--yaw 30] [--pitch 15] [--size 1200x800]
import argparse
import os
import numpy as np
from PIL import Image

parser = argparse.ArgumentParser()
parser.add_argument('obj')
parser.add_argument('out')
parser.add_argument('--yaw', type=float, default=30.0)
parser.add_argument('--pitch', type=float, default=15.0)
parser.add_argument('--size', default='1200x800')
parser.add_argument('--background', default='#e8e8ec')
parser.add_argument('--distance', type=float, default=1.35, help='camera distance as a multiple of the model radius; large values approach orthographic')
args = parser.parse_args()
w, h = (int(v) for v in args.size.split('x'))


def loadMtl(path):
    mats = {}
    current = None
    for line in open(path):
        parts = line.split()
        if not parts:
            continue
        if parts[0] == 'newmtl':
            current = {'Kd': (0.8, 0.8, 0.8), 'Ks': (0.0, 0.0, 0.0), 'Ns': 10.0, 'd': 1.0}
            mats[parts[1]] = current
        elif current is not None and parts[0] in ('Kd', 'Ks'):
            current[parts[0]] = tuple(float(v) for v in parts[1:4])
        elif current is not None and parts[0] in ('Ns', 'd'):
            current[parts[0]] = float(parts[1])
    return mats


verts = []
tris = []
mats = {}
material = None
for line in open(args.obj):
    parts = line.split()
    if not parts:
        continue
    if parts[0] == 'mtllib':
        mats = loadMtl(os.path.join(os.path.dirname(args.obj), parts[1]))
    elif parts[0] == 'v':
        verts.append([float(v) for v in parts[1:4]])
    elif parts[0] == 'usemtl':
        material = parts[1]
    elif parts[0] == 'f':
        idx = [int(p.split('/')[0]) - 1 for p in parts[1:]]
        for i in range(1, len(idx) - 1):
            tris.append((idx[0], idx[i], idx[i + 1], material))
verts = np.array(verts)
print('%d vertices, %d triangles' % (len(verts), len(tris)))

# Camera orbiting the model centre.
centre = (verts.min(axis=0) + verts.max(axis=0)) / 2
radius = np.linalg.norm(verts.max(axis=0) - verts.min(axis=0)) / 2
yaw = np.radians(args.yaw)
pitch = np.radians(args.pitch)
eye = centre + radius * args.distance * np.array([np.sin(yaw) * np.cos(pitch), np.sin(pitch), np.cos(yaw) * np.cos(pitch)])
forward = centre - eye
forward = forward / np.linalg.norm(forward)
right = np.cross(forward, [0.0, 1.0, 0.0])
right = right / np.linalg.norm(right)
up = np.cross(right, forward)
focal = 1.6 * args.distance / 1.35
keyLight = np.array([-0.4, 0.8, 0.6])
keyLight = keyLight / np.linalg.norm(keyLight)
fillLight = np.array([0.6, 0.2, 0.5])
fillLight = fillLight / np.linalg.norm(fillLight)


def project(p):
    d = p - eye
    x = d @ right
    y = d @ up
    z = d @ forward
    return np.array([w / 2 + x / z * focal * h / 2, h / 2 - y / z * focal * h / 2, z])


bg = tuple(int(args.background[i:i + 2], 16) for i in (1, 3, 5))
colour = np.zeros((h, w, 3), dtype=float) + np.array(bg, dtype=float)
zbuf = np.full((h, w), np.inf)


def shade(a, b, c, m):
    n = np.cross(b - a, c - a)
    n = n / max(1e-9, np.linalg.norm(n))
    view = eye - (a + b + c) / 3
    view = view / np.linalg.norm(view)
    kd = np.array(m['Kd'])
    ks = np.array(m['Ks'])
    diffuse = max(0.0, n @ keyLight) * 0.85 + max(0.0, n @ fillLight) * 0.35
    half = keyLight + view
    half = half / np.linalg.norm(half)
    spec = max(0.0, n @ half) ** (m['Ns'] / 4) * 0.6
    return np.clip((kd * (0.18 + diffuse) + ks * spec) * 255, 0, 255)


def raster(tri, write, alpha):
    a, b, c, name = tri
    m = mats.get(name, {'Kd': (0.8, 0.8, 0.8), 'Ks': (0, 0, 0), 'Ns': 10.0, 'd': 1.0})
    pa, pb, pc = project(verts[a]), project(verts[b]), project(verts[c])
    if min(pa[2], pb[2], pc[2]) <= 0:
        return
    xs = [pa[0], pb[0], pc[0]]
    ys = [pa[1], pb[1], pc[1]]
    xMin, xMax = max(0, int(min(xs))), min(w - 1, int(max(xs)) + 1)
    yMin, yMax = max(0, int(min(ys))), min(h - 1, int(max(ys)) + 1)
    if xMin >= xMax or yMin >= yMax:
        return
    gx, gy = np.meshgrid(np.arange(xMin, xMax + 1) + 0.5, np.arange(yMin, yMax + 1) + 0.5)
    det = (pb[0] - pa[0]) * (pc[1] - pa[1]) - (pc[0] - pa[0]) * (pb[1] - pa[1])
    if abs(det) < 1e-9:
        return
    l1 = ((pb[0] - gx) * (pc[1] - gy) - (pc[0] - gx) * (pb[1] - gy)) / det
    l2 = ((pc[0] - gx) * (pa[1] - gy) - (pa[0] - gx) * (pc[1] - gy)) / det
    l3 = 1 - l1 - l2
    inside = (l1 >= 0) & (l2 >= 0) & (l3 >= 0)
    if not inside.any():
        return
    depth = l1 * pa[2] + l2 * pb[2] + l3 * pc[2]
    region = zbuf[yMin:yMax + 1, xMin:xMax + 1]
    visible = inside & (depth < region)
    if not visible.any():
        return
    rgb = shade(verts[a], verts[b], verts[c], m)
    target = colour[yMin:yMax + 1, xMin:xMax + 1]
    target[visible] = target[visible] * (1 - alpha) + rgb * alpha
    if write:
        region[visible] = depth[visible]


opaque = [t for t in tris if mats.get(t[3], {'d': 1.0})['d'] >= 1.0]
clear = [t for t in tris if mats.get(t[3], {'d': 1.0})['d'] < 1.0]
for t in opaque:
    raster(t, True, 1.0)
# Translucent faces back to front, against the opaque depth buffer.
clear.sort(key=lambda t: -np.mean([project(verts[i])[2] for i in t[:3]]))
for t in clear:
    raster(t, False, mats[t[3]]['d'])
for t in clear:
    raster(t, True, 0.0)
Image.fromarray(colour.astype(np.uint8)).save(args.out)
print('wrote', args.out)