#!/usr/bin/env python3
"""
render_preview.py -- CPU cross-check of the viewer pipeline.

Consumes exactly what the WebGL viewer consumes (bundle.build_model output:
merged positions/normals/uvs/indices + decoded SVT texture) and software-rasterizes
one model to a PNG. If this looks right, the WebGL viewer -- fed identical data --
renders the same. Verification, not the product.

    python viewer/render_preview.py 0 out.png     # model index (default 0)
"""
import base64
import math
import os
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
sys.path.insert(0, os.path.join(os.path.dirname(HERE), "parser"))
import bundle
import svt

W, H = 640, 480


def mv(m, v):
    return [sum(m[r][c] * v[c] for c in range(4)) for r in range(4)]


def matmul(a, b):
    return [[sum(a[r][k] * b[k][c] for k in range(4)) for c in range(4)] for r in range(4)]


def look_at(eye, ctr, up):
    def sub(a, b): return [a[0]-b[0], a[1]-b[1], a[2]-b[2]]
    def crs(a, b): return [a[1]*b[2]-a[2]*b[1], a[2]*b[0]-a[0]*b[2], a[0]*b[1]-a[1]*b[0]]
    def nrm(a):
        l = math.hypot(*a) or 1.0; return [a[0]/l, a[1]/l, a[2]/l]
    def dot(a, b): return a[0]*b[0]+a[1]*b[1]+a[2]*b[2]
    z = nrm(sub(eye, ctr)); x = nrm(crs(up, z)); y = crs(z, x)
    return [[x[0], x[1], x[2], -dot(x, eye)],
            [y[0], y[1], y[2], -dot(y, eye)],
            [z[0], z[1], z[2], -dot(z, eye)],
            [0, 0, 0, 1]]


def persp(fovy, asp, n, f):
    t = 1/math.tan(fovy/2)
    return [[t/asp, 0, 0, 0], [0, t, 0, 0],
            [0, 0, (f+n)/(n-f), 2*f*n/(n-f)], [0, 0, -1, 0]]


def render(md, outpath, cam=None):
    # decode textures to (w,h,rgba bytes)
    texs = []
    for t in md["textures"]:
        texs.append((t["w"], t["h"], base64.b64decode(t["rgba"])))

    b = md["bounds"]
    ctr = [(b["min"][i]+b["max"][i])/2 for i in range(3)]
    rad = math.hypot(*(b["max"][i]-b["min"][i] for i in range(3)))/2 or 1.0
    up = [0, 0, 1] if md["up"] == "z" else [0, 1, 0]
    az, el, dist = 0.9, 0.55, rad*2.6
    ce = math.cos(el)
    if md["up"] == "z":
        eye = [ctr[0]+dist*ce*math.sin(az), ctr[1]+dist*ce*math.cos(az), ctr[2]+dist*math.sin(el)]
    else:
        eye = [ctr[0]+dist*ce*math.sin(az), ctr[1]+dist*math.sin(el), ctr[2]+dist*ce*math.cos(az)]
    if cam is not None:
        eye, ctr, up = cam["eye"], cam["center"], cam["up"]
        P = matmul(persp(1.0, W/H, 50, rad*5+2000), look_at(eye, ctr, up))
    else:
        P = matmul(persp(1.0, W/H, dist*0.02, dist*10+rad*4), look_at(eye, ctr, up))
    L = [0.4, 0.7, 0.6]; ll = math.hypot(*L); L = [c/ll for c in L]

    img = bytearray(b"\x1a\x14\x10" * (W*H))   # dark ground
    zbuf = [1e30]*(W*H)

    for g in md["groups"]:
        pos, nrm, uv, idx = g["positions"], g["normals"], g["uvs"], g["indices"]
        col = g["color"]; ti = g["texture"]
        tex = texs[ti] if ti >= 0 else None
        # project all verts
        sx = [0.0]*(len(pos)//3); sy = list(sx); sz = list(sx); sw = list(sx)
        for i in range(len(pos)//3):
            c = mv(P, [pos[i*3], pos[i*3+1], pos[i*3+2], 1.0])
            w = c[3] if abs(c[3]) > 1e-9 else 1e-9
            sw[i] = w
            sx[i] = (c[0]/w*0.5+0.5)*W
            sy[i] = (1-(c[1]/w*0.5+0.5))*H
            sz[i] = c[2]/w
        for t in range(0, len(idx), 3):
            a, bb, cc = idx[t], idx[t+1], idx[t+2]
            if sw[a] <= 0 or sw[bb] <= 0 or sw[cc] <= 0:
                continue
            x0, y0, x1, y1, x2, y2 = sx[a], sy[a], sx[bb], sy[bb], sx[cc], sy[cc]
            area = (x1-x0)*(y2-y0)-(x2-x0)*(y1-y0)
            if abs(area) < 1e-9:
                continue
            minx = max(0, int(min(x0, x1, x2))); maxx = min(W-1, int(max(x0, x1, x2))+1)
            miny = max(0, int(min(y0, y1, y2))); maxy = min(H-1, int(max(y0, y1, y2))+1)
            iw = [1.0/sw[a], 1.0/sw[bb], 1.0/sw[cc]]
            # world normals for lighting
            nl = []
            for vi in (a, bb, cc):
                nx, ny, nz = nrm[vi*3], nrm[vi*3+1], nrm[vi*3+2]
                nl.append(abs(nx*L[0]+ny*L[1]+nz*L[2]))
            for py in range(miny, maxy+1):
                for px in range(minx, maxx+1):
                    w0 = ((x1-px)*(y2-py)-(x2-px)*(y1-py))/area
                    w1 = ((x2-px)*(y0-py)-(x0-px)*(y2-py))/area
                    w2 = 1-w0-w1
                    if w0 < 0 or w1 < 0 or w2 < 0:
                        continue
                    z = w0*sz[a]+w1*sz[bb]+w2*sz[cc]
                    o = py*W+px
                    if z >= zbuf[o]:
                        continue
                    zbuf[o] = z
                    lit = 0.35+0.65*(w0*nl[0]+w1*nl[1]+w2*nl[2])
                    if tex:
                        # perspective-correct uv
                        pw = w0*iw[0]+w1*iw[1]+w2*iw[2]
                        u = (w0*uv[a*2]*iw[0]+w1*uv[bb*2]*iw[1]+w2*uv[cc*2]*iw[2])/pw
                        v = (w0*uv[a*2+1]*iw[0]+w1*uv[bb*2+1]*iw[1]+w2*uv[cc*2+1]*iw[2])/pw
                        tw, th, td = tex
                        tx = int((u % 1.0)*tw) % tw; tyv = int((v % 1.0)*th) % th
                        tp = (tyv*tw+tx)*4
                        r, gg, bl = td[tp], td[tp+1], td[tp+2]
                        r *= col[0]; gg *= col[1]; bl *= col[2]
                    else:
                        r, gg, bl = col[0]*255, col[1]*255, col[2]*255
                    # Division DAC gamma 1.7 (out = in^(1/1.25))
                    img[o*3] = min(255, int(((r*lit/255.0)**(1/1.25))*255))
                    img[o*3+1] = min(255, int(((gg*lit/255.0)**(1/1.25))*255))
                    img[o*3+2] = min(255, int(((bl*lit/255.0)**(1/1.25))*255))

    # write PNG (reuse svt's encoder, but it expects RGBA)
    rgba = bytearray(W*H*4)
    for i in range(W*H):
        rgba[i*4:i*4+3] = img[i*3:i*3+3]; rgba[i*4+3] = 255
    svt.write_png(outpath, W, H, bytes(rgba))
    print("wrote %s  (%dx%d, model=%r, %d tris)" % (outpath, W, H, md["name"], md["tris"]))


if __name__ == "__main__":
    i = int(sys.argv[1]) if len(sys.argv) > 1 else 0
    out = sys.argv[2] if len(sys.argv) > 2 else os.path.join(HERE, "preview_%d.png" % i)
    render(bundle.build_model(bundle.SHOWCASE[i]), out)