"""Measured numbers for the card, run on the local GPU from source."""
import os
import sys
import time

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)

import torch

import load_local

ptd = load_local.load()

print("GPU:", torch.cuda.get_device_name(0))


def quad(a, b, c, d):
    return [[a, b, c], [a, c, d]]


def add_box(verts, faces, mat, lo, hi, mat_id):
    x0, y0, z0 = lo
    x1, y1, z1 = hi
    base = len(verts)
    verts.extend([(x0, y0, z0), (x1, y0, z0), (x1, y0, z1), (x0, y0, z1),
                  (x0, y1, z0), (x1, y1, z0), (x1, y1, z1), (x0, y1, z1)])
    i = lambda k: base + k
    for q in [quad(i(0), i(1), i(2), i(3)), quad(i(4), i(7), i(6), i(5)),
              quad(i(0), i(4), i(5), i(1)), quad(i(3), i(2), i(6), i(7)),
              quad(i(0), i(3), i(7), i(4)), quad(i(1), i(5), i(6), i(2))]:
        faces.extend(q)
        mat.extend([mat_id, mat_id])


def cornell(with_box=True):
    verts, faces, mat = [], [], []

    def wall(a, b, c, d, m):
        base = len(verts)
        verts.extend([a, b, c, d])
        faces.extend(quad(base, base + 1, base + 2, base + 3))
        mat.extend([m, m])

    X, Y, Z = 5.56, 5.488, 5.592
    wall((0, 0, 0), (X, 0, 0), (X, 0, Z), (0, 0, Z), 0)
    wall((0, Y, 0), (0, Y, Z), (X, Y, Z), (X, Y, 0), 0)
    wall((0, 0, Z), (X, 0, Z), (X, Y, Z), (0, Y, Z), 0)
    wall((0, 0, 0), (0, 0, Z), (0, Y, Z), (0, Y, 0), 1)
    wall((X, 0, 0), (X, Y, 0), (X, Y, Z), (X, 0, Z), 2)
    wall((2.13, Y - 0.001, 2.27), (3.43, Y - 0.001, 2.27),
         (3.43, Y - 0.001, 3.32), (2.13, Y - 0.001, 3.32), 3)
    if with_box:
        add_box(verts, faces, mat, (1.3, 0.0, 0.65), (2.4, 1.65, 1.7), 0)
    albedo = [[0.73, 0.73, 0.73], [0.65, 0.05, 0.05], [0.12, 0.45, 0.15],
              [0.78, 0.78, 0.78]]
    emission = [[0, 0, 0], [0, 0, 0], [0, 0, 0], [18.4, 15.6, 8.0]]
    return ptd.Scene(verts, faces, mat, albedo=albedo, emission=emission)


cam = ptd.Camera(position=(2.78, 2.73, -8.0), look_at=(2.78, 2.73, 2.8),
                 vfov_deg=39.0)

# ---- furnace exactness
a, E, B = 0.5, 0.25, 8
verts, faces, mat = [], [], []
add_box(verts, faces, mat, (-2, -2, -2), (2, 2, 2), 0)
fs = ptd.Scene(verts, faces, mat, albedo=[[a] * 3], emission=[[E] * 3])
fc = ptd.Camera(position=(0, 0, 0), look_at=(0, 0, 1), vfov_deg=60.0)
img = ptd.render(fs, fc, 64, 64, spp=4, max_bounces=B, nee=False, seed=0)
exp = E * (1 - a ** B) / (1 - a)
print(f"furnace: analytic {exp:.9f}, max abs err {(img - exp).abs().max().item():.2e}")

# ---- estimator agreement
verts, faces, mat = [], [], []
add_box(verts, faces, mat, (0, 0, 0), (5, 5, 5), 0)
base = len(verts)
verts.extend([(1.0, 4.999, 1.0), (4.5, 4.999, 1.0), (4.5, 4.999, 4.5),
              (1.0, 4.999, 4.5)])
faces.extend(quad(base, base + 1, base + 2, base + 3))
mat.extend([1, 1])
es = ptd.Scene(verts, faces, mat, albedo=[[0.6] * 3, [0.0] * 3],
               emission=[[0] * 3, [3.0] * 3])
ec = ptd.Camera(position=(2.5, 2.5, 0.2), look_at=(2.5, 2.5, 5.0), vfov_deg=70.0)
m_nee = ptd.render(es, ec, 128, 128, spp=128, max_bounces=4, nee=True, seed=1).mean().item()
m_brdf = ptd.render(es, ec, 128, 128, spp=512, max_bounces=4, nee=False, seed=2).mean().item()
print(f"estimators: NEE {m_nee:.5f} vs BRDF {m_brdf:.5f}, rel diff {abs(m_nee-m_brdf)/m_brdf:.4%}")

# ---- FD gradient agreement
scene = cornell(with_box=False)
H = W = 48
torch.manual_seed(0)
Wr = torch.rand(H, W, 3, device="cuda")


def loss_of(alb, emi):
    scene.albedo = alb
    scene.emission = emi
    return (ptd.render(scene, cam, H, W, spp=8, max_bounces=3, seed=7) * Wr).sum()


alb0 = scene.albedo.clone()
emi0 = scene.emission.clone()
alb = alb0.clone().requires_grad_(True)
emi = emi0.clone().requires_grad_(True)
loss_of(alb, emi).backward()
worst = 0.0
for (i, c) in [(0, 0), (0, 1), (0, 2), (1, 0), (2, 1), (2, 2)]:
    h = 2e-3
    ap = alb0.clone(); ap[i, c] += h
    am = alb0.clone(); am[i, c] -= h
    fd = (loss_of(ap, emi0) - loss_of(am, emi0)).item() / (2 * h)
    worst = max(worst, abs(alb.grad[i, c].item() - fd) / abs(fd))
h = 0.05
for c in range(3):
    ep = emi0.clone(); ep[3, c] += h
    em = emi0.clone(); em[3, c] -= h
    fd = (loss_of(alb0, ep) - loss_of(alb0, em)).item() / (2 * h)
    worst = max(worst, abs(emi.grad[3, c].item() - fd) / abs(fd))
print(f"gradients vs central differences: max rel err {worst:.2e} over 9 entries")

# ---- inverse rendering
target_row = torch.tensor([0.2, 0.5, 0.7], device="cuda")
scene = cornell(with_box=False)
t_alb = scene.albedo.clone(); t_alb[1] = target_row
scene.albedo = t_alb
target = ptd.render(scene, cam, 48, 48, spp=8, max_bounces=3, seed=3).detach()
alb = t_alb.clone(); alb[1] = torch.tensor([0.5, 0.5, 0.5], device="cuda")
alb = alb.requires_grad_(True)
opt = torch.optim.Adam([alb], lr=0.05)
first = None
for it in range(80):
    opt.zero_grad()
    scene.albedo = alb
    loss = (ptd.render(scene, cam, 48, 48, spp=8, max_bounces=3, seed=3) - target).square().mean()
    loss.backward()
    alb.grad[0] = 0; alb.grad[2] = 0; alb.grad[3] = 0
    opt.step()
    with torch.no_grad():
        alb.clamp_(0.02, 0.98)
    if first is None:
        first = loss.item()
print(f"inverse: target {target_row.tolist()} recovered "
      f"{[round(x, 3) for x in alb[1].detach().tolist()]}, "
      f"max abs err {(alb[1].detach() - target_row).abs().max().item():.3f}, "
      f"loss {first:.2e} -> {loss.item():.2e} in 80 steps")

# ---- per-texel albedo
def cornell_tex(back_albedo):
    verts, faces, mat, uvs = [], [], [], []

    def wall(a, b, c, d, m, uv=False):
        base = len(verts)
        verts.extend([a, b, c, d])
        faces.extend(quad(base, base + 1, base + 2, base + 3))
        mat.extend([m, m])
        if uv:
            uvs.extend([[(0, 0), (1, 0), (1, 1)], [(0, 0), (1, 1), (0, 1)]])
        else:
            uvs.extend([[(0, 0)] * 3, [(0, 0)] * 3])

    X, Y, Z = 5.56, 5.488, 5.592
    wall((0, 0, 0), (X, 0, 0), (X, 0, Z), (0, 0, Z), 0)
    wall((0, Y, 0), (0, Y, Z), (X, Y, Z), (X, Y, 0), 0)
    wall((0, 0, Z), (X, 0, Z), (X, Y, Z), (0, Y, Z), 4, uv=True)
    wall((0, 0, 0), (0, 0, Z), (0, Y, Z), (0, Y, 0), 1)
    wall((X, 0, 0), (X, Y, 0), (X, Y, Z), (X, 0, Z), 2)
    wall((2.13, Y - 0.001, 2.27), (3.43, Y - 0.001, 2.27),
         (3.43, Y - 0.001, 3.32), (2.13, Y - 0.001, 3.32), 3)
    albedo = [torch.tensor([0.73, 0.73, 0.73]),
              torch.tensor([0.65, 0.05, 0.05]),
              torch.tensor([0.12, 0.45, 0.15]),
              torch.tensor([0.78, 0.78, 0.78]),
              back_albedo]
    emission = [[0, 0, 0], [0, 0, 0], [0, 0, 0], [18.4, 15.6, 8.0], [0, 0, 0]]
    return ptd.Scene(verts, faces, mat, albedo=albedo, emission=emission,
                     uvs=uvs)


torch.manual_seed(2)
base_tex = (torch.rand(4, 4, 3, device="cuda") * 0.6 + 0.2)
torch.manual_seed(0)
Wr48 = torch.rand(48, 48, 3, device="cuda")


def tex_loss(t):
    sc = cornell_tex(t)
    return (ptd.render(sc, cam, 48, 48, spp=8, max_bounces=3, seed=7) * Wr48).sum()


tt = base_tex.clone().requires_grad_(True)
tex_loss(tt).backward()
worst = 0.0
for (y, x, c) in [(0, 0, 0), (1, 2, 1), (3, 3, 2), (2, 1, 0), (0, 3, 1)]:
    h = 2e-3
    tp = base_tex.clone(); tp[y, x, c] += h
    tm = base_tex.clone(); tm[y, x, c] -= h
    fd = (tex_loss(tp) - tex_loss(tm)).item() / (2 * h)
    worst = max(worst, abs(tt.grad[y, x, c].item() - fd) / abs(fd))
print(f"texel gradients vs central differences: max rel err {worst:.2e} over 5 texels")

torch.manual_seed(4)
target_tex = (torch.rand(8, 8, 3, device="cuda") * 0.7 + 0.15)
target = ptd.render(cornell_tex(target_tex), cam, 64, 64, spp=8,
                    max_bounces=3, seed=3).detach()
t = torch.full((8, 8, 3), 0.5, device="cuda", requires_grad=True)
sc_opt = cornell_tex(t)
opt = torch.optim.Adam([t], lr=0.1)
first = None
for _ in range(200):
    opt.zero_grad()
    loss = (ptd.render(sc_opt, cam, 64, 64, spp=8, max_bounces=3, seed=3)
            - target).square().mean()
    loss.backward()
    opt.step()
    with torch.no_grad():
        t.clamp_(0.02, 0.98)
    if first is None:
        first = loss.item()
err = (t.detach() - target_tex).abs()
print(f"texture inverse: 8x8x3 (192 unknowns) recovered to mean abs err "
      f"{err.mean().item():.3f} (max {err.max().item():.3f}), "
      f"loss {first:.2e} -> {loss.item():.2e} in 200 steps")

# ---- Fresnel slab analytic
verts, faces, mat = [], [], []


def quad_w(a, b, c, d, m):
    base = len(verts)
    verts.extend([a, b, c, d])
    faces.extend(quad(base, base + 1, base + 2, base + 3))
    mat.extend([m, m])


quad_w((-10, -10, 5), (10, -10, 5), (10, 10, 5), (-10, 10, 5), 0)
add_box(verts, faces, mat, (-10, -10, 2.0), (10, 10, 2.2), 1)
slab = ptd.Scene(verts, faces, mat, albedo=[[0, 0, 0], [0, 0, 0]],
                 emission=[[1.0] * 3, [0, 0, 0]],
                 material_types=[ptd.DIFFUSE, ptd.DIELECTRIC], ior=[1.5, 1.5])
scam = ptd.Camera(position=(0, 0, 0), look_at=(0, 0, 1), vfov_deg=10.0)
img = ptd.render(slab, scam, 32, 32, spp=1024, max_bounces=10,
                 estimator="brdf", seed=5)
R = (0.5 / 2.5) ** 2
exp_T = (1 - R) / (1 + R)
got = img[12:20, 12:20].mean().item()
print(f"fresnel slab (ior 1.5): analytic T {exp_T:.6f}, measured {got:.6f}, "
      f"rel err {abs(got-exp_T)/exp_T:.4%}")

# ---- conductor F0 FD (through GGX + MIS)
H = W = 48
torch.manual_seed(0)
Wr = torch.rand(H, W, 3, device="cuda")
base_alb = torch.tensor([[0.73] * 3, [0.9, 0.5, 0.2], [0.12, 0.45, 0.15],
                         [0.78] * 3], device="cuda")


def cond_loss(alb):
    verts, faces, mat = [], [], []

    def wall(a, b, c, d, m):
        b0 = len(verts)
        verts.extend([a, b, c, d])
        faces.extend(quad(b0, b0 + 1, b0 + 2, b0 + 3))
        mat.extend([m, m])

    X, Y, Z = 5.56, 5.488, 5.592
    wall((0, 0, 0), (X, 0, 0), (X, 0, Z), (0, 0, Z), 0)
    wall((0, Y, 0), (0, Y, Z), (X, Y, Z), (X, Y, 0), 0)
    wall((0, 0, Z), (X, 0, Z), (X, Y, Z), (0, Y, Z), 0)
    wall((0, 0, 0), (0, 0, Z), (0, Y, Z), (0, Y, 0), 1)
    wall((X, 0, 0), (X, Y, 0), (X, Y, Z), (X, 0, Z), 2)
    wall((2.13, Y - 0.001, 2.27), (3.43, Y - 0.001, 2.27),
         (3.43, Y - 0.001, 3.32), (2.13, Y - 0.001, 3.32), 3)
    sc = ptd.Scene(verts, faces, mat, albedo=alb,
                   emission=[[0, 0, 0], [0, 0, 0], [0, 0, 0],
                             [18.4, 15.6, 8.0]],
                   material_types=[0, 1, 0, 0], roughness=[0.3, 0.2, 0.3, 0.3])
    return (ptd.render(sc, cam, H, W, spp=8, max_bounces=3, seed=13) * Wr).sum()


alb = base_alb.clone().requires_grad_(True)
cond_loss(alb).backward()
worst = 0.0
for (i, c) in [(1, 0), (1, 2), (0, 1)]:
    h = 2e-3
    ap = base_alb.clone(); ap[i, c] += h
    am = base_alb.clone(); am[i, c] -= h
    fd = (cond_loss(ap) - cond_loss(am)).item() / (2 * h)
    worst = max(worst, abs(alb.grad[i, c].item() - fd) / abs(fd))
print(f"conductor F0 gradients vs central differences (GGX+MIS): "
      f"max rel err {worst:.2e} over 3 entries")

# ---- env texel FD
torch.manual_seed(3)
env0 = torch.rand(4, 8, 3, device="cuda") * 1.5 + 0.2
verts, faces, mat = [], [], []
b0 = len(verts)
verts.extend([(-5, 0, -5), (5, 0, -5), (5, 0, 5), (-5, 0, 5)])
faces.extend(quad(b0, b0 + 1, b0 + 2, b0 + 3))
mat.extend([0, 0])
ecam = ptd.Camera(position=(0, 2.0, 0.01), look_at=(0, 0, 0.5), vfov_deg=40.0)
torch.manual_seed(0)
Wre = torch.rand(32, 32, 3, device="cuda")
esc = ptd.Scene(verts, faces, mat, albedo=[[0.6] * 3], emission=[[0, 0, 0]],
                env=env0)


def env_loss(e):
    esc.env = e.to("cuda")
    return (ptd.render(esc, ecam, 32, 32, spp=8, max_bounces=3, seed=11)
            * Wre).sum()


e = env0.clone().requires_grad_(True)
env_loss(e).backward()
worst = 0.0
for (y, x, c) in [(0, 0, 0), (1, 4, 1), (1, 7, 2)]:
    h = 5e-3
    ep = env0.clone(); ep[y, x, c] += h
    em2 = env0.clone(); em2[y, x, c] -= h
    fd = (env_loss(ep) - env_loss(em2)).item() / (2 * h)
    worst = max(worst, abs(e.grad[y, x, c].item() - fd) / abs(fd))
print(f"environment texel gradients vs central differences: "
      f"max rel err {worst:.2e} over 3 texels (detached CDF)")


# ---- v3: plastic FD, rough-dielectric slab, emission texels, medium, geometry
def cornell_mats(mat1_type, rough1):
    verts, faces, mat = [], [], []

    def wall(a, b, c, d, m):
        b0 = len(verts)
        verts.extend([a, b, c, d])
        faces.extend(quad(b0, b0 + 1, b0 + 2, b0 + 3))
        mat.extend([m, m])

    X, Y, Z = 5.56, 5.488, 5.592
    wall((0, 0, 0), (X, 0, 0), (X, 0, Z), (0, 0, Z), 0)
    wall((0, Y, 0), (0, Y, Z), (X, Y, Z), (X, Y, 0), 0)
    wall((0, 0, Z), (X, 0, Z), (X, Y, Z), (0, Y, Z), 0)
    wall((0, 0, 0), (0, 0, Z), (0, Y, Z), (0, Y, 0), 1)
    wall((X, 0, 0), (X, Y, 0), (X, Y, Z), (X, 0, Z), 2)
    wall((2.13, Y - 0.001, 2.27), (3.43, Y - 0.001, 2.27),
         (3.43, Y - 0.001, 3.32), (2.13, Y - 0.001, 3.32), 3)
    return ptd.Scene(verts, faces, mat,
                     albedo=[[0.73] * 3, [0.9, 0.5, 0.2],
                             [0.12, 0.45, 0.15], [0.78] * 3],
                     emission=[[0, 0, 0], [0, 0, 0], [0, 0, 0],
                               [18.4, 15.6, 8.0]],
                     material_types=[0, mat1_type, 0, 0],
                     roughness=[0.3, rough1, 0.3, 0.3])


def pl_loss(alb):
    s = cornell_mats(ptd.PLASTIC, 0.25)
    s.albedo = alb
    return (ptd.render(s, cam, 48, 48, spp=8, max_bounces=3, seed=13)
            * Wr).sum()


base_alb = torch.tensor([[0.73] * 3, [0.9, 0.5, 0.2], [0.12, 0.45, 0.15],
                         [0.78] * 3], device="cuda")
alb = base_alb.clone().requires_grad_(True)
pl_loss(alb).backward()
worst = 0.0
for (i, c) in [(1, 0), (1, 2)]:
    h = 2e-3
    ap = base_alb.clone(); ap[i, c] += h
    am = base_alb.clone(); am[i, c] -= h
    fd = (pl_loss(ap) - pl_loss(am)).item() / (2 * h)
    worst = max(worst, abs(alb.grad[i, c].item() - fd) / abs(fd))
print(f"plastic albedo gradients vs central differences (mixed lobes): "
      f"max rel err {worst:.2e} over 2 entries")

verts, faces, mat = [], [], []
quad_w((-10, -10, 5), (10, -10, 5), (10, 10, 5), (-10, 10, 5), 0)
add_box(verts, faces, mat, (-10, -10, 2.0), (10, 10, 2.2), 1)
rslab = ptd.Scene(verts, faces, mat, albedo=[[0, 0, 0], [0, 0, 0]],
                  emission=[[1.0] * 3, [0, 0, 0]],
                  material_types=[ptd.DIFFUSE, ptd.ROUGH_DIELECTRIC],
                  roughness=[0.3, 0.02], ior=[1.5, 1.5])
img = ptd.render(rslab, scam, 32, 32, spp=1024, max_bounces=10,
                 estimator="brdf", seed=5)
got = img[12:20, 12:20].mean().item()
print(f"rough dielectric slab (alpha 0.02): smooth analytic T {exp_T:.6f}, "
      f"measured {got:.6f}, rel err {abs(got-exp_T)/exp_T:.4%}")

import math as _math  # noqa: E402

verts, faces, mat = [], [], []
b0 = len(verts)
verts.extend([(-10, -10, 4), (10, -10, 4), (10, 10, 4), (-10, 10, 4)])
faces.extend(quad(b0, b0 + 1, b0 + 2, b0 + 3))
mat.extend([0, 0])
mcam = ptd.Camera(position=(0, 0, 0), look_at=(0, 0, 1), vfov_deg=10.0)
msc = ptd.Scene(verts, faces, mat, albedo=[[0, 0, 0]],
                emission=[[2.0, 2.0, 2.0]],
                medium=(torch.tensor([0.25] * 3, device="cuda"),
                        torch.tensor([1e-6] * 3, device="cuda")))
img = ptd.render(msc, mcam, 32, 32, spp=1024, max_bounces=2,
                 estimator="brdf", seed=4)
exp_bl = 2.0 * _math.exp(-0.25 * 4.0)
print(f"medium Beer-Lambert: analytic {exp_bl:.6f}, measured "
      f"{img.mean().item():.6f}, rel err "
      f"{abs(img.mean().item()-exp_bl)/exp_bl:.4%}")

# geometry gradients: per-vertex FD on the occluder scene
def geo_scene(single=None):
    verts, faces, mat = [], [], []
    verts.extend([(-4, 0, -4), (4, 0, -4), (4, 0, 4), (-4, 0, 4)])
    faces.extend(quad(0, 1, 2, 3)); mat.extend([0, 0])
    verts.extend([(-0.8, 4.0, -0.8), (0.8, 4.0, -0.8), (0.8, 4.0, 0.8),
                  (-0.8, 4.0, 0.8)])
    faces.extend(quad(4, 5, 6, 7)); mat.extend([1, 1])
    verts.extend([(-1.0, 2.0, -1.0), (1.0, 2.0, -1.0), (1.0, 2.0, 1.0),
                  (-1.0, 2.0, 1.0)])
    faces.extend(quad(8, 9, 10, 11)); mat.extend([2, 2])
    if single is not None:
        vid, dx = single
        v = list(verts[vid]); v[0] += dx; verts[vid] = tuple(v)
    return ptd.Scene(verts, faces, mat,
                     albedo=[[0.7] * 3, [0.8] * 3, [0.5] * 3],
                     emission=[[0, 0, 0], [10.0] * 3, [0, 0, 0]])


gcam = ptd.Camera(position=(0, 5.0, 7.0), look_at=(0, 0, 0), vfov_deg=45.0)
torch.manual_seed(0)
Wg = torch.rand(64, 64, 3, device="cuda")


def geo_loss(s, seed):
    img = ptd.render(s, gcam, 64, 64, spp=512, max_bounces=2,
                     estimator="brdf", seed=seed)
    return (img * Wg).sum().item()


gv = ptd.geometry_grad(geo_scene(), gcam, Wg, spp=64, edge_samples=1 << 20,
                       seed=7)
worst = 0.0
h = 0.1
for vid in (8, 9, 10, 11):
    fds = [(geo_loss(geo_scene((vid, h)), s) -
            geo_loss(geo_scene((vid, -h)), s)) / (2 * h)
           for s in range(3, 11)]
    fd = sum(fds) / len(fds)
    worst = max(worst, abs(gv[vid, 0].item() - fd) / abs(fd))
print(f"geometry gradients (interior + shadow + primary silhouettes) vs "
      f"seed-averaged FD, occluder verts: max rel err {worst:.2f}")

# ---- throughput helper
def bench(scene, bcam, H, W, spp, B, label, grad_leaf=None):
    img = ptd.render(scene, bcam, H, W, spp=spp, max_bounces=B)
    torch.cuda.synchronize()
    ts = []
    for _ in range(3):
        t0 = time.perf_counter()
        img = ptd.render(scene, bcam, H, W, spp=spp, max_bounces=B)
        torch.cuda.synchronize()
        ts.append(time.perf_counter() - t0)
    fwd = sorted(ts)[1]
    paths = H * W * spp
    line = (f"{label}: forward {fwd*1e3:.1f} ms ({paths/fwd/1e6:.0f} Mpaths/s)")
    if grad_leaf is not None:
        g = torch.ones(H, W, 3, device="cuda")
        img = ptd.render(scene, bcam, H, W, spp=spp, max_bounces=B)
        img.backward(g)
        torch.cuda.synchronize()
        ts = []
        for _ in range(3):
            grad_leaf.grad = None
            img = ptd.render(scene, bcam, H, W, spp=spp, max_bounces=B)
            torch.cuda.synchronize()
            t0 = time.perf_counter()
            img.backward(g)
            torch.cuda.synchronize()
            ts.append(time.perf_counter() - t0)
        bwd = sorted(ts)[1]
        line += f"; backward {bwd*1e3:.1f} ms ({bwd/fwd:.2f}x)"
    print(line)


# cornell (continuity with v1 numbers)
scene = cornell(with_box=True)
alb = scene.albedo.clone().requires_grad_(True)
scene.albedo = alb
bench(scene, cam, 512, 512, 64, 4, "cornell 24 tris, 512x512 spp 64 b4", alb)


# ---- large scenes: displaced terrain + conductor box + checker + env
def terrain_scene(n):
    ax = torch.linspace(0, 10, n + 1)
    X, Z = torch.meshgrid(ax, ax, indexing="ij")
    Y = (0.35 * torch.sin(X * 1.7) * torch.cos(Z * 1.3) +
         0.2 * torch.sin(X * 4.1 + 1.0) * torch.sin(Z * 3.7) +
         0.08 * torch.sin(X * 9.3) * torch.cos(Z * 8.1))
    verts = torch.stack([X, Y, Z], dim=-1).reshape(-1, 3)
    ii = torch.arange(n)
    jj = torch.arange(n)
    I, J = torch.meshgrid(ii, jj, indexing="ij")
    v00 = (I * (n + 1) + J).reshape(-1)
    v10 = ((I + 1) * (n + 1) + J).reshape(-1)
    v01 = (I * (n + 1) + J + 1).reshape(-1)
    v11 = ((I + 1) * (n + 1) + J + 1).reshape(-1)
    f1 = torch.stack([v00, v10, v11], dim=1)
    f2 = torch.stack([v00, v11, v01], dim=1)
    faces = torch.cat([f1, f2], dim=0)
    mat = torch.zeros(faces.shape[0], dtype=torch.int64)
    uv = torch.stack([X / 10.0, Z / 10.0], dim=-1).reshape(-1, 2)

    verts_l, faces_l, mat_l = [verts], [faces], [mat]
    base = verts.shape[0]
    bx = torch.tensor([(4.2, 0.4, 4.2), (5.8, 0.4, 4.2), (5.8, 0.4, 5.8),
                       (4.2, 0.4, 5.8), (4.2, 2.2, 4.2), (5.8, 2.2, 4.2),
                       (5.8, 2.2, 5.8), (4.2, 2.2, 5.8)])
    verts_l.append(bx)
    bq = [[0, 1, 2], [0, 2, 3], [4, 7, 6], [4, 6, 5], [0, 4, 5], [0, 5, 1],
          [3, 2, 6], [3, 6, 7], [0, 3, 7], [0, 7, 4], [1, 5, 6], [1, 6, 2]]
    faces_l.append(torch.tensor(bq, dtype=torch.int64) + base)
    mat_l.append(torch.ones(12, dtype=torch.int64))
    verts = torch.cat(verts_l)
    faces = torch.cat(faces_l)
    mat = torch.cat(mat_l)
    uv = torch.cat([uv, torch.zeros(8, 2)])

    yy, xx = torch.meshgrid(torch.arange(256), torch.arange(256),
                            indexing="ij")
    checker = (((yy // 16 + xx // 16) % 2).float() * 0.45 + 0.3)
    tex = torch.stack([checker, checker * 0.9, checker * 0.7], dim=-1)
    env = torch.full((16, 32, 3), 0.9)
    t0 = time.perf_counter()
    sc = ptd.Scene(verts, faces, mat,
                   albedo=[tex.cuda(), torch.tensor([0.9, 0.65, 0.35])],
                   emission=[[0, 0, 0], [0, 0, 0]],
                   material_types=[ptd.DIFFUSE, ptd.CONDUCTOR],
                   roughness=[0.3, 0.15], uvs=uv, env=env)
    build = time.perf_counter() - t0
    return sc, build


tcam = ptd.Camera(position=(11.0, 4.5, 11.0), look_at=(5.0, 0.5, 5.0),
                  vfov_deg=45.0)
for n in (511, 723):
    sc, build = terrain_scene(n)
    ntri = sc.n_faces
    tex_leaf = sc.albedo_textures[0].requires_grad_(True)
    sc.albedo_textures[0] = tex_leaf
    bench(sc, tcam, 512, 512, 16, 4,
          f"terrain {ntri} tris (SAH build {build:.1f} s), 512x512 spp 16 b4",
          tex_leaf)