import math

import pytest
import torch

import kernels

ptd = kernels.get_kernel("phanerozoic/pathtracer-diff", version=1,
                         trust_remote_code=True)

requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(),
                                   reason="CUDA required")


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)),   # bottom
              quad(i(4), i(7), i(6), i(5)),   # top
              quad(i(0), i(4), i(5), i(1)),   # z0
              quad(i(3), i(2), i(6), i(7)),   # z1
              quad(i(0), i(3), i(7), i(4)),   # x0
              quad(i(1), i(5), i(6), i(2))]:  # x1
        faces.extend(q)
        mat.extend([mat_id, mat_id])


def furnace_scene(albedo=0.5, emission=0.25):
    verts, faces, mat = [], [], []
    add_box(verts, faces, mat, (-2, -2, -2), (2, 2, 2), 0)
    return ptd.Scene(verts, faces, mat,
                     albedo=[[albedo] * 3], emission=[[emission] * 3])


def cornell(with_box=True, albedo=None, emission=None):
    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)          # floor
    wall((0, Y, 0), (0, Y, Z), (X, Y, Z), (X, Y, 0), 0)          # ceiling
    wall((0, 0, Z), (X, 0, Z), (X, Y, Z), (0, Y, Z), 0)          # back
    wall((0, 0, 0), (0, 0, Z), (0, Y, Z), (0, Y, 0), 1)          # left (red)
    wall((X, 0, 0), (X, Y, 0), (X, Y, Z), (X, 0, Z), 2)          # right (green)
    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)    # light
    if with_box:
        add_box(verts, faces, mat, (1.3, 0.0, 0.65), (2.4, 1.65, 1.7), 0)
    if albedo is None:
        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]]
    if emission is None:
        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)


def cornell_camera():
    return ptd.Camera(position=(2.78, 2.73, -8.0), look_at=(2.78, 2.73, 2.8),
                      vfov_deg=39.0)


@requires_cuda
@pytest.mark.kernels_ci
def test_furnace_analytic_zero_variance():
    """Closed uniform box, BRDF sampling only: with cosine sampling the
    diffuse throughput multiplier is exactly the albedo, so every path
    returns exactly E * sum_{k<B} a^k and every pixel equals the analytic
    series regardless of the random walk."""
    a, E, B = 0.5, 0.25, 8
    scene = furnace_scene(a, E)
    cam = ptd.Camera(position=(0, 0, 0), look_at=(0, 0, 1), vfov_deg=60.0)
    img = ptd.render(scene, cam, 32, 32, spp=2, max_bounces=B, nee=False,
                     seed=0)
    expected = E * (1 - a ** B) / (1 - a)
    assert torch.isfinite(img).all()
    assert (img - expected).abs().max().item() < 1e-3, \
        (img.min().item(), img.max().item(), expected)


@requires_cuda
@pytest.mark.kernels_ci
def test_forward_bitwise_deterministic():
    scene = cornell()
    cam = cornell_camera()
    i1 = ptd.render(scene, cam, 96, 96, spp=8, max_bounces=4, seed=11)
    i2 = ptd.render(scene, cam, 96, 96, spp=8, max_bounces=4, seed=11)
    assert torch.equal(i1, i2)
    i3 = ptd.render(scene, cam, 96, 96, spp=8, max_bounces=4, seed=12)
    assert not torch.equal(i1, i3)


@requires_cuda
@pytest.mark.kernels_ci
def test_nee_matches_brdf_estimator():
    """The NEE-only and BRDF-only estimators target the same integral; their
    image means must agree within Monte Carlo tolerance."""
    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])
    scene = ptd.Scene(verts, faces, mat,
                      albedo=[[0.6, 0.6, 0.6], [0.0, 0.0, 0.0]],
                      emission=[[0, 0, 0], [3.0, 3.0, 3.0]])
    cam = 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(scene, cam, 64, 64, spp=64, max_bounces=4, nee=True,
                       seed=1).mean().item()
    m_brdf = ptd.render(scene, cam, 64, 64, spp=256, max_bounces=4, nee=False,
                        seed=2).mean().item()
    assert abs(m_nee - m_brdf) / m_brdf < 0.03, (m_nee, m_brdf)


@requires_cuda
@pytest.mark.kernels_ci
def test_gradients_match_finite_differences():
    """Sampling directions never depend on albedo/emission, so with a fixed
    seed the loss is smooth in the parameters along the same paths and
    central differences must match the analytic replay gradients."""
    scene = cornell(with_box=False)
    cam = cornell_camera()
    H = W = 48
    torch.manual_seed(0)
    Wr = torch.rand(H, W, 3, device="cuda")

    def loss_of(albedo, emission):
        scene.albedo = albedo
        scene.emission = emission
        img = ptd.render(scene, cam, H, W, spp=8, max_bounces=3, nee=True,
                         seed=7)
        return (img * 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()

    checks = []
    for (i, c) in [(0, 1), (1, 0), (2, 2)]:  # white g, red r, green b
        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)
        an = alb.grad[i, c].item()
        checks.append(("albedo", i, c, an, fd))
    h = 0.05
    ep = emi0.clone(); ep[3, 1] += h
    em = emi0.clone(); em[3, 1] -= h
    fd = (loss_of(alb0, ep) - loss_of(alb0, em)).item() / (2 * h)
    an = emi.grad[3, 1].item()
    checks.append(("emission", 3, 1, an, fd))

    for kind, i, c, an, fd in checks:
        assert fd != 0.0, (kind, i, c)
        assert abs(an - fd) / max(abs(fd), 1e-6) < 2e-2, \
            (kind, i, c, an, fd)


@requires_cuda
@pytest.mark.kernels_ci
def test_inverse_rendering_recovers_wall_albedo():
    """Recover the left-wall albedo from a rendered target by gradient
    descent through the kernel: the flagship no-framework-in-the-loop use."""
    target_row = torch.tensor([0.2, 0.5, 0.7], device="cuda")
    scene = cornell(with_box=False)
    cam = cornell_camera()
    H = W = 48

    tgt_alb = scene.albedo.clone()
    tgt_alb[1] = target_row
    scene.albedo = tgt_alb
    target = ptd.render(scene, cam, H, W, spp=8, max_bounces=3, seed=3).detach()

    alb = scene.albedo.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)
    losses = []
    for _ in range(80):
        opt.zero_grad()
        scene.albedo = alb
        img = ptd.render(scene, cam, H, W, spp=8, max_bounces=3, seed=3)
        loss = (img - 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)
        losses.append(loss.item())
    err = (alb[1].detach() - target_row).abs().max().item()
    assert err < 0.05, (alb[1].detach().tolist(), losses[0], losses[-1])
    assert losses[-1] < losses[0] * 0.1, (losses[0], losses[-1])


def cornell_tex(back_albedo):
    """Cornell walls with the back wall on its own material carrying a
    texture; per-corner UVs map the back quad to [0,1]^2."""
    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)   # back
    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)


@requires_cuda
@pytest.mark.kernels_ci
def test_texture_constants_equivalence():
    """Per-material constants and 1x1 textures share one code path: images
    agree to float rounding and the 1x1 texel gradient equals the constant
    gradient."""
    cam = cornell_camera()
    s1 = cornell(with_box=False)
    i1 = ptd.render(s1, cam, 64, 64, spp=8, max_bounces=3, seed=9)

    consts = [[0.73, 0.73, 0.73], [0.65, 0.05, 0.05], [0.12, 0.45, 0.15],
              [0.78, 0.78, 0.78]]
    texs = [torch.tensor(c, device="cuda").reshape(1, 1, 3) for c in consts]
    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)
    emission = [[0, 0, 0], [0, 0, 0], [0, 0, 0], [18.4, 15.6, 8.0]]
    s2 = ptd.Scene(verts, faces, mat, albedo=texs, emission=emission)
    i2 = ptd.render(s2, cam, 64, 64, spp=8, max_bounces=3, seed=9)
    assert (i1 - i2).abs().max().item() < 5e-6

    torch.manual_seed(1)
    Wr = torch.rand(64, 64, 3, device="cuda")
    alb = s1.albedo.clone().requires_grad_(True)
    s1.albedo = alb
    (ptd.render(s1, cam, 64, 64, spp=8, max_bounces=3, seed=9) * Wr).sum().backward()
    texs_g = [t.clone().requires_grad_(True) for t in texs]
    s2.albedo_textures = texs_g
    (ptd.render(s2, cam, 64, 64, spp=8, max_bounces=3, seed=9) * Wr).sum().backward()
    for m in range(4):
        a = alb.grad[m]
        b = texs_g[m].grad.reshape(3)
        assert (a - b).abs().max().item() < 1e-3 + 1e-3 * a.abs().max().item()


@requires_cuda
@pytest.mark.kernels_ci
def test_texture_gradients_match_finite_differences():
    """Texel gradients through the bilinear-footprint adjoint match central
    differences along the same paths."""
    torch.manual_seed(2)
    base_tex = (torch.rand(4, 4, 3, device="cuda") * 0.6 + 0.2)
    H = W = 48
    cam = cornell_camera()
    torch.manual_seed(0)
    Wr = torch.rand(H, W, 3, device="cuda")

    def loss_of(t):
        scene = cornell_tex(t)
        img = ptd.render(scene, cam, H, W, spp=8, max_bounces=3, seed=7)
        return (img * Wr).sum()

    t = base_tex.clone().requires_grad_(True)
    loss_of(t).backward()

    worst = 0.0
    for (y, x, c) in [(0, 0, 0), (1, 2, 1), (3, 3, 2), (2, 1, 0)]:
        h = 2e-3
        tp = base_tex.clone(); tp[y, x, c] += h
        tm = base_tex.clone(); tm[y, x, c] -= h
        fd = (loss_of(tp) - loss_of(tm)).item() / (2 * h)
        an = t.grad[y, x, c].item()
        assert fd != 0.0, (y, x, c)
        worst = max(worst, abs(an - fd) / abs(fd))
    assert worst < 2e-2, worst


@requires_cuda
@pytest.mark.kernels_ci
def test_texture_inverse_recovery():
    """Recover an 8x8 back-wall texture from a rendered target: 192 unknowns,
    the regime per-material constants cannot express."""
    torch.manual_seed(4)
    target_tex = (torch.rand(8, 8, 3, device="cuda") * 0.7 + 0.15)
    cam = cornell_camera()
    H = W = 64
    scene = cornell_tex(target_tex)
    target = ptd.render(scene, cam, H, W, spp=8, max_bounces=3, seed=3).detach()

    t = torch.full((8, 8, 3), 0.5, device="cuda", requires_grad=True)
    scene_opt = cornell_tex(t)
    opt = torch.optim.Adam([t], lr=0.1)
    first = None
    for _ in range(200):
        opt.zero_grad()
        img = ptd.render(scene_opt, cam, H, W, spp=8, max_bounces=3, seed=3)
        loss = (img - 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()
    assert loss.item() < first * 0.05, (first, loss.item())
    assert err.mean().item() < 0.08, err.mean().item()


@requires_cuda
@pytest.mark.kernels_ci
def test_estimators_agree_three_ways():
    """BRDF-only, NEE-only, and MIS target the same integral at the same
    bounce budget; their image means must agree within Monte Carlo
    tolerance."""
    scene = cornell(with_box=True)
    cam = cornell_camera()
    m_mis = ptd.render(scene, cam, 64, 64, spp=64, max_bounces=4,
                       estimator="mis", seed=1).mean().item()
    m_nee = ptd.render(scene, cam, 64, 64, spp=64, max_bounces=4,
                       estimator="nee", seed=2).mean().item()
    m_brdf = ptd.render(scene, cam, 64, 64, spp=1024, max_bounces=4,
                        estimator="brdf", seed=3).mean().item()
    assert abs(m_mis - m_nee) / m_nee < 0.03, (m_mis, m_nee)
    assert abs(m_mis - m_brdf) / m_brdf < 0.05, (m_mis, m_brdf)


@requires_cuda
@pytest.mark.kernels_ci
def test_conductor_under_env_estimators_agree():
    """A GGX conductor under a constant environment: MIS and BRDF-only
    means must agree, which jointly validates the VNDF weight, the GGX pdf,
    and the environment pdf."""
    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((-4, 0, -4), (4, 0, -4), (4, 0, 4), (-4, 0, 4), 0)   # floor
    add_box(verts, faces, mat, (-1, 0, -1), (1, 2, 1), 1)       # metal box
    env = torch.full((8, 16, 3), 1.0)
    scene = ptd.Scene(verts, faces, mat,
                      albedo=[[0.5, 0.5, 0.5], [0.9, 0.6, 0.3]],
                      emission=[[0, 0, 0], [0, 0, 0]],
                      material_types=[ptd.DIFFUSE, ptd.CONDUCTOR],
                      roughness=[0.3, 0.25], env=env)
    cam = ptd.Camera(position=(5, 3, 5), look_at=(0, 1, 0), vfov_deg=45.0)
    m_mis = ptd.render(scene, cam, 64, 64, spp=64, max_bounces=4,
                       estimator="mis", seed=1).mean().item()
    m_brdf = ptd.render(scene, cam, 64, 64, spp=512, max_bounces=4,
                        estimator="brdf", seed=2).mean().item()
    assert abs(m_mis - m_brdf) / m_brdf < 0.04, (m_mis, m_brdf)


@requires_cuda
@pytest.mark.kernels_ci
def test_fresnel_slab_analytic():
    """A smooth glass slab in front of a uniform emitter, viewed at normal
    incidence: the transmitted fraction sums the internal-bounce series to
    (1-R)/(1+R) with R = ((n-1)/(n+1))^2."""
    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)  # emitter
    add_box(verts, faces, mat, (-10, -10, 2.0), (10, 10, 2.2), 1)      # slab
    n_ior = 1.5
    scene = ptd.Scene(verts, faces, mat,
                      albedo=[[0, 0, 0], [0, 0, 0]],
                      emission=[[1.0, 1.0, 1.0], [0, 0, 0]],
                      material_types=[ptd.DIFFUSE, ptd.DIELECTRIC],
                      ior=[1.5, n_ior])
    cam = ptd.Camera(position=(0, 0, 0), look_at=(0, 0, 1), vfov_deg=10.0)
    img = ptd.render(scene, cam, 32, 32, spp=256, max_bounces=10,
                     estimator="brdf", seed=5)
    R = ((n_ior - 1) / (n_ior + 1)) ** 2
    expected = (1 - R) / (1 + R)
    got = img[12:20, 12:20].mean().item()
    assert abs(got - expected) / expected < 0.01, (got, expected)


@requires_cuda
@pytest.mark.kernels_ci
def test_env_furnace_diffuse():
    """A diffuse plane under a constant environment c: with BRDF sampling
    every sample returns exactly albedo * c (cosine pdf cancellation plus a
    guaranteed env hit), and MIS agrees in expectation."""
    verts, faces, mat = [], [], []
    base = len(verts)
    verts.extend([(-5, 0, -5), (5, 0, -5), (5, 0, 5), (-5, 0, 5)])
    faces.extend(quad(base, base + 1, base + 2, base + 3))
    mat.extend([0, 0])
    c = 0.8
    a = 0.6
    env = torch.full((8, 16, 3), c)
    scene = ptd.Scene(verts, faces, mat, albedo=[[a] * 3],
                      emission=[[0, 0, 0]], env=env)
    cam = ptd.Camera(position=(0, 2.0, 0.01), look_at=(0, 0, 0.2),
                     vfov_deg=30.0)
    img = ptd.render(scene, cam, 32, 32, spp=4, max_bounces=3,
                     estimator="brdf", seed=0)
    assert (img - a * c).abs().max().item() < 2e-3, \
        (img.min().item(), img.max().item(), a * c)
    m_mis = ptd.render(scene, cam, 32, 32, spp=256, max_bounces=3,
                       estimator="mis", seed=1).mean().item()
    assert abs(m_mis - a * c) / (a * c) < 0.02, (m_mis, a * c)


@requires_cuda
@pytest.mark.kernels_ci
def test_env_texel_gradients_match_fd():
    """Environment texels are linear in the estimate and the sampling CDF is
    detached, so same-seed central differences are exact."""
    torch.manual_seed(3)
    env0 = torch.rand(4, 8, 3, device="cuda") * 1.5 + 0.2
    verts, faces, mat = [], [], []
    base = len(verts)
    verts.extend([(-5, 0, -5), (5, 0, -5), (5, 0, 5), (-5, 0, 5)])
    faces.extend(quad(base, base + 1, base + 2, base + 3))
    mat.extend([0, 0])
    cam = ptd.Camera(position=(0, 2.0, 0.01), look_at=(0, 0, 0.5),
                     vfov_deg=40.0)
    torch.manual_seed(0)
    Wr = torch.rand(32, 32, 3, device="cuda")
    scene = ptd.Scene(verts, faces, mat, albedo=[[0.6] * 3],
                      emission=[[0, 0, 0]], env=env0)

    def loss_of(e):
        # swap texel values under the construction-time (detached) CDF, the
        # frozen-sampling semantics the kernel defines
        scene.env = e.to("cuda")
        img = ptd.render(scene, cam, 32, 32, spp=8, max_bounces=3, seed=11)
        return (img * Wr).sum()

    e = env0.clone().requires_grad_(True)
    loss_of(e).backward()
    worst = 0.0
    for (y, x, c) in [(0, 0, 0), (1, 4, 1), (1, 7, 2)]:  # upper hemisphere
        h = 5e-3
        ep = env0.clone(); ep[y, x, c] += h
        em = env0.clone(); em[y, x, c] -= h
        fd = (loss_of(ep) - loss_of(em)).item() / (2 * h)
        an = e.grad[y, x, c].item()
        assert fd != 0.0, (y, x, c)
        worst = max(worst, abs(an - fd) / abs(fd))
    assert worst < 2e-2, worst
    # a texel no path can see gets exactly zero from both sides
    assert e.grad[3, 0, 0].item() == 0.0


@requires_cuda
@pytest.mark.kernels_ci
def test_conductor_albedo_gradients_match_fd():
    """Conductor F0 enters every term through Schlick Fresnel, which is
    affine in F0, so replay gradients must match same-seed central
    differences through GGX sampling and MIS."""
    cam = cornell_camera()
    H = W = 48
    torch.manual_seed(0)
    Wr = torch.rand(H, W, 3, device="cuda")
    base_alb = torch.tensor([[0.73, 0.73, 0.73], [0.9, 0.5, 0.2],
                             [0.12, 0.45, 0.15], [0.78, 0.78, 0.78]],
                            device="cuda")

    def loss_of(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)   # conductor
        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)
        scene = 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=[ptd.DIFFUSE, ptd.CONDUCTOR,
                                          ptd.DIFFUSE, ptd.DIFFUSE],
                          roughness=[0.3, 0.2, 0.3, 0.3])
        img = ptd.render(scene, cam, H, W, spp=8, max_bounces=3, seed=13)
        return (img * Wr).sum()

    alb = base_alb.clone().requires_grad_(True)
    loss_of(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 = (loss_of(ap) - loss_of(am)).item() / (2 * h)
        an = alb.grad[i, c].item()
        assert fd != 0.0, (i, c)
        worst = max(worst, abs(an - fd) / abs(fd))
    assert worst < 2e-2, worst


def cornell_mats(mat1_type, rough1=0.2, albedo1=(0.9, 0.5, 0.2)):
    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)
    albedo = [[0.73] * 3, list(albedo1), [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]]
    return ptd.Scene(verts, faces, mat, albedo=albedo, emission=emission,
                     material_types=[0, mat1_type, 0, 0],
                     roughness=[0.3, rough1, 0.3, 0.3])


@requires_cuda
@pytest.mark.kernels_ci
def test_plastic_estimators_and_fd():
    """Rough plastic: MIS and BRDF-only agree (validating the mixed-lobe
    pdf), and the albedo gradient through both lobes matches central
    differences (the diffuse part is affine, the coat constant)."""
    cam = cornell_camera()
    scene = cornell_mats(ptd.PLASTIC, rough1=0.25)
    m_mis = ptd.render(scene, cam, 64, 64, spp=64, max_bounces=4,
                       estimator="mis", seed=1).mean().item()
    m_brdf = ptd.render(scene, cam, 64, 64, spp=768, max_bounces=4,
                        estimator="brdf", seed=2).mean().item()
    assert abs(m_mis - m_brdf) / m_brdf < 0.04, (m_mis, m_brdf)

    H = W = 48
    torch.manual_seed(0)
    Wr = torch.rand(H, W, 3, device="cuda")
    base = torch.tensor([[0.73] * 3, [0.9, 0.5, 0.2], [0.12, 0.45, 0.15],
                         [0.78] * 3], device="cuda")

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

    alb = base.clone().requires_grad_(True)
    loss_of(alb).backward()
    worst = 0.0
    for (i, c) in [(1, 0), (1, 2)]:
        h = 2e-3
        ap = base.clone(); ap[i, c] += h
        am = base.clone(); am[i, c] -= h
        fd = (loss_of(ap) - loss_of(am)).item() / (2 * h)
        assert fd != 0.0
        worst = max(worst, abs(alb.grad[i, c].item() - fd) / abs(fd))
    assert worst < 2e-2, worst


@requires_cuda
@pytest.mark.kernels_ci
def test_rough_dielectric_smooth_limit():
    """A rough-dielectric slab at alpha = 0.02 must transmit the smooth
    analytic series (1-R)/(1+R) at normal incidence within tolerance."""
    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)
    scene = 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])
    cam = ptd.Camera(position=(0, 0, 0), look_at=(0, 0, 1), vfov_deg=10.0)
    img = ptd.render(scene, cam, 32, 32, spp=512, max_bounces=10,
                     estimator="brdf", seed=5)
    R = (0.5 / 2.5) ** 2
    expected = (1 - R) / (1 + R)
    got = img[12:20, 12:20].mean().item()
    assert abs(got - expected) / expected < 0.02, (got, expected)


@requires_cuda
@pytest.mark.kernels_ci
def test_emission_texture_fd_and_inverse():
    """Per-texel emission: gradients are linear-exact against central
    differences, and an emissive panel texture recovers from the room it
    lights."""
    cam = cornell_camera()
    H = W = 48
    torch.manual_seed(0)
    Wr = torch.rand(H, W, 3, device="cuda")

    def build(etex):
        verts, faces, mat = [], [], []

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

        uvs = []
        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, uv=True)
        albedo = [[0.73] * 3, [0.65, 0.05, 0.05], [0.12, 0.45, 0.15],
                  [0.78] * 3]
        emission = [torch.zeros(3), torch.zeros(3), torch.zeros(3), etex]
        return ptd.Scene(verts, faces, mat, albedo=albedo, emission=emission,
                         uvs=uvs)

    torch.manual_seed(5)
    e0 = torch.rand(4, 8, 3, device="cuda") * 12.0 + 3.0
    scene = build(e0)

    def loss_of(e):
        scene.emission_textures[3] = e.to("cuda")
        img = ptd.render(scene, cam, H, W, spp=8, max_bounces=3, seed=7)
        return (img * Wr).sum()

    e = e0.clone().requires_grad_(True)
    loss_of(e).backward()
    worst = 0.0
    for (y, x, c) in [(0, 0, 0), (2, 5, 1), (3, 7, 2)]:
        h = 0.1
        ep = e0.clone(); ep[y, x, c] += h
        em = e0.clone(); em[y, x, c] -= h
        fd = (loss_of(ep) - loss_of(em)).item() / (2 * h)
        assert fd != 0.0, (y, x, c)
        worst = max(worst, abs(e.grad[y, x, c].item() - fd) / abs(fd))
    assert worst < 2e-2, worst

    target = ptd.render(build(e0), cam, H, W, spp=8, max_bounces=3,
                        seed=3).detach()
    t = torch.full((4, 8, 3), 8.0, device="cuda", requires_grad=True)
    s_opt = build(t)
    opt = torch.optim.Adam([t], lr=0.4)
    first = None
    for _ in range(150):
        opt.zero_grad()
        loss = (ptd.render(s_opt, cam, H, W, spp=8, max_bounces=3, seed=3)
                - target).square().mean()
        loss.backward()
        opt.step()
        with torch.no_grad():
            t.clamp_(0.1, 30.0)
        if first is None:
            first = loss.item()
    assert loss.item() < first * 0.05, (first, loss.item())
    rel = ((t.detach() - e0).abs() / e0).mean().item()
    assert rel < 0.2, rel


@requires_cuda
@pytest.mark.kernels_ci
def test_medium_beer_lambert_and_fd():
    """Gray purely-absorbing medium: the mean transmitted radiance follows
    Beer-Lambert, and sigma gradients match central differences exactly
    (the sampling rate is detached)."""
    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])
    cam = ptd.Camera(position=(0, 0, 0), look_at=(0, 0, 1), vfov_deg=10.0)

    def build(sa, ss):
        return ptd.Scene(verts, faces, mat, albedo=[[0, 0, 0]],
                         emission=[[2.0, 2.0, 2.0]], medium=(sa, ss))

    sa0 = torch.tensor([0.25, 0.25, 0.25], device="cuda")
    ss0 = torch.tensor([1e-6, 1e-6, 1e-6], device="cuda")
    scene = build(sa0, ss0)
    img = ptd.render(scene, cam, 32, 32, spp=1024, max_bounces=2,
                     estimator="brdf", seed=4)
    expected = 2.0 * math.exp(-0.25 * 4.0)
    got = img.mean().item()
    assert abs(got - expected) / expected < 0.02, (got, expected)

    torch.manual_seed(0)
    Wr = torch.rand(32, 32, 3, device="cuda")

    def loss_of(sa, ss):
        s = build(sa, ss)
        s.med_sbar = scene.med_sbar  # keep the frozen rate identical
        img = ptd.render(s, cam, 32, 32, spp=16, max_bounces=2, seed=9)
        return (img * Wr).sum()

    sa = sa0.clone().requires_grad_(True)
    ss = ss0.clone().requires_grad_(True)
    loss_of(sa, ss).backward()
    h = 5e-3
    sap = sa0.clone(); sap[1] += h
    sam = sa0.clone(); sam[1] -= h
    fd = (loss_of(sap, ss0) - loss_of(sam, ss0)).item() / (2 * h)
    assert fd != 0.0
    assert abs(sa.grad[1].item() - fd) / abs(fd) < 2e-2, \
        (sa.grad[1].item(), fd)


@requires_cuda
@pytest.mark.kernels_ci
def test_medium_scattering_estimators_and_fd():
    """Foggy Cornell: MIS and BRDF-only agree in the mean, and the
    sigma_s gradient matches central differences under detached sampling."""
    def build(sa, ss):
        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.65, 0.05, 0.05],
                                 [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]], medium=(sa, ss))

    cam = cornell_camera()
    sa0 = torch.tensor([0.02, 0.02, 0.02], device="cuda")
    ss0 = torch.tensor([0.06, 0.06, 0.06], device="cuda")
    fog = build(sa0, ss0)
    m_mis = ptd.render(fog, cam, 64, 64, spp=64, max_bounces=6,
                       estimator="mis", seed=1).mean().item()
    m_brdf = ptd.render(fog, cam, 64, 64, spp=1024, max_bounces=6,
                        estimator="brdf", seed=2).mean().item()
    assert abs(m_mis - m_brdf) / m_brdf < 0.05, (m_mis, m_brdf)

    torch.manual_seed(0)
    Wr = torch.rand(48, 48, 3, device="cuda")
    sbar = fog.med_sbar

    def loss_of(sa, ss):
        s = build(sa, ss)
        s.med_sbar = sbar
        img = ptd.render(s, cam, 48, 48, spp=8, max_bounces=4, seed=11)
        return (img * Wr).sum()

    sa = sa0.clone().requires_grad_(True)
    ss = ss0.clone().requires_grad_(True)
    loss_of(sa, ss).backward()
    h = 2e-3
    ssp = ss0.clone(); ssp[0] += h
    ssm = ss0.clone(); ssm[0] -= h
    fd = (loss_of(sa0, ssp) - loss_of(sa0, ssm)).item() / (2 * h)
    assert fd != 0.0
    assert abs(ss.grad[0].item() - fd) / abs(fd) < 2e-2, \
        (ss.grad[0].item(), fd)


def _geo_scene(light_y=4.0, occluder_dx=None, single_vert=None):
    verts, faces, mat = [], [], []
    b0 = len(verts)
    verts.extend([(-4, 0, -4), (4, 0, -4), (4, 0, 4), (-4, 0, 4)])  # floor
    faces.extend(quad(b0, b0 + 1, b0 + 2, b0 + 3))
    mat.extend([0, 0])
    b0 = len(verts)
    verts.extend([(-0.8, light_y, -0.8), (0.8, light_y, -0.8),
                  (0.8, light_y, 0.8), (-0.8, light_y, 0.8)])       # light
    faces.extend(quad(b0, b0 + 1, b0 + 2, b0 + 3))
    mat.extend([1, 1])
    if occluder_dx is not None:
        d = occluder_dx
        b0 = len(verts)
        verts.extend([(-1.0 + d, 2.0, -1.0), (1.0 + d, 2.0, -1.0),
                      (1.0 + d, 2.0, 1.0), (-1.0 + d, 2.0, 1.0)])   # blocker
        faces.extend(quad(b0, b0 + 1, b0 + 2, b0 + 3))
        mat.extend([2, 2])
    if single_vert is not None:
        vid, dx = single_vert
        v = list(verts[vid])
        v[0] += dx
        verts[vid] = tuple(v)
    albedo = [[0.7, 0.7, 0.7], [0.8] * 3, [0.5] * 3]
    emission = [[0, 0, 0], [10.0, 10.0, 10.0], [0, 0, 0]]
    nm = max(mat) + 1
    return ptd.Scene(verts, faces, mat, albedo=albedo[:nm],
                     emission=emission[:nm])


def _geo_loss(scene, cam, Wr, seed=3):
    # brdf at B=2 renders exactly the full direct-lighting integral (no MIS
    # truncation, no indirect), the transport term the geometry pass
    # differentiates
    img = ptd.render(scene, cam, 64, 64, spp=64, max_bounces=2,
                     estimator="brdf", seed=seed)
    return (img * Wr).sum()


@requires_cuda
@pytest.mark.kernels_ci
def test_geometry_interior_matches_fd_smooth():
    """Translating the area light vertically: the dual-number interior term
    (shading + light-area measure) plus the light's own primary-silhouette
    term must match seed-averaged central differences of the direct
    image."""
    cam = ptd.Camera(position=(0, 5.0, 7.0), look_at=(0, 0, 0),
                     vfov_deg=45.0)
    torch.manual_seed(0)
    Wr = torch.rand(64, 64, 3, device="cuda")
    h = 5e-2
    fds = []
    for seed in range(3, 9):
        fds.append((_geo_loss(_geo_scene(light_y=4.0 + h), cam, Wr, seed=seed)
                    - _geo_loss(_geo_scene(light_y=4.0 - h), cam, Wr,
                                seed=seed)).item() / (2 * h))
    fd = sum(fds) / len(fds)
    scene = _geo_scene(light_y=4.0)
    gv = ptd.geometry_grad(scene, cam, Wr, spp=128, edge_samples=1 << 18,
                           seed=11)
    an = gv[4:8, 1].sum().item()  # the light's four verts, y components
    assert fd != 0.0
    assert abs(an - fd) / abs(fd) < 0.15, (an, fd)


@requires_cuda
@pytest.mark.kernels_ci
def test_geometry_boundary_per_vertex_fd():
    """Moving each occluder vertex sweeps its shadow (boundary term) and
    its own image outline (primary-silhouette term); the interior term is
    zero on the occluder (translating a plane inside itself changes
    nothing). Each vertex's x-gradient must match the finite-difference
    derivative averaged over seeds (a single-seed FD of a visibility
    discontinuity is flip noise)."""
    cam = ptd.Camera(position=(0, 5.0, 7.0), look_at=(0, 0, 0),
                     vfov_deg=45.0)
    torch.manual_seed(0)
    Wr = torch.rand(64, 64, 3, device="cuda")
    scene = _geo_scene(occluder_dx=0.0)
    gv = ptd.geometry_grad(scene, cam, Wr, spp=64, edge_samples=1 << 20,
                           seed=7)
    h = 0.1
    for vid in (8, 9, 10, 11):
        fds = []
        for seed in range(3, 9):
            fp = _geo_scene(occluder_dx=0.0, single_vert=(vid, h))
            fm = _geo_scene(occluder_dx=0.0, single_vert=(vid, -h))
            fds.append((_geo_loss(fp, cam, Wr, seed=seed) -
                        _geo_loss(fm, cam, Wr, seed=seed)).item() / (2 * h))
        fd = sum(fds) / len(fds)
        an = gv[vid, 0].item()
        assert fd != 0.0
        assert an * fd > 0.0, (vid, an, fd)
        assert abs(an - fd) / abs(fd) < 0.35, (vid, an, fd)


@requires_cuda
@pytest.mark.kernels_ci
def test_validation():
    verts, faces, mat = [], [], []
    add_box(verts, faces, mat, (0, 0, 0), (1, 1, 1), 0)
    with pytest.raises(ValueError):
        ptd.Scene(verts, faces, mat, albedo=[[0.5] * 3] * 65,
                  emission=[[0.0] * 3] * 65)
    scene = ptd.Scene(verts, faces, mat, albedo=[[0.5] * 3],
                      emission=[[1.0] * 3])
    cam = ptd.Camera(position=(0.5, 0.5, 0.5), look_at=(0.5, 0.5, 1.0))
    with pytest.raises(ValueError):
        ptd.render(scene, cam, 8, 8, spp=1, max_bounces=17)
    with pytest.raises(ValueError):
        ptd.render(scene, cam, 8, 8, spp=1, estimator="mis", nee=True)
    with pytest.raises(ValueError):
        ptd.Scene(verts, faces, mat, albedo=[[0.5] * 3],
                  emission=[[1.0] * 3], material_types=[7])
    env = torch.full((4, 8, 3), 1.0)
    scene_env = ptd.Scene(verts, faces, mat, albedo=[[0.5] * 3],
                          emission=[[1.0] * 3], env=env)
    with pytest.raises(ValueError):
        ptd.render(scene_env, cam, 8, 8, spp=1, estimator="nee")
    img = ptd.render(scene, cam, 8, 8, spp=2, max_bounces=2)
    assert img.shape == (8, 8, 3) and img.dtype == torch.float32
    assert torch.isfinite(img).all()