"""Unit tests for reward model implementation."""

import torch
import torch.nn.functional as F
from modules.gpt import GPTConfig, GPTTrainConfig, GPT


def compute_loss_end_indices(Y: torch.Tensor, seq_len: int, semantic_pad_token: int = 4000) -> list[int]:
    """Compute where valid data ends for each sample (copied for testing)."""
    batch_size = Y.shape[0]
    loss_end_index_list = []

    for i in range(batch_size):
        y_sample = Y[i, 0, :]  # First codebook (semantic)

        # Check for BOTH -1 (marked by mask_middle_padding) AND actual pad token
        non_pad_mask = (y_sample != -1) & (y_sample != semantic_pad_token)

        if non_pad_mask.any():
            last_valid = torch.where(non_pad_mask)[0][-1].item()
            end_idx = last_valid + 2  # Y-shift: Y[j]=pad → X[j+1] padding
        else:
            end_idx = seq_len

        loss_end_index_list.append(end_idx)

    return loss_end_index_list


def extract_scalar_rewards(reward_logits, loss_start_index_list, loss_end_index_list=None):
    """Extract scalar rewards by averaging between start and end indices (copied for testing)."""
    batch_size, seq_len = reward_logits.shape
    device = reward_logits.device

    # Position indices
    positions = torch.arange(seq_len, device=device).unsqueeze(0).expand(batch_size, -1)

    # Start indices
    start_indices = (
        torch.tensor(loss_start_index_list, device=device).unsqueeze(1)
        if loss_start_index_list
        else torch.zeros(batch_size, 1, dtype=torch.long, device=device)
    )

    # End indices
    end_indices = (
        torch.tensor(loss_end_index_list, device=device).unsqueeze(1)
        if loss_end_index_list
        else torch.full((batch_size, 1), seq_len, dtype=torch.long, device=device)
    )

    # Mask: start <= position < end
    mask = (positions >= start_indices) & (positions < end_indices)

    # Average
    masked_rewards = reward_logits * mask
    valid_counts = mask.sum(dim=1, keepdim=True).clamp(min=1)
    scalar_rewards = masked_rewards.sum(dim=1) / valid_counts.squeeze(1)

    return scalar_rewards


def token_level_reward_loss(
    reward_logits_chosen,
    reward_logits_rejected,
    loss_start_index_list_chosen,
    loss_start_index_list_rejected,
    loss_end_index_list_chosen,
    loss_end_index_list_rejected,
    beta=1.0,
):
    """Token-level Bradley-Terry loss (copied for testing)."""
    batch_size, seq_len = reward_logits_chosen.shape
    device = reward_logits_chosen.device

    # Create masks inline
    positions = torch.arange(seq_len, device=device).unsqueeze(0).expand(batch_size, -1)

    start_c = torch.tensor(loss_start_index_list_chosen, device=device).unsqueeze(1)
    end_c = torch.tensor(loss_end_index_list_chosen, device=device).unsqueeze(1)
    mask_chosen = (positions >= start_c) & (positions < end_c)

    start_r = torch.tensor(loss_start_index_list_rejected, device=device).unsqueeze(1)
    end_r = torch.tensor(loss_end_index_list_rejected, device=device).unsqueeze(1)
    mask_rejected = (positions >= start_r) & (positions < end_r)

    mask = mask_chosen & mask_rejected

    token_diff = reward_logits_chosen - reward_logits_rejected
    token_losses = -F.logsigmoid(beta * token_diff)

    masked_losses = token_losses * mask
    loss = masked_losses.sum() / mask.sum().clamp(min=1)

    token_correct = (reward_logits_chosen > reward_logits_rejected) & mask
    accuracy = token_correct.sum().float() / mask.sum().float().clamp(min=1)

    return loss, accuracy


def reward_model_loss(reward_chosen, reward_rejected, label_smoothing=0.0):
    """Bradley-Terry loss with label smoothing (copied for testing)."""
    losses = (
        -F.logsigmoid(reward_chosen - reward_rejected) * (1 - label_smoothing)
        - F.logsigmoid(reward_rejected - reward_chosen) * label_smoothing
    )
    loss = losses.mean()
    accuracy = (reward_chosen > reward_rejected).float().mean()
    return loss, accuracy


def test_reward_head_initialization():
    """Test that reward head is properly initialized when use_reward_head=True."""
    config = GPTConfig(
        n_layer=2,
        n_head=4,
        d_head=64,
        block_size=256,
        t_text=128,
        t_audio=128,
        use_reward_head=True,
        output_paradigm="gpt",
        output_distribution="semantic",
    )
    train_config = GPTTrainConfig()
    model = GPT(config, train_config)

    # Check reward head exists
    assert "reward_head" in model.output_modules
    assert hasattr(model.output_modules["reward_head"], "reward_proj")
    assert hasattr(model.output_modules["reward_head"], "ln")

    # Verify reward head is initialized with small weights (std=0.01)
    reward_head = model.output_modules["reward_head"]
    weight_norm = reward_head.reward_proj.weight.norm().item()
    # Expected norm: sqrt(n_embd) * std with std=0.01
    # For small test model, this will be small; for full model ~0.64
    # Just verify it's not too large (< 2.0) and not zero
    assert 0.05 < weight_norm < 2.0, f"Reward head weight norm should be small, got {weight_norm}"

    # Test that it outputs near-zero rewards with small variance
    batch_size = 2
    seq_len = 10
    x = torch.randn(batch_size, seq_len, config.n_embd)
    with torch.no_grad():
        rewards = reward_head(x)
    assert rewards.shape == (batch_size, seq_len, 1)
    reward_mean = rewards.mean().item()
    reward_std = rewards.std().item()
    # Mean should be close to 0, std should be small
    assert abs(reward_mean) < 1.0, f"Reward mean should be near 0, got {reward_mean:.4f}"
    assert reward_std < 2.0, f"Reward std should be small, got {reward_std:.4f}"


def test_reward_head_forward():
    """Test that reward head returns correct shape for sequence output."""
    # Test directly on reward head module (skip full model forward which needs GPU)
    config = GPTConfig(
        n_layer=2,
        n_head=4,
        d_head=64,
        block_size=256,
        t_text=128,
        t_audio=128,
        use_reward_head=True,
        output_paradigm="gpt",
        output_distribution="semantic",
    )
    train_config = GPTTrainConfig()
    model = GPT(config, train_config)
    model.eval()

    # Test reward head on sequence of hidden states
    batch_size = 4
    seq_len = 100
    n_embd = config.n_embd
    hidden_states = torch.randn(batch_size, seq_len, n_embd)

    # Get rewards from reward head (should be one reward per token)
    with torch.no_grad():
        rewards = model.output_modules["reward_head"](hidden_states)

    # Check shape: should be (batch, seq_len, 1)
    assert rewards.shape == (
        batch_size,
        seq_len,
        1,
    ), f"Expected shape ({batch_size}, {seq_len}, 1), got {rewards.shape}"

    # Squeeze and check again
    rewards = rewards.squeeze(-1)
    assert rewards.shape == (
        batch_size,
        seq_len,
    ), f"Expected shape ({batch_size}, {seq_len}), got {rewards.shape}"


def test_bradley_terry_loss():
    """Test Bradley-Terry loss computation."""
    # Create dummy rewards
    reward_chosen = torch.tensor([1.0, 2.0, 3.0, 4.0])
    reward_rejected = torch.tensor([0.5, 1.5, 2.5, 3.5])

    loss, accuracy = reward_model_loss(reward_chosen, reward_rejected)

    # Loss should be positive
    assert loss.item() >= 0

    # Accuracy should be 1.0 (all chosen > rejected)
    assert accuracy.item() == 1.0

    # Test case where some are wrong
    reward_chosen_mixed = torch.tensor([1.0, 0.5, 3.0, 2.0])
    reward_rejected_mixed = torch.tensor([0.5, 1.0, 2.5, 3.0])

    loss_mixed, accuracy_mixed = reward_model_loss(reward_chosen_mixed, reward_rejected_mixed)

    # Accuracy should be 0.5 (2 out of 4 correct)
    assert accuracy_mixed.item() == 0.5


def test_reward_head_gradient_flow():
    """Test that gradients flow through reward head."""
    config = GPTConfig(
        n_layer=2,
        n_head=4,
        d_head=64,
        block_size=256,
        t_text=128,
        t_audio=128,
        use_reward_head=True,
        output_paradigm="gpt",
        output_distribution="semantic",
    )
    train_config = GPTTrainConfig()
    model = GPT(config, train_config)
    model.train()

    # Test gradient flow through reward head directly
    batch_size = 2
    n_embd = config.n_embd
    hidden_states = torch.randn(batch_size, n_embd, requires_grad=True)

    # Forward pass through reward head
    rewards = model.output_modules["reward_head"](hidden_states).squeeze(-1)

    # Compute dummy loss
    loss = rewards.mean()
    loss.backward()

    # Check that reward head has gradients
    assert model.output_modules["reward_head"].reward_proj.weight.grad is not None
    assert model.output_modules["reward_head"].ln.weight.grad is not None
    # Also check input has gradients
    assert hidden_states.grad is not None


def test_freeze_base_model():
    """Test freezing base model while training reward head."""
    config = GPTConfig(
        n_layer=2,
        n_head=4,
        d_head=64,
        block_size=256,
        t_text=128,
        t_audio=128,
        use_reward_head=True,
        output_paradigm="gpt",
        output_distribution="semantic",
    )
    train_config = GPTTrainConfig()
    model = GPT(config, train_config)

    # Freeze everything except reward head
    for name, param in model.named_parameters():
        if "reward_head" not in name:
            param.requires_grad = False

    # Count trainable params
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total = sum(p.numel() for p in model.parameters())

    # Reward head should be much smaller than full model
    assert trainable < total * 0.01, "Reward head should be <1% of total params"

    # Check specific params
    assert model.output_modules["reward_head"].reward_proj.weight.requires_grad
    # Check that transformer layers are frozen (use c_proj which exists in all models)
    assert not model.transformer.h[0].attn.c_proj.weight.requires_grad


def test_paired_preference_data():
    """Test that paired preference data structure is correct."""
    # Simulate paired batch: [neg0, pos0, neg1, pos1, ...]
    batch_size = 8
    rewards = torch.tensor([0.5, 1.5, 0.3, 1.2, 0.7, 1.8, 0.4, 1.3])

    # Split into chosen and rejected
    reward_rejected = rewards[::2]  # Even indices: [0.5, 0.3, 0.7, 0.4]
    reward_chosen = rewards[1::2]  # Odd indices: [1.5, 1.2, 1.8, 1.3]

    assert reward_rejected.shape == (batch_size // 2,)
    assert reward_chosen.shape == (batch_size // 2,)
    assert torch.all(reward_chosen > reward_rejected), "Chosen should be > rejected in test data"


def test_reward_extraction_from_generated_portion():
    """Test that reward is extracted from generated portion only (after loss_start_index)."""
    config = GPTConfig(
        n_layer=2,
        n_head=4,
        d_head=64,
        block_size=256,
        t_text=128,
        t_audio=128,
        use_reward_head=True,
        output_paradigm="gpt",
        output_distribution="semantic",
    )
    train_config = GPTTrainConfig()
    model = GPT(config, train_config)
    model.eval()

    # Test the extraction logic directly without running full model (which needs GPU)
    batch_size = 4
    seq_len = 256
    n_embd = config.n_embd
    x = torch.randn(batch_size, seq_len, n_embd)

    # Simulate different generation start positions
    loss_start_index_list = [50, 100, 75, 120]  # Different for each sample

    # Create dummy reward logits (one per token) from reward head
    reward_logits = torch.randn(batch_size, seq_len)

    # Create Y with padding tokens (-1)
    n_codebooks = config.semantic_n_codebooks + config.coarse_n_codebooks
    Y = torch.randint(0, 1000, (batch_size, n_codebooks, seq_len - 1))
    # Add padding at different positions
    Y[0, :, 180:] = -1  # Sample 0: padding after position 180
    Y[1, :, 200:] = -1  # Sample 1: padding after position 200
    Y[2, :, 150:] = -1  # Sample 2: padding after position 150
    Y[3, :, 220:] = -1  # Sample 3: padding after position 220

    # Extract scalar rewards (specify end indices based on Y padding)
    # For tests with no padding, use None
    scalar_rewards = extract_scalar_rewards(reward_logits, loss_start_index_list, None)

    # Check shape
    assert scalar_rewards.shape == (
        batch_size,
    ), f"Expected shape ({batch_size},), got {scalar_rewards.shape}"

    # Verify each is a scalar
    for i in range(batch_size):
        assert scalar_rewards[i].dim() == 0, f"Sample {i} should be scalar"

    # Verify that different sequences produce different rewards (usually true for random data)
    # This is a basic sanity check that extraction works


def test_extract_scalar_rewards_correctness():
    """Comprehensive test to verify extract_scalar_rewards masking is correct."""
    batch_size = 3
    seq_len = 20
    n_codebooks = 1

    # Create simple reward logits where reward[i, j] = j (position-based for easy verification)
    reward_logits = torch.arange(seq_len, dtype=torch.float32).unsqueeze(0).expand(batch_size, -1)
    # reward_logits[i, j] = j for all samples

    # Create Y with known padding positions
    Y = torch.zeros(batch_size, n_codebooks, seq_len - 1, dtype=torch.long)

    # Sample 0: No padding, start at position 5
    # Should average positions [5, 6, 7, ..., 19]
    # Expected: sum(5..19) / 15 = (5+19)*15/2 / 15 = 12.0
    Y[0, :, :] = 1  # All valid (no -1)

    # Sample 1: Padding starts at Y[14], start at position 8
    # Y[14]=-1 means the target at position 14 is padding
    # Y[j] predicts from X[j] to X[j+1], so Y[14] corresponds to predicting X[15]
    # If Y[14]=-1, we can't use X[15] (it's in padding region)
    # Valid X positions: 0-14 (Y[0..13] are valid)
    # Should average positions [8, 9, 10, ..., 14]
    # Expected: sum(8..14) / 7 = 11.0
    Y[1, :, :14] = 1  # Valid up to Y[13]
    Y[1, :, 14:] = -1  # Y[14]=-1 → X[15] onwards is padding

    # Sample 2: Start at 10, padding at Y[11]
    # Y[11]=-1 excludes X[12] onwards
    # Valid X positions: 0-11 (Y[0..10] are valid)
    # Should average positions [10, 11]
    # Expected: sum(10..11) / 2 = 10.5
    Y[2, :, :11] = 1
    Y[2, :, 11:] = -1

    loss_start_index_list = [5, 8, 10]

    # Compute end indices from Y padding
    # Sample 0: No padding → end at 20
    # Sample 1: Y[14]=-1 → end at 15 (14+2 from Y-shift, but actually 15 is where it ends)
    # Sample 2: Y[11]=-1 → end at 12
    loss_end_index_list = [20, 15, 12]

    # Extract scalar rewards
    scalar_rewards = extract_scalar_rewards(reward_logits, loss_start_index_list, loss_end_index_list)

    # Verify shape
    assert scalar_rewards.shape == (batch_size,), f"Shape mismatch: {scalar_rewards.shape}"

    # Verify each reward is correct
    # Sample 0: mean of positions 5-19 (15 values)
    expected_0 = sum(range(5, 20)) / 15.0
    actual_0 = scalar_rewards[0].item()
    assert abs(actual_0 - expected_0) < 0.01, f"Sample 0: expected {expected_0:.2f}, got {actual_0:.2f}"

    # Sample 1: mean of positions 8-14 (7 values, Y[14]=-1 excludes X[15])
    expected_1 = sum(range(8, 15)) / 7.0
    actual_1 = scalar_rewards[1].item()
    assert abs(actual_1 - expected_1) < 0.01, f"Sample 1: expected {expected_1:.2f}, got {actual_1:.2f}"

    # Sample 2: mean of positions 10-11 (2 values, Y[11]=-1 excludes X[12])
    expected_2 = sum(range(10, 12)) / 2.0
    actual_2 = scalar_rewards[2].item()
    assert abs(actual_2 - expected_2) < 0.01, f"Sample 2: expected {expected_2:.2f}, got {actual_2:.2f}"

    print(f"  Sample 0: {actual_0:.2f} == {expected_0:.2f} ✓")
    print(f"  Sample 1: {actual_1:.2f} == {expected_1:.2f} ✓")
    print(f"  Sample 2: {actual_2:.2f} == {expected_2:.2f} ✓")


def test_extract_scalar_rewards_edge_cases():
    """Test edge cases for extract_scalar_rewards."""
    batch_size = 4
    seq_len = 100
    n_codebooks = 1

    reward_logits = torch.randn(batch_size, seq_len)
    Y = torch.randint(0, 1000, (batch_size, n_codebooks, seq_len - 1))

    # Edge case 1: loss_start_index at 0 (full sequence)
    loss_start_index_list = [0, 0, 0, 0]
    Y[:, :, :] = 1  # No padding

    rewards = extract_scalar_rewards(reward_logits, loss_start_index_list, None)
    assert rewards.shape == (batch_size,)

    # Each should be mean of full sequence
    for i in range(batch_size):
        expected = reward_logits[i].mean()
        assert abs(rewards[i].item() - expected.item()) < 1e-5, f"Full sequence mean mismatch"

    # Edge case 2: loss_start_index very late (only last few tokens)
    loss_start_index_list = [95, 96, 97, 98]
    Y[:, :, :] = 1  # No padding

    rewards = extract_scalar_rewards(reward_logits, loss_start_index_list, None)
    assert rewards.shape == (batch_size,)

    # Each should be mean of last few tokens
    for i in range(batch_size):
        start = loss_start_index_list[i]
        expected = reward_logits[i, start:].mean()
        assert abs(rewards[i].item() - expected.item()) < 1e-5, f"Late start mean mismatch"

    # Edge case 3: Heavy padding (only first few tokens valid)
    Y[:, :, 5:] = -1  # Heavy padding - Y[5]=-1 means X[6] onwards is padding
    loss_start_index_list = [0, 0, 0, 0]
    # Y[5]=-1 → last_valid=4 → end_idx=4+2=6
    loss_end_index_list = [6, 6, 6, 6]

    rewards = extract_scalar_rewards(reward_logits, loss_start_index_list, loss_end_index_list)
    assert rewards.shape == (batch_size,)
    # Should average over only positions 0-5 (Y[5]=-1 → X[6] is padding, so valid up to X[5])
    for i in range(batch_size):
        expected = reward_logits[i, :6].mean()  # Positions 0-5
        assert abs(rewards[i].item() - expected.item()) < 1e-5, f"Heavy padding mean mismatch"

    # Edge case 4: None for loss_start_index_list (should default to 0)
    Y[:, :, :] = 1  # No padding
    rewards = extract_scalar_rewards(reward_logits, None, None)
    assert rewards.shape == (batch_size,)
    for i in range(batch_size):
        expected = reward_logits[i].mean()
        assert abs(rewards[i].item() - expected.item()) < 1e-5


def test_token_level_reward_loss():
    """Test token-level Bradley-Terry loss computation."""
    batch_size = 2
    seq_len = 10
    n_codebooks = 1

    # Create token-level rewards where chosen is consistently better
    # Chosen: [0.5, 0.6, 0.7, ..., 1.4]
    # Rejected: [0.4, 0.5, 0.6, ..., 1.3]
    reward_logits_chosen = (
        torch.arange(5, 15, dtype=torch.float32).unsqueeze(0).expand(batch_size, -1) / 10.0
    )
    reward_logits_rejected = (
        torch.arange(4, 14, dtype=torch.float32).unsqueeze(0).expand(batch_size, -1) / 10.0
    )

    # Create Y with no padding
    Y_chosen = torch.ones(batch_size, n_codebooks, seq_len - 1, dtype=torch.long)
    Y_rejected = torch.ones(batch_size, n_codebooks, seq_len - 1, dtype=torch.long)

    # Start from position 2
    loss_start_index_list_chosen = [2, 2]
    loss_start_index_list_rejected = [2, 2]

    # Compute token-level loss
    # Compute end indices (no padding, so use seq_len)
    seq_len = reward_logits_chosen.shape[1]
    loss_end_list = [seq_len, seq_len]

    loss, accuracy = token_level_reward_loss(
        reward_logits_chosen,
        reward_logits_rejected,
        loss_start_index_list_chosen,
        loss_start_index_list_rejected,
        loss_end_list,  # chosen end
        loss_end_list,  # rejected end
        beta=1.0,
    )

    # Loss should be positive
    assert loss.item() > 0

    # Accuracy should be 1.0 (all tokens have chosen > rejected)
    assert accuracy.item() == 1.0, f"Expected accuracy 1.0, got {accuracy.item()}"

    # Test with some positions wrong
    reward_logits_mixed_chosen = reward_logits_chosen.clone()
    reward_logits_mixed_chosen[:, 5] = 0.3  # Make position 5 worse than rejected

    loss_mixed, acc_mixed = token_level_reward_loss(
        reward_logits_mixed_chosen,
        reward_logits_rejected,
        loss_start_index_list_chosen,
        loss_start_index_list_rejected,
        loss_end_list,  # chosen end
        loss_end_list,  # rejected end
        beta=1.0,
    )

    # Accuracy should be less than 1.0 now
    # Positions 2-9 = 8 tokens, position 5 is wrong, so 7/8 = 0.875
    assert 0.8 < acc_mixed.item() < 0.9, f"Expected accuracy ~0.875, got {acc_mixed.item()}"


def test_token_level_with_padding():
    """Test token-level loss correctly handles padding."""
    batch_size = 2
    seq_len = 20
    n_codebooks = 1

    # Position-based rewards for easy verification
    reward_logits_chosen = torch.arange(seq_len, dtype=torch.float32).unsqueeze(0).expand(batch_size, -1)
    reward_logits_rejected = (
        (torch.arange(seq_len, dtype=torch.float32) - 0.5).unsqueeze(0).expand(batch_size, -1)
    )

    # Create Y with padding
    Y_chosen = torch.ones(batch_size, n_codebooks, seq_len - 1, dtype=torch.long)
    Y_rejected = torch.ones(batch_size, n_codebooks, seq_len - 1, dtype=torch.long)

    # Add padding at position 10 for both
    Y_chosen[:, :, 10:] = -1
    Y_rejected[:, :, 10:] = -1

    # Start from position 5
    loss_start_list = [5, 5]

    # Compute end indices from Y (padding at position 10)
    seq_len = reward_logits_chosen.shape[1]
    # Y[10]=-1 means X[11] padding, so end at 12 (exclusive)
    loss_end_list = [12, 12]

    loss, accuracy = token_level_reward_loss(
        reward_logits_chosen,
        reward_logits_rejected,
        loss_start_list,
        loss_start_list,
        loss_end_list,
        loss_end_list,
        beta=1.0,
    )

    # Should only compute loss on positions 5-11 (7 positions total)
    # Y[10]=-1 means X[11] is last valid, but since Y padding starts at 10,
    # valid positions are 5-10 (6 positions)
    # All should have chosen > rejected
    assert accuracy.item() == 1.0
    assert loss.item() > 0


def test_label_smoothing_loss():
    """Test Bradley-Terry loss with label smoothing."""
    reward_chosen = torch.tensor([1.0, 2.0, 3.0, 4.0])
    reward_rejected = torch.tensor([0.5, 1.5, 2.5, 3.5])

    # Loss without smoothing
    loss_no_smooth, acc_no_smooth = reward_model_loss(
        reward_chosen, reward_rejected, label_smoothing=0.0
    )

    # Loss with smoothing
    loss_smooth, acc_smooth = reward_model_loss(reward_chosen, reward_rejected, label_smoothing=0.1)

    # Smoothed loss should be higher (less confident)
    assert loss_smooth.item() > loss_no_smooth.item(), "Smoothed loss should be higher"

    # Accuracy should be the same (smoothing doesn't change predictions)
    assert acc_smooth.item() == acc_no_smooth.item() == 1.0


def test_compute_loss_end_indices():
    """Test compute_loss_end_indices with various padding scenarios."""
    semantic_pad_token = 4000
    seq_len = 100
    batch_size = 6
    n_codebooks = 1

    # Create Y with shape (batch, n_codebooks, seq_len-1)
    Y = torch.zeros(batch_size, n_codebooks, seq_len - 1, dtype=torch.long)

    # Case 1: No padding - full sequence
    Y[0, :, :] = 1  # All valid data

    # Case 2: Padding with 1 token at end (only marked with semantic_pad_token=4000)
    # Last valid data at Y[97], then Y[98]=4000 (1 pad token)
    # This simulates mask_middle_padding NOT marking it as -1 (only 1 pad token)
    Y[1, :, :98] = 1
    Y[1, :, 98:] = semantic_pad_token

    # Case 3: Padding with 2 tokens at end (only marked with semantic_pad_token=4000)
    # Last valid data at Y[96], then Y[97:99]=4000 (2 pad tokens)
    Y[2, :, :97] = 1
    Y[2, :, 97:] = semantic_pad_token

    # Case 4: Padding with 3+ tokens at end (marked with -1 by mask_middle_padding)
    # Last valid data at Y[94], then Y[95:99]=4000, but middle ones marked as -1
    Y[3, :, :95] = 1
    Y[3, :, 95] = semantic_pad_token  # First pad
    Y[3, :, 96:98] = -1  # Middle pads marked as -1
    Y[3, :, 98] = semantic_pad_token  # Last pad

    # Case 5: Different pair - short sequence (like chosen vs rejected)
    # Last valid at Y[49], rest is padding
    Y[4, :, :50] = 1
    Y[4, :, 50:] = semantic_pad_token

    # Case 6: Very short sequence with only -1 padding
    Y[5, :, :20] = 1
    Y[5, :, 20:] = -1

    # Compute end indices
    end_indices = compute_loss_end_indices(Y, seq_len, semantic_pad_token)

    # Verify results
    # Case 1: No padding → end at seq_len
    assert end_indices[0] == seq_len, f"Case 1: Expected {seq_len}, got {end_indices[0]}"

    # Case 2: Last valid Y[97] → end_idx = 97 + 2 = 99
    # Y[98]=4000 means X[99] is padding
    assert end_indices[1] == 99, f"Case 2 (1 pad): Expected 99, got {end_indices[1]}"

    # Case 3: Last valid Y[96] → end_idx = 96 + 2 = 98
    # Y[97]=4000 means X[98] is padding
    assert end_indices[2] == 98, f"Case 3 (2 pads): Expected 98, got {end_indices[2]}"

    # Case 4: Last valid Y[94] → end_idx = 94 + 2 = 96
    # Y[95]=4000 or -1 means X[96] is padding
    assert end_indices[3] == 96, f"Case 4 (3+ pads with -1): Expected 96, got {end_indices[3]}"

    # Case 5: Last valid Y[49] → end_idx = 49 + 2 = 51
    assert end_indices[4] == 51, f"Case 5 (short seq): Expected 51, got {end_indices[4]}"

    # Case 6: Last valid Y[19] → end_idx = 19 + 2 = 21
    assert end_indices[5] == 21, f"Case 6 (only -1): Expected 21, got {end_indices[5]}"

    print(f"  All cases validated:")
    print(f"    No padding: {end_indices[0]}")
    print(f"    1 pad token: {end_indices[1]}")
    print(f"    2 pad tokens: {end_indices[2]}")
    print(f"    3+ pads (-1): {end_indices[3]}")
    print(f"    Short sequence: {end_indices[4]}")
    print(f"    Only -1 padding: {end_indices[5]}")


def test_loss_indices_for_preference_pairs():
    """Test that chosen/rejected pairs can have different loss indices."""
    semantic_pad_token = 4000
    seq_len = 100
    batch_size = 4  # 2 pairs
    n_codebooks = 1

    # Create Y for 2 preference pairs (interleaved: [rej_0, cho_0, rej_1, cho_1])
    Y = torch.zeros(batch_size, n_codebooks, seq_len - 1, dtype=torch.long)
    loss_start_index_list = [10, 10, 20, 20]  # Same start for each pair

    # Pair 0: Rejected (short) vs Chosen (long)
    Y[0, :, :60] = 1  # Rejected: valid to Y[59]
    Y[0, :, 60:] = semantic_pad_token

    Y[1, :, :80] = 1  # Chosen: valid to Y[79]
    Y[1, :, 80:] = semantic_pad_token

    # Pair 1: Rejected (long) vs Chosen (short)
    Y[2, :, :90] = 1  # Rejected: valid to Y[89]
    Y[2, :, 90:] = -1

    Y[3, :, :50] = 1  # Chosen: valid to Y[49]
    Y[3, :, 50:] = semantic_pad_token

    # Compute end indices
    end_indices = compute_loss_end_indices(Y, seq_len, semantic_pad_token)

    # Verify pairs have DIFFERENT lengths
    # Pair 0: Rejected valid to Y[59] → end = 59+2 = 61
    rej_0_end = end_indices[0]
    cho_0_end = end_indices[1]  # Chosen valid to Y[79] → end = 79+2 = 81
    assert rej_0_end == 61, f"Pair 0 rejected: Expected 61, got {rej_0_end}"
    assert cho_0_end == 81, f"Pair 0 chosen: Expected 81, got {cho_0_end}"
    assert rej_0_end != cho_0_end, "Pair 0: rejected and chosen should have DIFFERENT end indices!"

    # Pair 1: Rejected valid to Y[89] → end = 89+2 = 91
    rej_1_end = end_indices[2]
    cho_1_end = end_indices[3]  # Chosen valid to Y[49] → end = 49+2 = 51
    assert rej_1_end == 91, f"Pair 1 rejected: Expected 91, got {rej_1_end}"
    assert cho_1_end == 51, f"Pair 1 chosen: Expected 51, got {cho_1_end}"
    assert rej_1_end != cho_1_end, "Pair 1: rejected and chosen should have DIFFERENT end indices!"

    # Verify they use different amounts of data
    rej_0_len = rej_0_end - loss_start_index_list[0]
    cho_0_len = cho_0_end - loss_start_index_list[1]
    rej_1_len = rej_1_end - loss_start_index_list[2]
    cho_1_len = cho_1_end - loss_start_index_list[3]

    print(f"  Pair 0: rejected_len={rej_0_len}, chosen_len={cho_0_len} ✓ DIFFERENT")
    print(f"  Pair 1: rejected_len={rej_1_len}, chosen_len={cho_1_len} ✓ DIFFERENT")


if __name__ == "__main__":
    # Run tests
    print("Testing reward head initialization...")
    test_reward_head_initialization()
    print("✓ Passed")

    print("Testing reward head forward pass...")
    test_reward_head_forward()
    print("✓ Passed")

    print("Testing Bradley-Terry loss...")
    test_bradley_terry_loss()
    print("✓ Passed")

    print("Testing gradient flow...")
    test_reward_head_gradient_flow()
    print("✓ Passed")

    print("Testing freeze base model...")
    test_freeze_base_model()
    print("✓ Passed")

    print("Testing paired preference data...")
    test_paired_preference_data()
    print("✓ Passed")

    print("Testing reward extraction from generated portion...")
    test_reward_extraction_from_generated_portion()
    print("✓ Passed")

    print("Testing extract_scalar_rewards correctness...")
    test_extract_scalar_rewards_correctness()
    print("✓ Passed")

    print("Testing extract_scalar_rewards edge cases...")
    test_extract_scalar_rewards_edge_cases()
    print("✓ Passed")

    print("Testing label smoothing loss...")
    test_label_smoothing_loss()
    print("✓ Passed")

    print("Testing token-level reward loss...")
    test_token_level_reward_loss()
    print("✓ Passed")

    print("Testing token-level loss with padding...")
    test_token_level_with_padding()
    print("✓ Passed")

    print("Testing compute_loss_end_indices...")
    test_compute_loss_end_indices()
    print("✓ Passed")

    print("Testing loss indices for preference pairs...")
    test_loss_indices_for_preference_pairs()
    print("✓ Passed")

    print("\n✅ All tests passed!")