"""
Test script for PAT loss functions
"""

import torch
from pat._internal.losses import (
    VQLoss,
    RVQLoss,
    ReconstructionLoss,
    CodebookDiversityLoss,
    ChunkContinuityLoss,
    ActionCodecLoss,
    OATLoss,
    load_info,
)


def test_vq_loss():
    """Test VQ Loss"""
    print("=== Testing VQLoss ===")
    vq_loss = VQLoss(commitment_weight=0.25, codebook_weight=1.0)
    
    z_e = torch.randn(4, 16, 512, requires_grad=True)
    z_q = torch.randn(4, 16, 512, requires_grad=True)
    
    loss_dict = vq_loss(z_e, z_q)
    
    print(f"  Total loss: {loss_dict['total'].item():.6f}")
    print(f"  Commitment loss: {loss_dict['commitment'].item():.6f}")
    print(f"  Codebook loss: {loss_dict['codebook'].item():.6f}")
    
    # Test backward pass
    loss_dict['total'].backward()
    assert z_e.grad is not None, "Gradient should flow to z_e"
    assert z_q.grad is not None, "Gradient should flow to z_q"
    print("  ✓ Backward pass successful\n")


def test_rvq_loss():
    """Test RVQ Loss"""
    print("=== Testing RVQLoss ===")
    rvq_loss = RVQLoss(n_codebooks=4, commitment_weight=0.25, quantizer_dropout=0.25)
    
    batch_size = 4
    z_e_list = [torch.randn(batch_size, 16, 512, requires_grad=True) for _ in range(4)]
    z_q_list = [torch.randn(batch_size, 16, 512, requires_grad=True) for _ in range(4)]
    
    # Training mode with dropout
    rvq_loss.train()
    loss_dict = rvq_loss(z_e_list, z_q_list)
    
    print(f"  Total loss: {loss_dict['total'].item():.6f}")
    print(f"  Commitment loss: {loss_dict['commitment'].item():.6f}")
    print(f"  Codebook loss: {loss_dict['codebook'].item():.6f}")
    print(f"  Number of per-quantizer losses: {len(loss_dict['per_quantizer'])}")
    
    # Test backward pass
    loss_dict['total'].backward()
    for i, z_e in enumerate(z_e_list):
        assert z_e.grad is not None, f"Gradient should flow to z_e[{i}]"
    print("  ✓ Backward pass successful\n")


def test_reconstruction_loss():
    """Test Reconstruction Loss with dual-arm data format"""
    print("=== Testing ReconstructionLoss ===")
    
    # Data format: [batch, seq, 20]
    # Left: xyz(0:3) + 6D(3:9) + gripper(9)
    # Right: xyz(10:13) + 6D(13:19) + gripper(19)
    batch_size = 4
    seq_len = 16
    
    # Test MSE loss
    recon_loss_mse = ReconstructionLoss(loss_type='mse')
    x = torch.randn(batch_size, seq_len, 20, requires_grad=True)
    x_recon = torch.randn(batch_size, seq_len, 20, requires_grad=True)
    
    loss = recon_loss_mse(x, x_recon)
    print(f"  MSE loss: {loss.item():.6f}")
    loss.backward()
    assert x.grad is not None
    print("  ✓ MSE backward pass successful")
    
    # Test weighted loss with gripper weight
    x.grad = None
    x_recon.grad = None
    recon_loss_weighted = ReconstructionLoss(
        loss_type='weighted',
        position_weight=1.0,
        rotation_weight=1.0,
        gripper_weight=0.5,  # Less emphasis on gripper
    )
    loss_weighted = recon_loss_weighted(x, x_recon)
    print(f"  Weighted loss (gripper_weight=0.5): {loss_weighted.item():.6f}")
    loss_weighted.backward()
    print("  ✓ Weighted loss backward pass successful")
    
    # Test with mask
    x.grad = None
    x_recon.grad = None
    mask = torch.ones(batch_size, seq_len, dtype=torch.bool)
    mask[2:, 8:] = False  # Mask out some timesteps
    
    loss_masked = recon_loss_mse(x, x_recon, mask=mask)
    print(f"  Masked MSE loss: {loss_masked.item():.6f}")
    print("  ✓ Masked loss successful\n")


def test_geodesic_loss():
    """Test Geodesic Loss for dual-arm rotations"""
    print("=== Testing GeodesicLoss ===")
    
    recon_loss_geo = ReconstructionLoss(
        loss_type='geodesic',
        position_weight=1.0,
        rotation_weight=1.0,
        gripper_weight=1.0,
    )
    
    # Dual-arm data format: [batch, seq, 20]
    # Left: xyz(3) + 6D(6) + gripper(1) = 10
    # Right: xyz(3) + 6D(6) + gripper(1) = 10
    batch_size = 4
    seq_len = 16
    
    # Create leaf tensors directly (don't multiply by 0.1 as it creates non-leaf tensors)
    x = torch.randn(batch_size, seq_len, 20, requires_grad=True)
    x_recon = torch.randn(batch_size, seq_len, 20, requires_grad=True)
    
    loss = recon_loss_geo(x, x_recon)
    print(f"  Geodesic loss (combined): {loss.item():.6f}")
    loss.backward()
    
    # Check gradients flow
    if x.grad is not None:
        print("  ✓ Gradient flows to x")
    else:
        print("  ✗ Warning: x.grad is None (may be due to non-leaf tensor)")
    
    if x_recon.grad is not None:
        print("  ✓ Gradient flows to x_recon")
    else:
        print("  ✗ Warning: x_recon.grad is None (may be due to non-leaf tensor)")
    
    print("  ✓ Geodesic backward pass successful\n")


def test_codebook_diversity_loss():
    """Test Codebook Diversity Loss"""
    print("=== Testing CodebookDiversityLoss ===")
    
    diversity_loss = CodebookDiversityLoss(weight=0.1)
    
    # Uniform usage (should have low loss)
    codes_uniform = torch.arange(0, 512).unsqueeze(0).repeat(4, 1)
    loss_uniform = diversity_loss(codes_uniform, codebook_size=512)
    print(f"  Uniform usage loss: {loss_uniform.item():.6f}")
    
    # Skewed usage (should have higher loss)
    codes_skewed = torch.zeros(4, 512, dtype=torch.long)
    codes_skewed[:, :10] = torch.randint(0, 10, (4, 10))  # Only use first 10 codes
    loss_skewed = diversity_loss(codes_skewed, codebook_size=512)
    print(f"  Skewed usage loss: {loss_skewed.item():.6f}")
    
    assert loss_skewed > loss_uniform, "Skewed usage should have higher loss"
    print("  ✓ Diversity loss distinguishes usage patterns\n")


def test_chunk_continuity_loss():
    """Test Chunk Continuity Loss"""
    print("=== Testing ChunkContinuityLoss ===")
    
    continuity_loss = ChunkContinuityLoss(
        boundary_weight=1.0,
        temporal_weight=0.1,
        chunk_size=16,
    )
    
    # Smooth sequence (should have low loss)
    smooth = torch.linspace(0, 1, 64).unsqueeze(-1).repeat(4, 1, 7)
    loss_smooth = continuity_loss(smooth)
    print(f"  Smooth sequence loss: {loss_smooth['total'].item():.6f}")
    print(f"    Boundary: {loss_smooth['boundary'].item():.6f}")
    print(f"    Temporal: {loss_smooth['temporal'].item():.6f}")
    
    # Discontinuous sequence (should have higher loss)
    discontinuous = torch.randn(4, 64, 7)
    loss_discontinuous = continuity_loss(discontinuous)
    print(f"  Discontinuous sequence loss: {loss_discontinuous['total'].item():.6f}")
    print(f"    Boundary: {loss_discontinuous['boundary'].item():.6f}")
    print(f"    Temporal: {loss_discontinuous['temporal'].item():.6f}")
    
    print("  ✓ Continuity loss computed successfully\n")


def test_actioncodec_loss():
    """Test combined ActionCodec Loss with dual-arm data"""
    print("=== Testing ActionCodecLoss ===")
    
    loss_fn = ActionCodecLoss(
        reconstruction_weight=1.0,
        commitment_weight=0.25,
        codebook_weight=1.0,
        diversity_weight=0.0,
        continuity_weight=0.0,
        n_codebooks=4,
    )
    
    batch_size = 4
    seq_len = 16
    # Dual-arm data: [batch, seq, 20]
    x = torch.randn(batch_size, seq_len, 20, requires_grad=True)
    x_recon = torch.randn(batch_size, seq_len, 20, requires_grad=True)
    
    # Simulate RVQ outputs
    z_e_list = [torch.randn(batch_size, seq_len, 512, requires_grad=True) for _ in range(4)]
    z_q_list = [torch.randn(batch_size, seq_len, 512, requires_grad=True) for _ in range(4)]
    
    loss_dict = loss_fn(
        x, x_recon,
        z_e_list=z_e_list,
        z_q_list=z_q_list,
    )
    
    print(f"  Total loss: {loss_dict['total'].item():.6f}")
    print(f"  Reconstruction loss: {loss_dict['reconstruction'].item():.6f}")
    print(f"  RVQ commitment loss: {loss_dict['rvq/commitment'].item():.6f}")
    print(f"  RVQ codebook loss: {loss_dict['rvq/codebook'].item():.6f}")
    
    # Test backward pass
    loss_dict['total'].backward()
    assert x.grad is not None, "Gradient should flow to x"
    assert x_recon.grad is not None, "Gradient should flow to x_recon"
    for i, z_e in enumerate(z_e_list):
        assert z_e.grad is not None, f"Gradient should flow to z_e[{i}]"
    print("  ✓ Backward pass successful\n")


def test_loss_info():
    """Test loading loss function info"""
    print("=== Testing load_info ===")
    info = load_info()
    for name, description in info.items():
        print(f"  {name}: {description}")
    print()


def test_oat_loss():
    """Test OAT Loss with all components"""
    print("=== Testing OATLoss ===")
    
    # Create OAT loss with default parameters
    oat_loss = OATLoss(
        position_weight=1.0,
        rotation_weight=1.0,
        gripper_weight=0.5,
        temporal_weight=0.1,
    )
    
    batch_size = 4
    seq_len = 64  # Longer sequence for temporal smoothness
    x = torch.randn(batch_size, seq_len, 20, requires_grad=True)
    x_recon = torch.randn(batch_size, seq_len, 20, requires_grad=True)
    
    # Compute loss
    loss_dict = oat_loss(x, x_recon)
    
    print(f"  Total loss: {loss_dict['total'].item():.6f}")
    print(f"  Reconstruction loss: {loss_dict['reconstruction'].item():.6f}")
    print(f"  Temporal smoothness: {loss_dict['temporal'].item():.6f}")
    
    # Test backward pass
    loss_dict['total'].backward()
    assert x.grad is not None, "Gradient should flow to x"
    assert x_recon.grad is not None, "Gradient should flow to x_recon"
    print("  ✓ Backward pass successful\n")


if __name__ == "__main__":
    print("=" * 60)
    print("PAT Loss Functions Test Suite")
    print("=" * 60 + "\n")
    
    test_loss_info()
    test_vq_loss()
    test_rvq_loss()
    test_reconstruction_loss()
    test_geodesic_loss()
    test_codebook_diversity_loss()
    test_chunk_continuity_loss()
    test_actioncodec_loss()
    test_oat_loss()
    
    print("=" * 60)
    print("All tests passed! ✓")
    print("=" * 60)
