"""Sliding-window decode helpers used by the Gradio demo."""

import torch
from torch_geometric.data import Batch


def decode_to_hatD_with_sliding_window(
    model,
    latent_codes,
    src_graphs_list,
    total_frames,
    window_size,
    overlap,
    device,
    model_type,
):
    """Decode a long latent/code sequence in overlapping windows.

    Overlapping regions keep the earlier window prediction and drop the later
    window overlap. This matches the preprocessing convention used by the demo
    checkpoints.
    """
    if model_type not in {"vae", "rvq"}:
        raise ValueError(f"model_type must be 'vae' or 'rvq', got: {model_type}")

    stride = window_size - overlap
    if stride <= 0:
        raise ValueError("window_size must be greater than overlap")

    if len(src_graphs_list) != total_frames:
        raise ValueError(
            f"src_graphs_list length ({len(src_graphs_list)}) != total_frames ({total_frames})"
        )
    if latent_codes.shape[0] != total_frames:
        raise ValueError(
            f"latent_codes length ({latent_codes.shape[0]}) != total_frames ({total_frames})"
        )

    if total_frames <= window_size:
        num_windows = 1
    else:
        num_windows = (total_frames - window_size + stride - 1) // stride + 1

    num_nodes_per_frame = src_graphs_list[0].skel_x.shape[0]
    hatD_parts = []

    for window_idx in range(num_windows):
        start_frame = window_idx * stride
        if start_frame >= total_frames:
            break

        end_frame = min(start_frame + window_size, total_frames)
        window_frames = end_frame - start_frame
        latent_window = latent_codes[start_frame:end_frame].to(device)

        src_window_graphs = src_graphs_list[start_frame:end_frame]
        src_window = Batch.from_data_list(src_window_graphs).to(device)

        with torch.no_grad():
            if model_type == "vae":
                hatD_win = model.decode(latent_window, src_window, window_frames)
            else:
                hatD_win, _ = model.decode_from_codes(latent_window, src_window, window_frames)

        hatD_dim = hatD_win.shape[1]
        hatD_win_reshaped = hatD_win.view(window_frames, num_nodes_per_frame, hatD_dim)

        if window_idx == 0:
            keep_start_idx = 0
        else:
            keep_start_idx = overlap
        hatD_parts.append(hatD_win_reshaped[keep_start_idx:window_frames].cpu())

        del src_window, hatD_win, hatD_win_reshaped

    hatD_full_3d = torch.cat(hatD_parts, dim=0)
    actual_frames = hatD_full_3d.shape[0]
    hatD_full = hatD_full_3d.view(actual_frames * num_nodes_per_frame, -1).to(device)
    src_batch_full = Batch.from_data_list(src_graphs_list).to(device)

    return hatD_full, src_batch_full, actual_frames, num_nodes_per_frame