import torch

from speculators.train.data import BatchType

__all__ = ["shift_batch"]


def shift_batch(batch: BatchType):
    input_ids = batch["input_ids"]  # shape: [seq_len]
    # [x0, x1, x2, x3, x4, x5, x6, x7, x8, x9]
    hidden_states = batch["hidden_states"]  # shape: [seq_len, hidden_size]
    # [g0, g1, g2, g3, g4, g5, g6, g7, g8, g9]
    verifier_last_hidden_states = batch[
        "verifier_last_hidden_states"
    ]  # shape: [seq_len, hidden_size]
    # [y0, y1, y2, y3, y4, y5, y6, y7, y8, y9]
    loss_mask = batch["loss_mask"]  # shape: [seq_len]
    # [l0, l1, l2, l3, l4, l5, l6, l7, l8, l9]
    lengths = batch["lengths"]  # shape: [1]
    # [10]
    position_ids = batch["position_ids"]  # shape: [seq_len]
    # [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

    # Need to align (x1, g0, y1, l1)

    # Drop x0, g(-1), y0, l0, reduce seq_len by 1

    input_ids = input_ids[1:]
    hidden_states = hidden_states[:-1]
    verifier_last_hidden_states = verifier_last_hidden_states[1:]
    loss_mask = loss_mask[1:]
    # Clamp to 0 so the collate_fn fallback empty sample (seq_len=0) doesn't
    # produce a negative length.
    lengths = torch.clamp(lengths - 1, min=0)
    position_ids = position_ids[1:]  # Note: position_ids now start at 1

    return {
        "input_ids": input_ids,
        "hidden_states": hidden_states,
        "verifier_last_hidden_states": verifier_last_hidden_states,
        "loss_mask": loss_mask,
        "lengths": lengths,
        "position_ids": position_ids,
    }