"""Small physical shadow-ray oracle with an exact finite-light reference.

Lambertian point receivers on z=0, three opaque sphere occluders, and a 6x6
array of point lights at z=2.2. This is direct illumination only. The scene
does not contain indirect transport, specular receivers, or camera antialiasing.
"""
from dataclasses import dataclass
import numpy as np


@dataclass
class Scene:
    seed: int
    spheres: np.ndarray  # [3,4] = x,y,z,radius, renderer-owned geometry

    def __post_init__(self):
        self.spheres=np.asarray(self.spheres,dtype=float)
        if self.spheres.shape!=(3,4) or not np.isfinite(self.spheres).all() or (self.spheres[:,3]<=0).any():
            raise ValueError("Reference scene requires three finite spheres with positive radii")

    @classmethod
    def create(cls, seed):
        rng = np.random.default_rng(seed)
        xyz = rng.uniform([-0.72,-0.72,0.45], [0.72,0.72,0.95], size=(3,3))
        radius = rng.uniform(0.17,0.36, size=(3,1))
        spheres = np.concatenate([xyz, radius], axis=1)
        return cls(int(seed), spheres[np.argsort(spheres[:, 0])])

    def changed(self):
        moved = self.spheres.copy()
        moved[0, 0] += 0.65
        moved[1, 1] -= 0.6
        moved[2, 3] *= 1.25
        return Scene(self.seed, moved)

    def visibility(self, receivers, lights):
        """Exact segment/sphere visibility for matching broadcastable arrays."""
        p, l = np.broadcast_arrays(np.asarray(receivers, float), np.asarray(lights, float))
        if p.shape[-1]!=3 or not np.isfinite(p).all() or not np.isfinite(l).all():
            raise ValueError("Physical segment endpoints must be finite 3D points")
        d = l-p
        a = (d*d).sum(-1)
        if (a <= 0).any():
            raise ValueError("A physical shadow segment must have positive length")
        visible = np.ones(a.shape, bool)
        for sphere in self.spheres:
            oc = p-sphere[:3]
            b = (oc*d).sum(-1)
            c = (oc*oc).sum(-1)-sphere[3]**2
            discriminant = b*b-a*c
            root = np.sqrt(np.maximum(discriminant,0))
            near, far = (-b-root)/a, (-b+root)/a
            intersects = (discriminant >= 0) & (far > 1e-6) & (near < 1-1e-6)
            visible &= ~intersects
        return visible.astype(np.float64)

    def features(self, receivers, lights):
        p, l = np.broadcast_arrays(np.asarray(receivers, float), np.asarray(lights, float))
        geom = np.broadcast_to(self.spheres.ravel(), p.shape[:-1]+(12,))
        return np.concatenate([p[..., :2], l[..., :2], geom], axis=-1).astype(np.float32)


def receiver_grid(height=32, width=64):
    x = np.linspace(-1,1,width)
    y = np.linspace(-1,1,height)
    xx,yy = np.meshgrid(x,y)
    return np.stack([xx,yy,np.zeros_like(xx)],axis=-1).reshape(-1,3)


def light_grid(side=6):
    x = np.linspace(-0.85,0.85,side)
    xx,yy = np.meshgrid(x,x)
    return np.stack([xx,yy,np.full_like(xx,2.2)],axis=-1).reshape(-1,3)


def albedo(points, phase=0):
    x,y = points[:,0],points[:,1]
    checker = ((np.floor((x+1)*10)+np.floor((y+1)*10))%2)
    r = 0.22+0.42*checker
    g = 0.25+0.25*(0.5+0.5*np.sin(11*x+phase))
    b = 0.2+0.38*(0.5+0.5*np.sin(13*y-0.3+phase))
    return np.stack([r,g,b],axis=-1)


def lighting(lights, changed=False):
    x,y=lights[:,0],lights[:,1]
    if changed:
        rgb=np.stack([0.3+1.7*(x>0),0.4+0.3*np.cos(3*y)**2,0.5+1.2*(x<0)],axis=-1)
    else:
        rgb=np.stack([0.8+0.4*np.cos(2*x)**2,0.9+0.2*np.sin(y)**2,0.7+0.2*np.cos(3*y)**2],axis=-1)
    return 11.0*rgb/len(lights)  # each emitter's radiant intensity


def unoccluded(points, lights, changed_lighting=False, material_phase=0):
    d=lights[None,:,:]-points[:,None,:]
    distance=np.linalg.norm(d,axis=-1)
    cosine=np.maximum(d[...,2]/distance,0)
    geometry=cosine/(np.pi*distance**2)
    return geometry[...,None]*albedo(points,material_phase)[:,None,:]*lighting(lights,changed_lighting)[None,:,:]


def physical_table(scene, points, lights, bound):
    """Privileged offline reference; do not pass this table to online policy."""
    vis=scene.visibility(points[:,None,:],lights[None,:,:])
    return bound*vis[...,None]


class VisibilityPrior:
    """Portable NumPy evaluation of the trained 16-48-48-1 MLP."""
    def __init__(self,path):
        with np.load(path,allow_pickle=False) as data:
            self.arrays={k:np.array(data[k]) for k in data.files}
        required={"mean","scale","w0","b0","w1","b1","w2","b2"}
        if set(self.arrays)!=required or not all(np.isfinite(v).all() for v in self.arrays.values()):
            raise ValueError("Malformed prior weights")
        shapes={"mean":(16,),"scale":(16,),"w0":(48,16),"b0":(48,),"w1":(48,48),"b1":(48,),"w2":(1,48),"b2":(1,)}
        if any(self.arrays[k].shape!=s for k,s in shapes.items()) or (self.arrays["scale"]<=0).any():
            raise ValueError("Invalid prior dimensions/scaling")

    def __call__(self,features):
        features=np.asarray(features,np.float32)
        if features.ndim<1 or features.shape[-1]!=16 or not np.isfinite(features).all():
            raise ValueError("Visibility prior requires finite 16-dimensional features")
        original_shape=features.shape[:-1]
        x=np.asarray(features,np.float32).reshape(-1,16)
        a=self.arrays
        outputs=[]
        for start in range(0,len(x),8192):
            h=(x[start:start+8192]-a["mean"])/a["scale"]
            h=np.maximum(h@a["w0"].T+a["b0"],0)
            h=np.maximum(h@a["w1"].T+a["b1"],0)
            h=h@a["w2"].T+a["b2"]
            outputs.append(1/(1+np.exp(-np.clip(h,-40,40))))
        if not outputs:
            return np.empty(original_shape,dtype=np.float32)
        return np.concatenate(outputs).reshape(original_shape)