"""Vectorized torch path tracer (the substrate for neural ray tracing).

Same design rule as the physics engine: keep the exact parts analytic —
ray intersections and next-event-estimated direct lighting — and reserve
learning for the expensive part (indirect transport, see W9 experiment).

Scene: a neon Cornell box. Diffuse surfaces, one area light on the
ceiling, a sphere and a box inside. All ray math is batched over
(N,3) tensors on the chosen device.
"""
import torch

EPS = 1e-4

# ---- scene definition (unit room, front face at z=+1 open) ----
LIGHT = dict(y=0.999, half=0.42, Le=torch.tensor([11.0, 11.5, 13.0]))
SPHERE = dict(c=torch.tensor([-0.42, -0.65, -0.22]), r=0.35,
              alb=torch.tensor([0.85, 0.85, 0.88]))
BOX = dict(mn=torch.tensor([0.12, -1.0, -0.10]),
           mx=torch.tensor([0.72, -0.34, 0.52]),
           alb=torch.tensor([0.30, 0.75, 0.95]))
WALLS = [  # (axis, value, normal-sign, albedo)
    (0, -1.0, +1, [0.25, 0.55, 0.95]),   # left  — neon blue
    (0, +1.0, -1, [0.95, 0.30, 0.40]),   # right — crimson
    (1, -1.0, +1, [0.70, 0.70, 0.72]),   # floor
    (1, +1.0, -1, [0.70, 0.70, 0.72]),   # ceiling
    (2, -1.0, +1, [0.55, 0.60, 0.75]),   # back
]


def _dev(dev):
    out = {"Le": LIGHT["Le"].to(dev), "sc": SPHERE["c"].to(dev),
           "sa": SPHERE["alb"].to(dev), "bmn": BOX["mn"].to(dev),
           "bmx": BOX["mx"].to(dev), "ba": BOX["alb"].to(dev),
           "wa": torch.tensor([w[3] for w in WALLS], device=dev)}
    return out


def intersect(o, d, S):
    """Closest hit for rays (N,3),(N,3) -> t, normal, albedo, is_light."""
    N = len(o)
    dev = o.device
    INF = torch.full((N,), 1e9, device=dev)
    best_t = INF.clone()
    n = torch.zeros(N, 3, device=dev)
    alb = torch.zeros(N, 3, device=dev)
    is_light = torch.zeros(N, dtype=torch.bool, device=dev)

    arange = torch.arange(N, device=dev)
    for wi, (ax, val, sgn, _) in enumerate(WALLS):
        denom = d[:, ax]
        t = (val - o[:, ax]) / torch.where(denom.abs() < 1e-9,
                                           torch.full_like(denom, 1e-9), denom)
        p = o + t[:, None] * d
        oth = [a for a in range(3) if a != ax]
        ok = (t > EPS) & (t < best_t) \
            & (p[:, oth[0]].abs() <= 1.0) & (p[:, oth[1]].abs() <= 1.0) \
            & (p[:, 2] <= 1.0)
        best_t = torch.where(ok, t, best_t)
        nw = torch.zeros_like(n); nw[:, ax] = float(sgn)
        n = torch.where(ok[:, None], nw, n)
        alb = torch.where(ok[:, None], S["wa"][wi], alb)
    # sphere
    oc = o - S["sc"]
    b = (oc * d).sum(1)
    c = (oc * oc).sum(1) - SPHERE["r"] ** 2
    disc = b * b - c
    sq = torch.sqrt(disc.clamp_min(0))
    t1 = -b - sq
    t2 = -b + sq
    ts = torch.where(t1 > EPS, t1, t2)
    ok = (disc > 0) & (ts > EPS) & (ts < best_t)
    best_t = torch.where(ok, ts, best_t)
    ps = o + ts[:, None] * d
    n = torch.where(ok[:, None], (ps - S["sc"]) / SPHERE["r"], n)
    alb = torch.where(ok[:, None], S["sa"], alb)
    # box (slabs)
    inv = 1.0 / torch.where(d.abs() < 1e-9, torch.full_like(d, 1e-9), d)
    t0s = (S["bmn"] - o) * inv
    t1s = (S["bmx"] - o) * inv
    tsm = torch.minimum(t0s, t1s).max(1).values
    tbg = torch.maximum(t0s, t1s).min(1).values
    ok = (tsm < tbg) & (tsm > EPS) & (tsm < best_t)
    best_t = torch.where(ok, tsm, best_t)
    pb = o + tsm[:, None] * d
    ctr = (S["bmn"] + S["bmx"]) / 2
    half = (S["bmx"] - S["bmn"]) / 2
    rel = (pb - ctr) / half
    axb = rel.abs().argmax(1)
    nb = torch.zeros_like(pb)
    nb[arange, axb] = torch.sign(rel[arange, axb])
    n = torch.where(ok[:, None], nb, n)
    alb = torch.where(ok[:, None], S["ba"], alb)
    # light flag: recomputed cleanly from the final hit (ceiling patch)
    p = o + best_t[:, None] * d
    is_light = (best_t < 1e8) & (n[:, 1] == -1.0) & (p[:, 1] > 0.99) \
        & (p[:, 0].abs() <= LIGHT["half"]) & (p[:, 2].abs() <= LIGHT["half"])
    return best_t, n, alb, is_light


def cosine_hemisphere(n, rng):
    """Cosine-weighted directions about normals n (N,3)."""
    N = len(n)
    u1 = torch.rand(N, device=n.device, generator=rng)
    u2 = torch.rand(N, device=n.device, generator=rng)
    r = torch.sqrt(u1)
    phi = 2 * torch.pi * u2
    x = r * torch.cos(phi)
    y = r * torch.sin(phi)
    z = torch.sqrt((1 - u1).clamp_min(0))
    a = torch.where(n[:, 0:1].abs() > 0.9,
                    torch.tensor([0.0, 1.0, 0.0], device=n.device).expand_as(n),
                    torch.tensor([1.0, 0.0, 0.0], device=n.device).expand_as(n))
    t = torch.linalg.cross(a, n)
    t = t / t.norm(dim=1, keepdim=True).clamp_min(1e-9)
    b = torch.linalg.cross(n, t)
    return x[:, None] * t + y[:, None] * b + z[:, None] * n


def nee(p, n, alb, S, rng):
    """Next-event estimation toward the ceiling light. Returns (N,3)."""
    N = len(p)
    dev = p.device
    u = (torch.rand(N, 2, device=dev, generator=rng) * 2 - 1) * LIGHT["half"]
    lp = torch.stack([u[:, 0], torch.full((N,), LIGHT["y"], device=dev),
                      u[:, 1]], 1)
    dl = lp - p
    dist = dl.norm(dim=1).clamp_min(1e-6)
    dl = dl / dist[:, None]
    cos_s = (n * dl).sum(1).clamp_min(0)
    # light normal is (0,-1,0); the emission cosine is between it and the
    # direction light->surface (=-dl), i.e. (0,-1,0)·(-dl) = +dl_y
    cos_l = dl[:, 1].clamp_min(0)
    t, _, _, _ = intersect(p + EPS * n, dl, S)
    vis = t > dist - 3e-3
    area = (2 * LIGHT["half"]) ** 2
    g = cos_s * cos_l / (dist ** 2)
    return (alb / torch.pi) * S["Le"] * (g * vis * area)[:, None]


def trace_split(o, d, S, rng, depth=5):
    """Path trace with NEE; returns (emitted, direct, indirect) per ray.

    emitted:  light seen directly by the given ray
    direct:   NEE at the FIRST hit (analytic given visibility)
    indirect: everything after the first bounce — the part the neural
              cache learns
    Also returns the first-hit geometry (t, n, albedo) for cache training.
    """
    N = len(o)
    dev = o.device
    emitted = torch.zeros(N, 3, device=dev)
    direct = torch.zeros(N, 3, device=dev)
    indirect = torch.zeros(N, 3, device=dev)
    tput = torch.ones(N, 3, device=dev)
    alive = torch.ones(N, dtype=torch.bool, device=dev)
    first = {}
    co, cd = o.clone(), d.clone()
    for depth_i in range(depth):
        t, n, alb, isl = intersect(co, cd, S)
        hit = (t < 1e8) & alive
        p = co + t[:, None] * cd
        if depth_i == 0:
            first = dict(t=t.clone(), n=n.clone(), alb=alb.clone(),
                         p=p.clone(), hit=hit.clone() & ~isl)
            emitted[hit & isl] = S["Le"]
        alive = alive & hit & ~isl
        # no .any() early-out: at these depths the surviving-ray count is
        # data-dependent and the sync costs more than the wasted work
        contrib = torch.zeros(N, 3, device=dev)
        contrib[alive] = nee(p[alive], n[alive], alb[alive], S, rng)
        contrib = contrib * tput
        if depth_i == 0:
            direct = direct + contrib
        else:
            indirect = indirect + contrib
        # bounce
        nd = torch.zeros_like(cd)
        nd[alive] = cosine_hemisphere(n[alive], rng)
        tput = tput * alb
        co = p + EPS * n
        cd = nd
    return emitted, direct, indirect, first


def camera_rays(res, jitter, dev, rng, cam=(0.0, 0.0, 3.05), fov=0.62):
    ys, xs = torch.meshgrid(
        torch.linspace(1, -1, res, device=dev),
        torch.linspace(-1, 1, res, device=dev), indexing="ij")
    if jitter:
        xs = xs + (torch.rand(res, res, device=dev, generator=rng) - .5) * (2 / res)
        ys = ys + (torch.rand(res, res, device=dev, generator=rng) - .5) * (2 / res)
    d = torch.stack([xs * fov, ys * fov, -torch.ones_like(xs)], -1).reshape(-1, 3)
    d = d / d.norm(dim=1, keepdim=True)
    o = torch.tensor(cam, device=dev).expand_as(d).contiguous()
    return o, d


def scene_tensors(dev):
    return _dev(dev)