"""Shared frame pooling, temporal pooling and reward heads for B1/C0/C1."""

from __future__ import annotations

import math
from dataclasses import dataclass

import torch
import torch.nn as nn


@dataclass(frozen=True)
class SpanPoolResult:
    frame_features: torch.Tensor
    frame_mask: torch.Tensor
    span_bounds: tuple[tuple[tuple[int, int], ...], ...]


def pool_image_token_spans(
    hidden_states: torch.Tensor,
    input_ids: torch.Tensor,
    image_token_id: int,
    image_grid_thw: torch.Tensor,
    frames_per_sample: list[int],
    spatial_merge_size: int,
) -> SpanPoolResult:
    """Mean-pool each image-token span and return padded per-frame features."""

    if hidden_states.ndim != 3 or input_ids.ndim != 2:
        raise ValueError("hidden_states/input_ids must have shapes [B,S,D] and [B,S]")
    if hidden_states.shape[:2] != input_ids.shape:
        raise ValueError("hidden_states and input_ids batch/sequence dimensions differ")
    if len(frames_per_sample) != hidden_states.shape[0] or any(
        count <= 0 for count in frames_per_sample
    ):
        raise ValueError(
            "frames_per_sample must contain one positive count per batch item"
        )
    if sum(frames_per_sample) != image_grid_thw.shape[0]:
        raise ValueError("frames_per_sample does not match image_grid_thw rows")
    if spatial_merge_size <= 0:
        raise ValueError("spatial_merge_size must be positive")

    max_frames = max(frames_per_sample)
    features = hidden_states.new_zeros(
        (hidden_states.shape[0], max_frames, hidden_states.shape[-1])
    )
    frame_mask = torch.zeros(
        (hidden_states.shape[0], max_frames),
        dtype=torch.bool,
        device=hidden_states.device,
    )
    all_bounds: list[tuple[tuple[int, int], ...]] = []
    grid_offset = 0
    merge_area = spatial_merge_size**2

    for batch_index, frame_count in enumerate(frames_per_sample):
        positions = torch.nonzero(
            input_ids[batch_index] == image_token_id, as_tuple=False
        ).flatten()
        token_offset = 0
        bounds: list[tuple[int, int]] = []
        for frame_index in range(frame_count):
            grid_t, grid_h, grid_w = (
                int(value) for value in image_grid_thw[grid_offset + frame_index]
            )
            raw_tokens = grid_t * grid_h * grid_w
            if raw_tokens % merge_area:
                raise ValueError(
                    "image token grid is not divisible by spatial merge area"
                )
            span_length = raw_tokens // merge_area
            frame_positions = positions[token_offset : token_offset + span_length]
            if frame_positions.numel() != span_length:
                raise ValueError("image token span is empty or truncated")
            if frame_positions.numel() > 1 and not torch.all(
                frame_positions[1:] == frame_positions[:-1] + 1
            ):
                raise ValueError("image token span is not contiguous")
            start = int(frame_positions[0])
            stop = int(frame_positions[-1]) + 1
            features[batch_index, frame_index] = hidden_states[
                batch_index, frame_positions
            ].mean(dim=0)
            frame_mask[batch_index, frame_index] = True
            bounds.append((start, stop))
            token_offset += span_length
        if token_offset != positions.numel():
            raise ValueError("unassigned image tokens remain after span pooling")
        all_bounds.append(tuple(bounds))
        grid_offset += frame_count

    return SpanPoolResult(features, frame_mask, tuple(all_bounds))


class MaskedQueryPool(nn.Module):
    """Single-query attention pooling over valid ordered-frame features."""

    def __init__(self, hidden_size: int) -> None:
        super().__init__()
        self.query = nn.Parameter(torch.empty(hidden_size))
        nn.init.normal_(self.query, mean=0.0, std=hidden_size**-0.5)

    def forward(
        self, frame_features: torch.Tensor, frame_mask: torch.Tensor
    ) -> torch.Tensor:
        if frame_features.ndim != 3 or frame_mask.shape != frame_features.shape[:2]:
            raise ValueError("expected frame_features [B,T,D] and frame_mask [B,T]")
        if not torch.all(frame_mask.any(dim=1)):
            raise ValueError("every trajectory must contain at least one valid frame")
        scores = torch.einsum("btd,d->bt", frame_features, self.query) / math.sqrt(
            frame_features.shape[-1]
        )
        scores = scores.masked_fill(~frame_mask, torch.finfo(scores.dtype).min)
        weights = torch.softmax(scores.float(), dim=1).to(frame_features.dtype)
        return torch.einsum("bt,btd->bd", weights, frame_features)


class MLPHead(nn.Module):
    def __init__(
        self, hidden_size: int, output_size: int, dropout: float = 0.1
    ) -> None:
        super().__init__()
        bottleneck = hidden_size // 2
        self.net = nn.Sequential(
            nn.Linear(hidden_size, bottleneck),
            nn.LayerNorm(bottleneck),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(bottleneck, output_size),
        )

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        return self.net(inputs)


class SharedRewardHeads(nn.Module):
    """The identical topology used for frozen probes and matched-LoRA runs."""

    def __init__(
        self, hidden_size: int, progress_bins: int = 10, dropout: float = 0.1
    ) -> None:
        super().__init__()
        self.temporal_pool = MaskedQueryPool(hidden_size)
        self.progress_head = MLPHead(hidden_size, progress_bins, dropout)
        self.success_head = MLPHead(hidden_size, 1, dropout)
        self.preference_head = MLPHead(hidden_size, 1, dropout)

    def forward(
        self, frame_features: torch.Tensor, frame_mask: torch.Tensor
    ) -> dict[str, torch.Tensor]:
        trajectory = self.temporal_pool(frame_features, frame_mask)
        return {
            "progress_logits": self.progress_head(frame_features),
            "success_logits": self.success_head(frame_features).squeeze(-1),
            "trajectory_features": trajectory,
            "trajectory_scores": self.preference_head(trajectory).squeeze(-1),
        }

    def preference_logit(
        self,
        frame_features_a: torch.Tensor,
        frame_mask_a: torch.Tensor,
        frame_features_b: torch.Tensor,
        frame_mask_b: torch.Tensor,
    ) -> torch.Tensor:
        score_a = self.preference_head(
            self.temporal_pool(frame_features_a, frame_mask_a)
        ).squeeze(-1)
        score_b = self.preference_head(
            self.temporal_pool(frame_features_b, frame_mask_b)
        ).squeeze(-1)
        return score_a - score_b