"""Neural radiance cache (W9 — neural ray tracing).

The same thesis as the physics engine, applied to light transport: keep
the exact parts analytic (ray/visibility, next-event direct lighting),
learn only the expensive part (multi-bounce indirect transport).

A path becomes:  L = emitted + direct(analytic NEE) + cache_θ(x, n)
where cache_θ is one tiny MLP, tied across the whole scene, that maps a
surface point + normal to its outgoing indirect radiance. This replaces
the random walk after the first bounce with a single network lookup —
Müller et al.'s "Real-time Neural Radiance Caching" idea, distilled to
the project's local-projection skeleton.

Positional (frequency) encoding is essential: a plain MLP on raw xyz
cannot fit the high-frequency shading near contact shadows and the
color-bleeding corners. The encoding is the analytic structure we keep;
the MLP only learns amplitudes.
"""
import numpy as np
import torch
import torch.nn as nn


class FreqEncoding(nn.Module):
    """NeRF-style sin/cos encoding: x -> [x, sin(2^k x), cos(2^k x)]_k."""

    def __init__(self, n_freq=6):
        super().__init__()
        self.register_buffer("bands", 2.0 ** torch.arange(n_freq) * torch.pi)
        self.out_mult = 1 + 2 * n_freq

    def forward(self, x):
        proj = x[..., None] * self.bands            # (...,D,F)
        enc = torch.cat([torch.sin(proj), torch.cos(proj)], -1)
        return torch.cat([x, enc.flatten(-2)], -1)


class RadianceCache(nn.Module):
    """Tied MLP: (position, normal) -> outgoing indirect radiance (RGB).

    Shared across every surface point in the scene (the tied-embedding
    thesis, now over spatial location instead of mesh element). Output is
    non-negative (softplus) — radiance cannot be negative, a guard baked
    into the architecture rather than clamped after the fact.
    """

    def __init__(self, n_freq=6, hidden=96):
        super().__init__()
        self.enc = FreqEncoding(n_freq)
        din = self.enc.out_mult * 3 + 3          # encoded position + raw normal
        self.net = nn.Sequential(
            nn.Linear(din, hidden), nn.SiLU(),
            nn.Linear(hidden, hidden), nn.SiLU(),
            nn.Linear(hidden, hidden), nn.SiLU(),
            nn.Linear(hidden, 3),
        )

    def forward(self, pos, nrm):
        h = torch.cat([self.enc(pos), nrm], -1)
        return torch.nn.functional.softplus(self.net(h))

    def n_params(self):
        return sum(p.numel() for p in self.parameters())