"""
ProtoVAR Smoke Test

A minimal test to verify the ProtoVAR implementation works correctly.
This test uses a small model and synthetic data for quick verification.
"""

import os
import sys
import time
import json
from pathlib import Path

import torch
import torch.nn as nn
import torch.nn.functional as F

# Add VAR to path
sys.path.insert(0, '/tmp/VAR')


class MiniVAR(nn.Module):
    """
    Miniature VAR model for smoke testing.
    
    This is a simplified version of VAR that can run on CPU
    without requiring pre-trained weights.
    """
    def __init__(
        self,
        num_classes: int = 10,
        embed_dim: int = 64,
        num_heads: int = 4,
        depth: int = 4,
        vocab_size: int = 256,
        patch_nums: tuple = (1, 2, 4, 8),
    ):
        super().__init__()
        self.num_classes = num_classes
        self.embed_dim = embed_dim
        self.patch_nums = patch_nums
        self.vocab_size = vocab_size
        
        # Embeddings
        self.class_emb = nn.Embedding(num_classes + 1, embed_dim)
        self.word_emb = nn.Embedding(vocab_size, embed_dim)
        
        # Position embeddings
        total_tokens = sum(pn**2 for pn in patch_nums)
        self.pos_emb = nn.Embedding(total_tokens, embed_dim)
        
        # Transformer blocks
        self.blocks = nn.ModuleList([
            nn.TransformerEncoderLayer(
                d_model=embed_dim,
                nhead=num_heads,
                dim_feedforward=embed_dim * 4,
                batch_first=True,
            )
            for _ in range(depth)
        ])
        
        # Output head
        self.head = nn.Linear(embed_dim, vocab_size)
    
    def forward(self, labels: torch.Tensor, tokens: torch.Tensor) -> torch.Tensor:
        """Forward pass for training."""
        B = labels.shape[0]
        
        # Class embedding
        cls_emb = self.class_emb(labels)  # (B, embed_dim)
        
        # Token embeddings
        token_emb = self.word_emb(tokens)  # (B, L, embed_dim)
        
        # Combine
        x = torch.cat([cls_emb.unsqueeze(1), token_emb], dim=1)
        
        # Add position embeddings
        pos = self.pos_emb(torch.arange(x.shape[1], device=x.device))
        x = x + pos.unsqueeze(0)
        
        # Transformer
        for block in self.blocks:
            x = block(x)
        
        return self.head(x)
    
    @torch.no_grad()
    def generate(self, class_label: int, num_samples: int = 1) -> torch.Tensor:
        """Generate samples for a class."""
        self.eval()
        
        batch_size = num_samples
        labels = torch.full((batch_size,), class_label, dtype=torch.long)
        
        # Start with class token
        x = self.class_emb(labels).unsqueeze(1)  # (B, 1, embed_dim)
        
        generated_tokens = []
        
        for scale_idx, pn in enumerate(self.patch_nums):
            # Generate tokens for this scale
            for i in range(pn * pn):
                # Add position embedding
                pos_idx = sum(p**2 for p in self.patch_nums[:scale_idx]) + i
                pos = self.pos_emb(torch.tensor([pos_idx], device=x.device))
                
                # Forward
                out = self.head(x[:, -1:, :] + pos)
                
                # Sample
                token = torch.argmax(out, dim=-1)  # (B, 1)
                generated_tokens.append(token)
                
                # Update x
                token_emb = self.word_emb(token)
                x = torch.cat([x, token_emb], dim=1)
        
        return torch.cat(generated_tokens, dim=1)  # (B, total_tokens)


class MiniPrototypeBank(nn.Module):
    """Miniature prototype bank for smoke testing."""
    def __init__(self, num_classes: int = 10, embed_dim: int = 64, num_scales: int = 4, input_dim: int = 256):
        super().__init__()
        self.num_classes = num_classes
        self.num_scales = num_scales
        
        self.prototypes = nn.Parameter(torch.randn(num_classes, num_scales, embed_dim) * 0.02)
        self.projections = nn.ModuleList([nn.Linear(input_dim, embed_dim) for _ in range(num_scales)])
    
    def get_guidance(self, class_idx: int, scale_idx: int) -> torch.Tensor:
        return self.prototypes[class_idx, scale_idx]


def smoke_test():
    """Run smoke test."""
    print("=" * 60)
    print("ProtoVAR Smoke Test")
    print("=" * 60)
    
    device = 'cpu'
    num_classes = 10
    
    # Create mini VAR model
    print("\n[1/3] Creating mini VAR model...")
    model = MiniVAR(
        num_classes=num_classes,
        embed_dim=64,
        num_heads=4,
        depth=4,
        vocab_size=256,
        patch_nums=(1, 2, 4, 8),
    )
    print(f"  Model parameters: {sum(p.numel() for p in model.parameters()) / 1e3:.1f}K")
    
    # Create prototype bank
    print("\n[2/3] Creating prototype bank...")
    proto_bank = MiniPrototypeBank(
        num_classes=num_classes,
        embed_dim=64,
        num_scales=4,
        input_dim=256,  # vocab_size
    )
    print(f"  Prototype parameters: {sum(p.numel() for p in proto_bank.parameters()) / 1e3:.1f}K")
    
    # Test generation
    print("\n[3/3] Testing generation...")
    model.eval()
    
    start_time = time.time()
    all_images = []
    all_labels = []
    
    for class_idx in range(num_classes):
        # Generate samples
        tokens = model.generate(class_label=class_idx, num_samples=5)
        all_images.append(tokens)
        all_labels.append(torch.full((5,), class_idx, dtype=torch.long))
        
        # Get prototype guidance
        guidance = proto_bank.get_guidance(class_idx, scale_idx=0)
        print(f"  Class {class_idx}: generated {tokens.shape[0]} samples, guidance shape: {guidance.shape}")
    
    elapsed = time.time() - start_time
    
    # Combine
    all_images = torch.cat(all_images, dim=0)
    all_labels = torch.cat(all_labels, dim=0)
    
    print(f"\n  Total generation time: {elapsed:.2f}s")
    print(f"  Generated dataset shape: {all_images.shape}")
    print(f"  Labels shape: {all_labels.shape}")
    
    # Test prototype matching loss
    print("\n[4/4] Testing prototype matching...")
    model.train()
    
    # Dummy forward pass
    dummy_labels = torch.randint(0, num_classes, (4,))
    dummy_tokens = torch.randint(0, 256, (4, 12))  # (B, L)
    logits = model(dummy_labels, dummy_tokens)
    
    # Compute prototype loss using logits projected to embed_dim
    total_loss = 0.0
    for scale_idx in range(4):
        # Project logits to embed_dim (from vocab_size)
        projected = proto_bank.projections[scale_idx](logits.mean(dim=1))  # (B, embed_dim)
        proto = proto_bank.prototypes[dummy_labels, scale_idx]  # (B, embed_dim)
        similarity = F.cosine_similarity(projected, proto, dim=-1)
        loss = (1 - similarity).mean()
        total_loss += loss
    
    print(f"  Prototype matching loss: {total_loss.item():.4f}")
    
    # Save results
    output_dir = './smoke_test_outputs'
    os.makedirs(output_dir, exist_ok=True)
    
    results = {
        'model_params': sum(p.numel() for p in model.parameters()),
        'proto_params': sum(p.numel() for p in proto_bank.parameters()),
        'generation_time': elapsed,
        'dataset_shape': list(all_images.shape),
        'prototype_loss': total_loss.item(),
    }
    
    with open(os.path.join(output_dir, 'smoke_test_results.json'), 'w') as f:
        json.dump(results, f, indent=2)
    
    print("\n" + "=" * 60)
    print("Smoke Test PASSED")
    print("=" * 60)
    print(f"Results saved to {output_dir}/smoke_test_results.json")
    
    return results


if __name__ == '__main__':
    smoke_test()