import torch
import time
from timm.models import create_model

# Import your model registry to ensure custom modules are recognized
import models_mamba_ecg 

def main():
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Benchmarking on device: {device}")

    # Parameters mirroring your exact bash script configuration
    model_name = 'ecg_vim_small_patch16_stride8_224_bimambav2_final_pool_mean_abs_pos_embed_with_midclstok_div2'
    batch_size = 20
    channels = 12
    seq_len = 8192

    print(f"Creating model: {model_name} (depth=16)...")
    model = create_model(
        model_name,
        pretrained=False,
        num_classes=26,
        drop_rate=0.0,
        drop_path_rate=0.0,
        drop_block_rate=None,
        block='VisionMamba',
        depth=16,
        fused_add_norm=False,
        if_divide_out=False,
        use_middle_cls_token=False,
        img_size=seq_len
    )
    
    model.to(device)
    model.eval() # Ensure dropout/norms are locked in inference mode

    # Create dummy data directly on the GPU to isolate pure model throughput
    print(f"Allocating dummy input tensor (Batch: {batch_size}, Channels: {channels}, Length: {seq_len})...")
    dummy_input = torch.randn(batch_size, channels, seq_len, device=device)

    warmup_iters = 20
    benchmark_iters = 100

    # 1. WARM-UP PHASE
    print(f"Executing {warmup_iters} warm-up iterations (discarding timing data)...")
    with torch.no_grad():
        for _ in range(warmup_iters):
            _ = model(dummy_input)

    # 2. BENCHMARK PHASE
    print(f"Executing {benchmark_iters} official benchmark iterations...")
    
    # CRITICAL: Force Python to wait until all warm-up kernels finish before starting the clock
    if device.type == 'cuda':
        torch.cuda.synchronize()
        
    start_time = time.time()

    with torch.no_grad():
        for _ in range(benchmark_iters):
            _ = model(dummy_input)

    # CRITICAL: Force Python to wait until the final forward pass finishes before stopping the clock
    if device.type == 'cuda':
        torch.cuda.synchronize()

    end_time = time.time()
    total_time = end_time - start_time

    # 3. METRICS CALCULATION
    total_files_processed = benchmark_iters * batch_size
    throughput = total_files_processed / total_time
    ms_per_batch = (total_time / benchmark_iters) * 1000

    print("\n" + "="*50)
    print(" THROUGHPUT BENCHMARK RESULTS")
    print("="*50)
    print(f"Batch Size:            {batch_size}")
    print(f"Total Time Elapsed:    {total_time:.4f} seconds")
    print(f"Total Files Processed: {total_files_processed} files")
    print(f"Model Throughput:      {throughput:.2f} files/s")
    print(f"Latency per Batch:     {ms_per_batch:.2f} ms")
    print("="*50)

if __name__ == '__main__':
    main()