"""Windowed inference for SignSegmenterV2.

Ported from /dev/4TB_SSD/segmentor_v2/models_Wilor.py::inference_wilor()
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import torch
from torch.nn.utils.rnn import pad_sequence

if TYPE_CHECKING:
    from sltk.segmentation._signsegmentor_v2.model import TransformerEncoderV2


def batch_data(data: torch.Tensor, window_size: int, stride: int) -> torch.Tensor:
    """Batch data into windows with given stride.

    Args:
        data: Input tensor of shape (T, D)
        window_size: Window size in frames
        stride: Stride between windows

    Returns:
        Batched tensor of shape (num_windows, window_size, D)
    """
    num_batches = (data.shape[0] - window_size) // stride + 1
    # Calculate the starting indices for each batch
    batch_start_indices = torch.arange(0, num_batches * stride, stride)
    # Use advanced indexing to extract the batches
    data = data[batch_start_indices.unsqueeze(1) + torch.arange(window_size)]
    return data


def inference_windowed(
    model: TransformerEncoderV2,
    features: torch.Tensor,
    window_size: int = 50,
    subsample: int = 2,
    device: str = "cuda",
) -> torch.Tensor:
    """Run windowed inference on WiLoR features.

    Splits features into overlapping windows, runs model inference,
    and reconstructs full sequence predictions.

    Args:
        model: Trained TransformerEncoderV2 model
        features: WiLoR features of shape (T, 192), assumed pre-subsampled
        window_size: Window size for batching (default 50)
        subsample: Subsample factor (default 2, but features should already be subsampled)
        device: Device for inference

    Returns:
        Predicted class labels of shape (T,)
    """
    model.eval()
    temp_pad = (model.temp_conv_kernel - 1) // 2

    # Get remainder frames
    a_overlap = features.shape[0] % window_size
    overlap_feature = features[-a_overlap:] if a_overlap > 0 else None

    # Batch full windows
    batch = batch_data(features, window_size, window_size)
    feat = []
    for b in batch:
        feat.append(b.squeeze(0))

    # Add remainder window
    if overlap_feature is not None:
        feat.append(overlap_feature)

    # Calculate sequence lengths
    src_len = torch.tensor([f.shape[0] for f in feat], device=device)

    # Pad sequences with edge replication (for temporal conv)
    src_padded_list = []
    for src in feat:
        left = src[0:1, :].expand(temp_pad, -1)
        right = src[-1:, :].expand(temp_pad, -1)
        src_padded = torch.cat([left, src, right], dim=0)
        src_padded_list.append(src_padded)

    # Pad to batch
    src = pad_sequence(src_padded_list, batch_first=True, padding_value=0.0).to(torch.float32)
    src = src.to(device)

    # Create mask (excludes padding)
    max_src_len = src.shape[1] - 2 * temp_pad
    src_mask = (torch.arange(max_src_len, device=device).unsqueeze(0) < src_len.unsqueeze(1)).bool()
    src_mask = src_mask.unsqueeze(1)  # (batch, 1, seq_len)

    # Run model
    with torch.no_grad():
        sign_level = model(src=src, src_length=src_len, src_mask=src_mask)

    # Reconstruct full sequence
    # Last batch contains overlap, extract only the overlap portion
    sign_crop = sign_level[-1, -a_overlap:]
    # Flatten all full windows
    sign_level = sign_level[:-1].flatten(0, 1)
    # Concatenate full windows + overlap
    sign_level = torch.cat((sign_level, sign_crop), dim=0).detach().cpu()

    # Get class predictions
    _, sign_level = torch.max(sign_level, dim=-1)

    return sign_level
