#!/usr/bin/env python3
"""
Benchmark Soprano TTS on GTX 1660 (CUDA)
Tests isolated performance before dual-GPU setup
"""

import time
import torch
from soprano import SopranoTTS

print("="*60)
print("Soprano Benchmark on GTX 1660")
print("="*60)

# Verify GPU
print(f"\nGPU: {torch.cuda.get_device_name(0)}")
print(f"CUDA Version: {torch.version.cuda}")
print(f"PyTorch: {torch.__version__}")

# Load model
print("\nLoading Soprano with lmdeploy backend...")
model = SopranoTTS(
    backend='lmdeploy',
    device='cuda',
    cache_size_mb=500,
    decoder_batch_size=2
)
print(f"✓ Loaded on {model.device}")

# Test texts of varying lengths
test_cases = [
    ("Short test for quick generation.", "Short"),
    ("This is a medium length test to measure the performance of Soprano text to speech on the GTX 1660 GPU.", "Medium"),
    ("This is a longer test sentence designed to evaluate sustained performance of the Soprano text to speech model running on the NVIDIA GTX 1660 graphics card with CUDA acceleration to see if it can achieve realtime or better performance.", "Long"),
    ("Extended performance test with even more text to really push the limits and see how the GTX 1660 handles longer generation tasks with the Soprano model using the lmdeploy backend which should provide optimized inference performance compared to the standard transformers backend.", "Very Long"),
]

print("\n" + "="*60)
print("Running Benchmarks...")
print("="*60)

results = []

for text, label in test_cases:
    print(f"\n[{label}] Text: \"{text[:50]}...\"")
    print(f"  Length: {len(text)} chars")
    
    # Warm up with first run
    if label == "Short":
        print("  (Warming up...)")
        _ = model.infer(text)
        time.sleep(1)
    
    # Actual benchmark - 3 runs
    times = []
    audio_lengths = []
    
    for run in range(3):
        start = time.time()
        audio = model.infer(text)
        elapsed = time.time() - start
        times.append(elapsed)
        
        # Calculate audio duration (32kHz, mono)
        audio_duration = len(audio) / 32000
        audio_lengths.append(audio_duration)
        
        realtime_factor = audio_duration / elapsed
        print(f"  Run {run+1}: {elapsed:.2f}s for {audio_duration:.2f}s audio = {realtime_factor:.2f}x realtime")
    
    # Average results
    avg_time = sum(times) / len(times)
    avg_audio = sum(audio_lengths) / len(audio_lengths)
    avg_realtime = avg_audio / avg_time
    
    results.append({
        'label': label,
        'text_len': len(text),
        'avg_time': avg_time,
        'avg_audio': avg_audio,
        'avg_realtime': avg_realtime
    })
    
    print(f"  → Average: {avg_time:.2f}s for {avg_audio:.2f}s audio = {avg_realtime:.2f}x realtime")

# Summary
print("\n" + "="*60)
print("SUMMARY")
print("="*60)
print(f"{'Text Length':<15} {'Gen Time':<12} {'Audio':<10} {'Realtime Factor':<15}")
print("-"*60)

for r in results:
    print(f"{r['label']:<15} {r['avg_time']:>8.2f}s   {r['avg_audio']:>6.2f}s   {r['avg_realtime']:>12.2f}x")

overall_avg = sum(r['avg_realtime'] for r in results) / len(results)
print("-"*60)
print(f"Overall Average: {overall_avg:.2f}x realtime")

# Verdict
print("\n" + "="*60)
print("VERDICT")
print("="*60)

if overall_avg >= 1.5:
    print("✅ EXCELLENT - GTX 1660 is fast enough!")
    print(f"   {overall_avg:.2f}x realtime is plenty for live generation")
elif overall_avg >= 1.0:
    print("✅ GOOD - GTX 1660 can handle realtime")
    print(f"   {overall_avg:.2f}x realtime is sufficient")
elif overall_avg >= 0.8:
    print("⚠️  MARGINAL - GTX 1660 is borderline")
    print(f"   {overall_avg:.2f}x realtime might work with optimization")
else:
    print("❌ TOO SLOW - GTX 1660 cannot keep up")
    print(f"   {overall_avg:.2f}x realtime is insufficient")
    print("   Consider using RX 6800 for Soprano instead")

print("="*60)
