"""Training-only context-subset geometry consistency for the unchanged V5 state."""

from __future__ import annotations

from collections.abc import Mapping
from dataclasses import dataclass

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

from mcss.measurements import FixedMeasurementRenderer
from mcss.types import Cameras, SceneState


@dataclass(frozen=True)
class ContextSubset:
    """One context tensor/camera set with a single view removed in place."""

    rgb: Tensor
    cameras: Cameras
    dropped_index: int
    retained_indices: tuple[int, ...]


@dataclass(frozen=True)
class ContextSubsetGeometryResult:
    """Differentiable consistency loss plus detached audit telemetry."""

    loss: Tensor
    terms: dict[str, Tensor]
    diagnostics: dict[str, float]
    audit_diagnostics: dict[str, float]
    per_sample_diagnostics: tuple[dict[str, float], ...]


def drop_context_view(
    context_rgb: Tensor,
    cameras: Cameras,
    *,
    dropped_index: int,
) -> ContextSubset:
    """Remove one view without changing the established coordinate frame or order."""

    if context_rgb.ndim != 5 or context_rgb.shape[2] != 3:
        raise ValueError("context_rgb must have shape [B, V, 3, H, W]")
    batch_size, view_count = context_rgb.shape[:2]
    if cameras.leading_shape != (batch_size, view_count):
        raise ValueError("context cameras must align with context_rgb")
    if cameras.device != context_rgb.device:
        raise ValueError("context cameras and RGB must be on the same device")
    if view_count < 2:
        raise ValueError("context view deletion requires at least two views")
    if not 0 <= dropped_index < view_count:
        raise ValueError("dropped_index must select an existing context view")

    retained_indices = tuple(index for index in range(view_count) if index != dropped_index)
    indices = torch.tensor(retained_indices, device=context_rgb.device, dtype=torch.long)
    return ContextSubset(
        rgb=context_rgb.index_select(1, indices),
        cameras=cameras.select_views(indices),
        dropped_index=dropped_index,
        retained_indices=retained_indices,
    )


def ramped_context_subset_geometry_weight(
    max_weight: float,
    *,
    optimizer_step: int,
    warmup_steps: int,
) -> float:
    """Ramp against optimizer updates, not gradient-accumulation micro-steps."""

    if max_weight < 0.0:
        raise ValueError("max_weight must be non-negative")
    if optimizer_step < 1:
        raise ValueError("optimizer_step must be positive")
    if warmup_steps < 0:
        raise ValueError("warmup_steps must be non-negative")
    if max_weight == 0.0 or warmup_steps == 0:
        return float(max_weight)
    return float(max_weight) * min(optimizer_step / warmup_steps, 1.0)


def compute_context_subset_geometry(
    full_state: SceneState,
    subset_state: SceneState,
    full_predictions: Mapping[str, Tensor],
    target_cameras: Cameras,
    renderer: FixedMeasurementRenderer,
    *,
    residual_saturation_threshold: float,
    collect_audit_diagnostics: bool = True,
) -> ContextSubsetGeometryResult:
    """Compare subset completion geometry against a detached full-context teacher."""

    if full_state.evidence is None or subset_state.evidence is None:
        raise ValueError("context-subset geometry requires V5 StateEvidence on both paths")
    if residual_saturation_threshold <= 0.0:
        raise ValueError("residual_saturation_threshold must be positive")
    if "depth" not in full_predictions or "visibility" not in full_predictions:
        raise ValueError("full_predictions must contain depth and visibility")

    full_evidence = full_state.evidence
    subset_evidence = subset_state.evidence
    reference_shape = full_state.density_logits.shape
    if subset_state.density_logits.shape != reference_shape:
        raise ValueError("full and subset geometry grids must have identical shapes")
    if full_evidence.density_residual.shape != subset_evidence.density_residual.shape:
        raise ValueError("full and subset density residuals must have identical shapes")

    shared_support = (
        subset_evidence.provenance.detach().gt(1e-6).sum(dim=1, keepdim=True) >= 2
    )
    stable_evidence = (
        1.0
        - (
            full_evidence.confidence.detach()
            - subset_evidence.confidence.detach()
        ).abs()
    ).clamp(0.0, 1.0)
    overlap_weight = (
        shared_support.to(dtype=full_state.density_logits.dtype) * stable_evidence
    ).detach()

    teacher_density = full_state.density_logits.detach()
    subset_on_full_evidence = (
        full_evidence.base_density_logits.detach()
        + full_evidence.completion_gate.detach() * subset_evidence.density_residual
    )
    counterfactual_density = (
        overlap_weight * subset_on_full_evidence
        + (1.0 - overlap_weight) * teacher_density
    )
    counterfactual_state = SceneState(
        density_logits=counterfactual_density,
        color=full_state.color.detach(),
        log_variance=full_state.log_variance.detach(),
        bounds=full_state.bounds.detach(),
    )
    counterfactual_predictions = renderer(
        counterfactual_state,
        target_cameras,
        {"depth", "visibility"},
    )

    teacher_depth = full_predictions["depth"].detach().float()
    teacher_visibility = full_predictions["visibility"].detach().float()
    counterfactual_depth = counterfactual_predictions["depth"].float()
    counterfactual_visibility = counterfactual_predictions["visibility"].float()
    voxel_spacing = _mean_metric_voxel_spacing(
        full_state.bounds.detach(),
        full_state.spatial_shape,
    )
    depth_error = functional.smooth_l1_loss(
        counterfactual_depth / voxel_spacing,
        teacher_depth / voxel_spacing,
        reduction="none",
    )
    visible_mass = teacher_visibility.sum()
    depth_loss = (depth_error * teacher_visibility).sum() / visible_mass.clamp_min(1.0)
    visibility_loss = functional.smooth_l1_loss(
        counterfactual_visibility,
        teacher_visibility,
    )
    loss = depth_loss + visibility_loss

    residual_difference = (
        full_evidence.density_residual.detach()
        - subset_evidence.density_residual.detach()
    ).abs()
    diagnostics = {
        "full_evidence_mean": _detached_mean(full_evidence.confidence),
        "subset_evidence_mean": _detached_mean(subset_evidence.confidence),
        "evidence_abs_diff": _detached_mean(
            (full_evidence.confidence - subset_evidence.confidence).abs()
        ),
        "full_completion_gate_mean": _detached_mean(full_evidence.completion_gate),
        "subset_completion_gate_mean": _detached_mean(subset_evidence.completion_gate),
        "shared_coverage": _detached_mean(shared_support),
        "overlap_weight_mean": _detached_mean(overlap_weight),
        "residual_abs_diff": _detached_weighted_mean(
            residual_difference,
            overlap_weight,
        ),
        "counterfactual_depth_drift_voxels": _detached_weighted_mean(
            (counterfactual_depth.detach() - teacher_depth).abs() / voxel_spacing,
            teacher_visibility,
        ),
        "visibility_drift": _detached_mean(
            (counterfactual_visibility.detach() - teacher_visibility).abs()
        ),
    }
    audit_diagnostics: dict[str, float] = {}
    per_sample_diagnostics: tuple[dict[str, float], ...] = ()
    if collect_audit_diagnostics:
        audit_diagnostics = {
            "full_residual_saturation_fraction": _detached_mean(
                full_evidence.density_residual.detach().abs()
                >= residual_saturation_threshold
            ),
            "subset_residual_saturation_fraction": _detached_mean(
                subset_evidence.density_residual.detach().abs()
                >= residual_saturation_threshold
            ),
            **_residual_quantile_diagnostics(
                "full_residual_abs",
                full_evidence.density_residual,
            ),
            **_residual_quantile_diagnostics(
                "subset_residual_abs",
                subset_evidence.density_residual,
            ),
        }
        per_sample_diagnostics = tuple(
            {
                "full_evidence_mean": _detached_mean(full_evidence.confidence[index]),
                "subset_evidence_mean": _detached_mean(subset_evidence.confidence[index]),
                "full_completion_gate_mean": _detached_mean(
                    full_evidence.completion_gate[index]
                ),
                "subset_completion_gate_mean": _detached_mean(
                    subset_evidence.completion_gate[index]
                ),
                "full_residual_abs": _detached_mean(
                    full_evidence.density_residual[index].abs()
                ),
                "subset_residual_abs": _detached_mean(
                    subset_evidence.density_residual[index].abs()
                ),
                "residual_abs_diff": _detached_weighted_mean(
                    residual_difference[index],
                    overlap_weight[index],
                ),
                "shared_coverage": _detached_mean(shared_support[index]),
                "overlap_weight_mean": _detached_mean(overlap_weight[index]),
            }
            for index in range(reference_shape[0])
        )
    return ContextSubsetGeometryResult(
        loss=loss,
        terms={"depth": depth_loss, "visibility": visibility_loss},
        diagnostics=diagnostics,
        audit_diagnostics=audit_diagnostics,
        per_sample_diagnostics=per_sample_diagnostics,
    )


def _mean_metric_voxel_spacing(
    bounds: Tensor,
    spatial_shape: tuple[int, int, int],
) -> Tensor:
    depth, height, width = spatial_shape
    xyz_counts = bounds.new_tensor([width, height, depth])
    spacing = ((bounds[:, 1] - bounds[:, 0]) / xyz_counts).mean(dim=-1)
    return spacing.float().view(-1, 1, 1, 1, 1).clamp_min(1e-6)


def _detached_mean(values: Tensor) -> float:
    return float(values.detach().float().mean().item())


def _detached_weighted_mean(values: Tensor, weights: Tensor) -> float:
    values_float = values.detach().float()
    weights_float = weights.detach().float()
    weight_sum = weights_float.sum()
    if weight_sum.item() == 0.0:
        return 0.0
    return float((values_float * weights_float).sum().div(weight_sum).item())


def _residual_quantile_diagnostics(prefix: str, residual: Tensor) -> dict[str, float]:
    probabilities = torch.arange(0.1, 1.0, 0.1, device=residual.device)
    quantiles = torch.quantile(residual.detach().float().abs().reshape(-1), probabilities)
    return {
        f"{prefix}_p{index * 10:02d}": float(value.item())
        for index, value in enumerate(quantiles, start=1)
    }