"""W9: neural radiance cache — neural ray tracing on the project skeleton.

Claim under test (the light-transport analog of the physics thesis): if
you keep visibility and direct lighting analytic and learn ONLY the
indirect transport with a tiny tied MLP, you match a many-bounce path
tracer at a fraction of the cost — one analytic ray + NEE + a network
lookup replaces the random walk after the first bounce.

Render decomposition:  L = emitted + direct(analytic) + indirect
  - reference:  indirect from high-spp path tracing (ground truth)
  - neural:     indirect from cache_theta(first-hit point, normal)
  - baseline:   an equal-COST path trace (few spp) for a fair comparison

Success criteria:
  A. cache accuracy on HELD-OUT surface points (never-seen view) — PSNR
     of predicted vs true indirect radiance
  B. full-image PSNR: neural-cached vs reference, and it beats the
     equal-cost path-traced baseline
  C. speedup at matched quality (spp the baseline needs to reach the
     neural render's PSNR)
"""
import sys, os, time
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import numpy as np
import torch
from engine3d.raytrace import (camera_rays, trace_split, scene_tensors,
                               intersect)
from engine3d.neural_rt import RadianceCache

torch.manual_seed(0)
HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
DEV = "cuda" if torch.cuda.is_available() else "cpu"
S = scene_tensors(DEV)


def rng_of(seed):
    g = torch.Generator(device=DEV); g.manual_seed(seed); return g


def render(res, spp, seed, cam=(0.0, 0.0, 3.05), depth=5):
    """Full reference render (emitted + direct + indirect), returns (res,res,3)."""
    g = rng_of(seed)
    acc = torch.zeros(res * res, 3, device=DEV)
    for _ in range(spp):
        o, d = camera_rays(res, True, DEV, g, cam=cam)
        em, di, ind, _ = trace_split(o, d, S, g, depth=depth)
        acc += em + di + ind
    return (acc / spp).reshape(res, res, 3)


def render_components(res, spp, seed, cam=(0.0, 0.0, 3.05)):
    """Per-pixel emitted+direct (low variance) and first-hit geometry, plus
    a high-spp indirect target. Used to compose the neural image."""
    g = rng_of(seed)
    ed = torch.zeros(res * res, 3, device=DEV)      # emitted + direct
    ind = torch.zeros(res * res, 3, device=DEV)
    o, d = camera_rays(res, False, DEV, g, cam=cam)  # no jitter: fixed hits
    for _ in range(spp):
        em, di, ii, first = trace_split(o, d, S, g, depth=5)
        ed += em + di
        ind += ii
    ed /= spp; ind /= spp
    o2, d2 = camera_rays(res, False, DEV, rng_of(seed), cam=cam)
    _, n, alb, isl = intersect(o2, d2, S)
    t, _, _, _ = intersect(o2, d2, S)
    p = o2 + t[:, None] * d2
    hit = (t < 1e8) & ~isl
    return ed, ind, p, n, hit


def psnr(a, b, mask=None):
    if mask is not None:
        a, b = a[mask], b[mask]
    mse = ((a.clamp(0, 4) - b.clamp(0, 4)) ** 2).mean().item()
    return 10 * np.log10(4.0 ** 2 / max(mse, 1e-12))


# ---------- training data: surface points with GT indirect ----------
print(f"=== generating radiance-cache training data on {DEV} ===")
t0 = time.time()
PTS, NRM, TGT = [], [], []
CAM_TRAIN = [(0.0, 0.0, 3.05), (0.7, 0.2, 2.9), (-0.7, 0.25, 2.9),
             (0.0, 0.6, 2.8), (0.4, -0.3, 3.0)]
for ci, cam in enumerate(CAM_TRAIN):
    g = rng_of(100 + ci)
    o, d = camera_rays(80, True, DEV, g, cam=cam)
    t, n, alb, isl = intersect(o, d, S)
    p = o + t[:, None] * d
    hit = (t < 1e8) & ~isl
    p, n = p[hit], n[hit]
    # GT indirect at these points: average many one-bounce-onward paths
    ind = torch.zeros(len(p), 3, device=DEV)
    SPP = 256
    for _ in range(SPP):
        from engine3d.raytrace import cosine_hemisphere, nee, EPS
        nd = cosine_hemisphere(n, g)
        # radiance arriving from the bounce dir, path-traced (depth 4),
        # times the cosine-weighted albedo throughput (albedo/pi * pi = albedo)
        _, ndi, nind, _ = trace_split(p + EPS * n, nd, S, g, depth=4)
        ind += (ndi + nind)
    ind /= SPP
    alb_h = alb[hit]
    ind = ind * alb_h                     # outgoing = albedo * incident indirect
    PTS.append(p); NRM.append(n); TGT.append(ind)
P = torch.cat(PTS); Nrm = torch.cat(NRM); Y = torch.cat(TGT)
print(f"  {len(P)} surface samples from {len(CAM_TRAIN)} views "
      f"({time.time()-t0:.1f}s)")

# ---------- train the cache ----------
cache = RadianceCache().to(DEV)
print(f"=== training radiance cache ({cache.n_params()} params) ===")
opt = torch.optim.Adam(cache.parameters(), lr=3e-3)
sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=1500)
# tone-mapped (log) loss: indirect radiance spans orders of magnitude and
# the eye is roughly logarithmic — regress log(1+L), like real NRC
logY = torch.log1p(Y)
t0 = time.time()
for it in range(1500):
    idx = torch.randint(0, len(P), (16384,), device=DEV)
    pred = cache(P[idx], Nrm[idx])
    loss = ((torch.log1p(pred) - logY[idx]) ** 2).mean()
    opt.zero_grad(); loss.backward(); opt.step(); sched.step()
    if it % 300 == 0 or it == 1499:
        print(f"  it {it:4d}  log-MSE {loss.item():.5f}")
print(f"  ({time.time()-t0:.1f}s)")

# ---------- criterion A: held-out points (novel view) ----------
g = rng_of(999)
cam_test = (0.35, -0.15, 2.95)
o, d = camera_rays(80, True, DEV, g, cam=cam_test)
t, n, alb, isl = intersect(o, d, S)
p = o + t[:, None] * d
hit = (t < 1e8) & ~isl
p, n, alb = p[hit], n[hit], alb[hit]
ind = torch.zeros(len(p), 3, device=DEV)
from engine3d.raytrace import cosine_hemisphere, EPS
for _ in range(256):
    nd = cosine_hemisphere(n, g)
    _, ndi, nind, _ = trace_split(p + EPS * n, nd, S, g, depth=4)
    ind += (ndi + nind)
ind = ind / 256 * alb
with torch.no_grad():
    pred = cache(p, n)
a_psnr = psnr(pred, ind)
print(f"\n[A] held-out indirect radiance PSNR: {a_psnr:.2f} dB "
      f"(novel view, {len(p)} pts)")

# ---------- criterion B & C: full-image renders ----------
RES = 112
REF_SPP = 256
print(f"\n=== full renders at {RES}x{RES} ===")
t0 = time.time(); REF = render(RES, REF_SPP, 7); ref_ms = (time.time()-t0)*1000
print(f"  reference ({REF_SPP} spp): {ref_ms:.0f} ms")

# neural: analytic emitted+direct at low spp + cache indirect
t0 = time.time()
ed, _, p, n, hit = render_components(RES, 4, 7)
with torch.no_grad():
    ind_pred = torch.zeros(RES * RES, 3, device=DEV)
    ind_pred[hit] = cache(p[hit], n[hit])
NEUR = (ed + ind_pred).reshape(RES, RES, 3)
neur_ms = (time.time()-t0)*1000
print(f"  neural (4 spp direct + cache): {neur_ms:.0f} ms")

# baseline: equal-cost full path trace
def spp_for_ms(target_ms):
    t0 = time.time(); render(RES, 4, 7); one = (time.time()-t0)/4*1000
    return max(1, round(target_ms / one)), one
base_spp, per = spp_for_ms(neur_ms)
BASE = render(RES, base_spp, 7)

pn = psnr(NEUR, REF)
pb = psnr(BASE, REF)
print(f"\n[B] full-image PSNR vs reference:")
print(f"     neural (cache)          : {pn:.2f} dB  ({neur_ms:.0f} ms)")
print(f"     path trace (equal cost) : {pb:.2f} dB  ({base_spp} spp)")

# criterion C: spp the baseline needs to match the neural PSNR
match_spp = None
for spp in [8, 16, 32, 64, 128]:
    if psnr(render(RES, spp, 7), REF) >= pn:
        match_spp = spp; break
base_ms = (match_spp or 128) * per
print(f"\n[C] path trace needs ~{match_spp or '>256'} spp "
      f"(~{base_ms:.0f} ms) to reach the neural PSNR "
      f"-> ~{base_ms/neur_ms:.1f}x speedup at matched quality")

# ---------- figure ----------
try:
    import matplotlib; matplotlib.use("Agg"); import matplotlib.pyplot as plt
    def tm(x): return np.clip((x.clamp(0, 4) ** (1/2.2)).cpu().numpy(), 0, 1)
    err = (NEUR - REF).abs().mean(-1).cpu().numpy()
    errb = (BASE - REF).abs().mean(-1).cpu().numpy()
    fig, ax = plt.subplots(2, 3, figsize=(12, 8))
    for a in ax.ravel(): a.set_xticks([]); a.set_yticks([])
    ax[0,0].imshow(tm(REF)); ax[0,0].set_title(f"reference ({REF_SPP} spp)")
    ax[0,1].imshow(tm(NEUR)); ax[0,1].set_title(f"neural cache — {pn:.1f} dB, {neur_ms:.0f} ms")
    ax[0,2].imshow(tm(BASE)); ax[0,2].set_title(f"equal-cost path trace ({base_spp} spp) — {pb:.1f} dB")
    ax[1,0].imshow(tm(ind_pred.reshape(RES,RES,3))); ax[1,0].set_title("learned indirect (cache only)")
    m=max(err.max(),errb.max())
    ax[1,1].imshow(err,cmap="inferno",vmax=m); ax[1,1].set_title("neural error")
    ax[1,2].imshow(errb,cmap="inferno",vmax=m); ax[1,2].set_title("path-trace error")
    fig.suptitle(f"Neural radiance cache — learn indirect, keep visibility+direct analytic "
                 f"(held-out cache PSNR {a_psnr:.1f} dB)", fontsize=12)
    fig.tight_layout()
    out=os.path.join(ROOT,"neural_raytrace_validation.png")
    fig.savefig(out,dpi=110); print(f"\nwrote {out}")
except Exception as e:
    print("plot skipped:", e)

torch.save({"state_dict": cache.state_dict(),
            "arch": {"n_freq": 6, "hidden": 96},
            "scene": "neon_cornell"},
           os.path.join(HERE, "radiance_cache.pt"))
print("saved experiments/radiance_cache.pt")