"""Predicted bounded spans over completed causal segments."""

from __future__ import annotations

from dataclasses import dataclass

import torch
from torch import nn

from strata.modeling.ph_pat.config import PHPATConfig
from strata.modeling.ph_pat.segment_commit import SegmentLayout


@dataclass(frozen=True, slots=True)
class SpanCompilerOutput:
    nodes: torch.Tensor
    scores: torch.Tensor
    starts: torch.Tensor
    ends: torch.Tensor
    segment_ids: torch.Tensor
    valid_mask: torch.Tensor
    boundary_logits: torch.Tensor


class SpanCompiler(nn.Module):
    def __init__(self, config: PHPATConfig) -> None:
        super().__init__()
        self.config = config
        self.start = nn.Linear(config.d_model, config.d_model, bias=False)
        self.end = nn.Linear(config.d_model, config.d_model, bias=False)
        self.pool = nn.Linear(config.d_model, config.d_model, bias=False)
        self.score = nn.Sequential(
            nn.Linear(3 * config.d_model, config.d_model),
            nn.SiLU(),
            nn.Linear(config.d_model, 1),
        )
        self.boundary = nn.Linear(config.d_model, 2)
        relative_starts, relative_ends = _candidate_bounds(0, config.segment_size, config.max_span_width, torch.device("cpu"))
        self.register_buffer("relative_starts", relative_starts, persistent=False)
        self.register_buffer("relative_ends", relative_ends, persistent=False)

    def forward(self, hidden: torch.Tensor, layout: SegmentLayout) -> SpanCompilerOutput:
        batch, _seq_len, dim = hidden.shape
        per_segment = self.config.spans_per_segment
        prefix = torch.cat((hidden.new_zeros(batch, 1, dim), hidden.cumsum(dim=1)), dim=1)
        boundary_logits = self.boundary(hidden)
        segment_offsets = layout.starts.view(1, layout.segment_count, 1)
        candidate_starts = segment_offsets + self.relative_starts.view(1, 1, -1)
        candidate_ends = segment_offsets + self.relative_ends.view(1, 1, -1)
        in_sequence = candidate_ends < hidden.shape[1]
        safe_starts = candidate_starts.clamp_max(hidden.shape[1] - 1)
        safe_ends = candidate_ends.clamp_max(hidden.shape[1] - 1)
        row = torch.arange(batch, device=hidden.device).view(batch, 1, 1)
        start_score = boundary_logits[row, safe_starts.expand(batch, -1, -1), 0]
        end_score = boundary_logits[row, safe_ends.expand(batch, -1, -1), 1]
        candidate_valid = layout.segment_valid.unsqueeze(-1) & in_sequence
        candidate_scores = (start_score + end_score).masked_fill(~candidate_valid, torch.finfo(hidden.dtype).min)
        take = min(per_segment, candidate_scores.shape[-1])
        top_scores, top_idx = torch.topk(candidate_scores, k=take, dim=-1)
        expanded_starts = candidate_starts.expand(batch, -1, -1)
        expanded_ends = candidate_ends.expand(batch, -1, -1)
        top_starts = expanded_starts.gather(2, top_idx).clamp_max(hidden.shape[1] - 1)
        top_ends = expanded_ends.gather(2, top_idx).clamp_max(hidden.shape[1] - 1)
        valid = candidate_valid.expand(batch, -1, -1).gather(2, top_idx)
        batch_index = torch.arange(batch, device=hidden.device).view(batch, 1, 1)
        span_sum = prefix[batch_index, top_ends + 1] - prefix[batch_index, top_starts]
        pooled = span_sum / (top_ends - top_starts + 1).to(hidden.dtype).unsqueeze(-1)
        start_hidden = hidden[batch_index, top_starts]
        end_hidden = hidden[batch_index, top_ends]
        top_nodes = (self.start(start_hidden) + self.end(end_hidden) + self.pool(pooled)) / 3.0
        selected_features = torch.cat((start_hidden, end_hidden, pooled), dim=-1)
        top_scores = top_scores + self.score(selected_features).squeeze(-1)

        if take < per_segment:
            pad = per_segment - take
            top_nodes = torch.cat((top_nodes, hidden.new_zeros(batch, layout.segment_count, pad, dim)), dim=2)
            top_scores = torch.cat((top_scores, hidden.new_full((batch, layout.segment_count, pad), torch.finfo(hidden.dtype).min)), dim=2)
            top_starts = torch.cat((top_starts, torch.zeros(batch, layout.segment_count, pad, device=hidden.device, dtype=torch.long)), dim=2)
            top_ends = torch.cat((top_ends, torch.zeros(batch, layout.segment_count, pad, device=hidden.device, dtype=torch.long)), dim=2)
            valid = torch.cat((valid, torch.zeros(batch, layout.segment_count, pad, device=hidden.device, dtype=torch.bool)), dim=2)
        segment_ids = torch.arange(layout.segment_count, device=hidden.device).view(1, -1, 1).expand(batch, -1, per_segment)

        return SpanCompilerOutput(
            nodes=top_nodes.flatten(1, 2),
            scores=top_scores.flatten(1, 2),
            starts=top_starts.flatten(1, 2),
            ends=top_ends.flatten(1, 2),
            segment_ids=segment_ids.flatten(1, 2),
            valid_mask=valid.flatten(1, 2),
            boundary_logits=boundary_logits,
        )


def _candidate_bounds(start: int, end: int, max_width: int, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]:
    starts: list[int] = []
    ends: list[int] = []
    for left in range(start, end):
        for right in range(left, min(end, left + max_width)):
            starts.append(left)
            ends.append(right)
    return torch.tensor(starts, device=device), torch.tensor(ends, device=device)


__all__ = ["SpanCompiler", "SpanCompilerOutput"]