import torch
import time
from cnn1 import UNet
from thop import profile
import numpy as np

def benchmark(device="cpu"):
    print(f"--- Benchmarking on {device.upper()} ---")
    
    # 1. Setup Model
    model = UNet(text_dim=512).to(device)
    model.eval()
    
    # 2. Setup Dummy Inputs (Standard 512x512 image + CLIP embedding)
    dummy_image = torch.randn(1, 3, 512, 512).to(device)
    dummy_text = torch.randn(1, 512).to(device) # Normalized text embedding
    
    # --- METRIC 1: FLOPs & Parameters ---
    if device == "cpu": # Only need to calc this once
        print("Calculating FLOPs and Params...")
        macs, params = profile(model, inputs=(dummy_image, dummy_text), verbose=False)
        print(f"Parameters: {params / 1e6:.2f} M")
        print(f"GFLOPs:     {macs / 1e9:.2f} G")
    
    # --- METRIC 2: Inference Time ---
    print("Measuring Inference Speed...")
    
    # Warmup (get cache ready)
    for _ in range(5):
        with torch.no_grad():
            _ = model(dummy_image, dummy_text)
            
    # Measure
    latencies = []
    with torch.no_grad():
        for _ in range(50): # Run 50 times
            start = time.time()
            _ = model(dummy_image, dummy_text)
            if device == "cuda":
                torch.cuda.synchronize() # Wait for GPU to finish
            end = time.time()
            latencies.append(end - start)
            
    avg_time = np.mean(latencies)
    print(f"Avg Inference Time: {avg_time:.4f} seconds")
    print(f"FPS: {1/avg_time:.2f}")
    print("-" * 30)

if __name__ == "__main__":
    # Run on CPU
    benchmark("cpu")
    
    # Run on GPU (If available, for comparison)
    if torch.cuda.is_available():
        benchmark("cuda")