"""Parameter-free composable measurements read from an explicit scene state."""

from __future__ import annotations

from collections.abc import Iterable

import torch
import torch.nn.functional as functional
from torch import Tensor, nn

from mcss.geometry import generate_rays, intersect_aabb
from mcss.types import Cameras, SceneState

MEASUREMENTS = frozenset({"rgb", "depth", "normal", "point", "visibility", "uncertainty"})


class FixedMeasurementRenderer(nn.Module):
    """Read all primary scene measurements with fixed grid sampling and ray integration.

    The module owns no learned tensors. It is deliberately separate from the context encoder so
    a no-bypass audit can verify that predicted targets originate from typed scene fields.
    """

    def __init__(
        self,
        *,
        n_samples: int = 64,
        ray_chunk_size: int = 16_384,
        background_color: tuple[float, float, float] = (0.0, 0.0, 0.0),
    ) -> None:
        super().__init__()
        if n_samples < 2:
            raise ValueError("n_samples must be at least two")
        if ray_chunk_size < 1:
            raise ValueError("ray_chunk_size must be positive")
        if len(background_color) != 3:
            raise ValueError("background_color must contain RGB values")
        self.n_samples = n_samples
        self.ray_chunk_size = ray_chunk_size
        self.background_color = tuple(float(value) for value in background_color)

    def forward(
        self,
        state: SceneState,
        cameras: Cameras,
        measurements: Iterable[str] | None = None,
    ) -> dict[str, Tensor]:
        """Render requested measurements as tensors shaped ``[B, V, C, H, W]``."""

        requested = MEASUREMENTS if measurements is None else frozenset(measurements)
        unknown = requested - MEASUREMENTS
        if unknown:
            raise ValueError(f"unknown measurement names: {sorted(unknown)}")
        outputs, _ = self._render(state, cameras, include_features=False)
        return {name: outputs[name] for name in sorted(requested)}

    def render_features(self, state: SceneState, cameras: Cameras) -> Tensor:
        """Integrate optional state features for matched learned-decoder controls only."""

        if state.features is None:
            raise ValueError("render_features requires SceneState.features")
        _, features = self._render(state, cameras, include_features=True)
        if features is None:
            raise RuntimeError("feature integration unexpectedly produced no features")
        return features

    def _render(
        self, state: SceneState, cameras: Cameras, *, include_features: bool
    ) -> tuple[dict[str, Tensor], Tensor | None]:
        if len(cameras.leading_shape) != 2:
            raise ValueError("renderer expects cameras with shape [B, V, ...]")
        batch_size, _ = cameras.leading_shape
        if batch_size != state.density_logits.shape[0]:
            raise ValueError("camera and state batch sizes must match")
        if state.density_logits.device != cameras.device:
            raise ValueError("state and cameras must be on the same device")

        origins, directions = generate_rays(cameras)
        _, view_count, height, width, _ = origins.shape
        ray_count = view_count * height * width
        origins = origins.reshape(batch_size, ray_count, 3)
        directions = directions.reshape(batch_size, ray_count, 3)
        near, far, hit = intersect_aabb(origins, directions, state.bounds)
        normal_grid = _state_normal_grid(state)
        color_volume = state.color if state.appearance is None else state.appearance.color

        rendered: dict[str, list[Tensor]] = {name: [] for name in MEASUREMENTS}
        rendered_features: list[Tensor] = []
        for start in range(0, ray_count, self.ray_chunk_size):
            stop = min(start + self.ray_chunk_size, ray_count)
            chunk_outputs, chunk_features = self._render_chunk(
                state,
                normal_grid,
                color_volume,
                origins[:, start:stop],
                directions[:, start:stop],
                near[:, start:stop],
                far[:, start:stop],
                hit[:, start:stop],
                include_features=include_features,
            )
            for name, value in chunk_outputs.items():
                rendered[name].append(value)
            if chunk_features is not None:
                rendered_features.append(chunk_features)

        output = {
            name: _reshape_image(torch.cat(parts, dim=1), batch_size, view_count, height, width)
            for name, parts in rendered.items()
        }
        feature_output = None
        if include_features:
            feature_output = _reshape_image(
                torch.cat(rendered_features, dim=1), batch_size, view_count, height, width
            )
        return output, feature_output

    def _render_chunk(
        self,
        state: SceneState,
        normal_grid: Tensor,
        color_volume: Tensor,
        origins: Tensor,
        directions: Tensor,
        near: Tensor,
        far: Tensor,
        hit: Tensor,
        *,
        include_features: bool,
    ) -> tuple[dict[str, Tensor], Tensor | None]:
        batch_size, ray_count, _ = origins.shape
        fractions = (
            torch.arange(self.n_samples, device=origins.device, dtype=origins.dtype) + 0.5
        ) / self.n_samples
        distances = near.unsqueeze(-1) + (far - near).unsqueeze(-1) * fractions
        points = origins.unsqueeze(-2) + directions.unsqueeze(-2) * distances.unsqueeze(-1)
        sample_valid = hit.unsqueeze(-1)

        density_logits = _sample_volume(state.density_logits, points, state.bounds).squeeze(-1)
        colors = _sample_volume(color_volume, points, state.bounds)
        log_variance = _sample_volume(state.log_variance, points, state.bounds).squeeze(-1)
        with torch.autocast(device_type=points.device.type, enabled=False):
            normals = _sample_volume(normal_grid, points.float(), state.bounds.float())
        density = torch.nn.functional.softplus(density_logits) * sample_valid
        deltas = ((far - near) / self.n_samples).unsqueeze(-1)
        alpha = 1.0 - torch.exp(-density * deltas)
        transmittance = torch.cumprod(
            torch.cat((torch.ones_like(alpha[..., :1]), 1.0 - alpha + 1e-10), dim=-1), dim=-1
        )[..., :-1]
        weights = transmittance * alpha
        visibility = weights.sum(dim=-1, keepdim=True)
        depth = (weights * distances).sum(dim=-1, keepdim=True)
        rgb = (weights.unsqueeze(-1) * colors).sum(dim=-2)
        background = torch.as_tensor(
            self.background_color, device=rgb.device, dtype=rgb.dtype
        ).view(1, 1, 3)
        rgb = rgb + (1.0 - visibility) * background
        with torch.autocast(device_type=weights.device.type, enabled=False):
            normal = _normalize((weights.float().unsqueeze(-1) * normals).sum(dim=-2))
        point = origins + directions * depth
        uncertainty = (weights * log_variance.exp()).sum(dim=-1, keepdim=True)

        outputs = {
            "depth": depth,
            "normal": normal,
            "point": point,
            "rgb": rgb,
            "uncertainty": uncertainty,
            "visibility": visibility,
        }
        features = None
        if include_features:
            if state.features is None:
                raise ValueError("state features are required for learned controls")
            sampled_features = _sample_volume(state.features, points, state.bounds)
            features = (weights.unsqueeze(-1) * sampled_features).sum(dim=-2)
        return outputs, features


def _sample_volume(volume: Tensor, points: Tensor, bounds: Tensor) -> Tensor:
    """Trilinearly sample ``[B, C, D, H, W]`` at world points ``[B, R, S, 3]``."""

    minimum = bounds[:, 0].view(-1, 1, 1, 3)
    extent = (bounds[:, 1] - bounds[:, 0]).view(-1, 1, 1, 3)
    grid = (points - minimum) / extent * 2.0 - 1.0
    batch_size, rays, samples, _ = grid.shape
    grid = grid.reshape(batch_size, rays * samples, 1, 1, 3)
    sampled = functional.grid_sample(
        volume, grid, mode="bilinear", padding_mode="zeros", align_corners=False
    )
    return sampled.reshape(batch_size, volume.shape[1], rays, samples).permute(0, 2, 3, 1)


def _state_normal_grid(state: SceneState) -> Tensor:
    """Finite-difference density gradients converted to outward world normals."""

    with torch.autocast(device_type=state.density_logits.device.type, enabled=False):
        density = torch.nn.functional.softplus(state.density_logits.float())
        depth, height, width = state.spatial_shape
        extent = state.bounds[:, 1].float() - state.bounds[:, 0].float()
        dx = functional.pad(density[..., 2:] - density[..., :-2], (1, 1, 0, 0, 0, 0))
        dy = functional.pad(density[:, :, :, 2:, :] - density[:, :, :, :-2, :], (0, 0, 1, 1, 0, 0))
        dz = functional.pad(density[:, :, 2:, :, :] - density[:, :, :-2, :, :], (0, 0, 0, 0, 1, 1))
        x_spacing = (extent[:, 0] / width).view(-1, 1, 1, 1, 1)
        y_spacing = (extent[:, 1] / height).view(-1, 1, 1, 1, 1)
        z_spacing = (extent[:, 2] / depth).view(-1, 1, 1, 1, 1)
        gradient = torch.cat(
            (dx / (2.0 * x_spacing), dy / (2.0 * y_spacing), dz / (2.0 * z_spacing)),
            dim=1,
        )
        return -_normalize(gradient, channel_dim=1)


def _reshape_image(values: Tensor, batch_size: int, views: int, height: int, width: int) -> Tensor:
    return values.reshape(batch_size, views, height, width, values.shape[-1]).permute(0, 1, 4, 2, 3)


def _normalize(values: Tensor, *, channel_dim: int = -1) -> Tensor:
    norm = torch.linalg.vector_norm(values, dim=channel_dim, keepdim=True)
    normalized = values / norm.clamp_min(1e-6)
    return torch.where(norm > 1e-6, normalized, torch.zeros_like(normalized))