#!/usr/bin/env python3
"""
🧪 VAD Demo - Pre-deployment Testing & Optimization Script

This script helps you test and optimize your VAD demo before deploying
to Hugging Face Spaces for your WASPAA 2025 presentation.

Usage:
    python test_and_optimize.py --test-all
    python test_and_optimize.py --optimize-models
    python test_and_optimize.py --benchmark
"""

import sys
import time
import traceback
import argparse
import numpy as np
import torch
import psutil
import subprocess
from pathlib import Path
from typing import Dict, List, Tuple
import warnings
warnings.filterwarnings('ignore')

# ===== PERFORMANCE TESTING =====

class VADTester:
    """Comprehensive testing suite for VAD demo"""
    
    def __init__(self):
        self.test_results = {}
        self.performance_metrics = {}
        
    def test_dependencies(self) -> bool:
        """Test all required dependencies"""
        print("🔍 Testing Dependencies...")
        
        dependencies = [
            'gradio', 'numpy', 'torch', 'librosa', 
            'plotly', 'scipy', 'soundfile'
        ]
        
        missing = []
        for dep in dependencies:
            try:
                __import__(dep)
                print(f"  ✅ {dep}")
            except ImportError:
                print(f"  ❌ {dep}")
                missing.append(dep)
        
        if missing:
            print(f"\n⚠️  Missing dependencies: {missing}")
            print("Run: pip install " + " ".join(missing))
            return False
            
        print("✅ All dependencies available")
        return True
    
    def test_audio_generation(self) -> bool:
        """Test synthetic audio generation"""
        print("\n🎵 Testing Audio Generation...")
        
        try:
            # Generate test audio signals
            sample_rate = 16000
            duration = 4.0
            t = np.linspace(0, duration, int(sample_rate * duration))
            
            # Test signals
            test_signals = {
                'silence': np.zeros_like(t),
                'noise': np.random.normal(0, 0.1, len(t)),
                'tone': np.sin(2 * np.pi * 440 * t) * 0.5,
                'speech_sim': np.sin(2 * np.pi * 200 * t) * np.exp(-t/2) * 0.3
            }
            
            for name, signal in test_signals.items():
                if len(signal) == int(sample_rate * duration):
                    print(f"  ✅ {name} signal generated")
                else:
                    print(f"  ❌ {name} signal incorrect length")
                    return False
                    
            self.test_audio = test_signals
            print("✅ Audio generation working")
            return True
            
        except Exception as e:
            print(f"❌ Audio generation failed: {e}")
            return False
    
    def test_model_loading(self) -> Dict[str, bool]:
        """Test individual model loading"""
        print("\n🤖 Testing Model Loading...")
        
        # Import models from main app
        try:
            sys.path.append('.')
            from app import (OptimizedSileroVAD, OptimizedWebRTCVAD, 
                           OptimizedEPANNs, OptimizedAST, OptimizedPANNs)
            
            models = {
                'Silero-VAD': OptimizedSileroVAD,
                'WebRTC-VAD': OptimizedWebRTCVAD,
                'E-PANNs': OptimizedEPANNs,
                'AST': OptimizedAST,
                'PANNs': OptimizedPANNs
            }
            
            results = {}
            for name, model_class in models.items():
                try:
                    start_time = time.time()
                    model = model_class()
                    load_time = time.time() - start_time
                    
                    print(f"  ✅ {name} loaded ({load_time:.2f}s)")
                    results[name] = True
                    
                except Exception as e:
                    print(f"  ❌ {name} failed: {str(e)[:50]}...")
                    results[name] = False
            
            return results
            
        except ImportError as e:
            print(f"❌ Cannot import models from app.py: {e}")
            return {}
    
    def test_model_inference(self, model_results: Dict[str, bool]) -> Dict[str, float]:
        """Test model inference speed"""
        print("\n⚡ Testing Model Inference...")
        
        if not hasattr(self, 'test_audio'):
            print("❌ No test audio available")
            return {}
        
        try:
            from app import (OptimizedSileroVAD, OptimizedWebRTCVAD, 
                           OptimizedEPANNs, OptimizedAST, OptimizedPANNs)
            
            models = {}
            if model_results.get('Silero-VAD', False):
                models['Silero-VAD'] = OptimizedSileroVAD()
            if model_results.get('WebRTC-VAD', False):
                models['WebRTC-VAD'] = OptimizedWebRTCVAD()
            if model_results.get('E-PANNs', False):
                models['E-PANNs'] = OptimizedEPANNs()
            if model_results.get('AST', False):
                models['AST'] = OptimizedAST()
            if model_results.get('PANNs', False):
                models['PANNs'] = OptimizedPANNs()
            
            inference_times = {}
            test_audio = self.test_audio['speech_sim']
            
            for name, model in models.items():
                try:
                    # Warm-up run
                    model.predict(test_audio[:1000])
                    
                    # Benchmark runs
                    times = []
                    for _ in range(5):
                        start = time.time()
                        result = model.predict(test_audio)
                        times.append(time.time() - start)
                    
                    avg_time = np.mean(times)
                    inference_times[name] = avg_time
                    
                    # Check if real-time capable
                    is_realtime = avg_time < 4.0  # 4 second audio
                    status = "✅" if is_realtime else "⚠️ "
                    
                    print(f"  {status} {name}: {avg_time:.3f}s (RTF: {avg_time/4.0:.3f})")
                    
                except Exception as e:
                    print(f"  ❌ {name} inference failed: {str(e)[:50]}...")
                    inference_times[name] = float('inf')
            
            return inference_times
            
        except Exception as e:
            print(f"❌ Inference testing failed: {e}")
            return {}
    
    def test_memory_usage(self) -> Dict[str, float]:
        """Test memory usage of models"""
        print("\n💾 Testing Memory Usage...")
        
        try:
            import gc
            from app import VADDemo
            
            # Baseline memory
            gc.collect()
            baseline_mb = psutil.virtual_memory().used / 1024 / 1024
            
            # Load demo
            demo = VADDemo()
            gc.collect()
            demo_mb = psutil.virtual_memory().used / 1024 / 1024
            
            memory_usage = {
                'baseline': baseline_mb,
                'with_demo': demo_mb,
                'demo_overhead': demo_mb - baseline_mb
            }
            
            print(f"  📊 Baseline: {baseline_mb:.0f}MB")
            print(f"  📊 With Demo: {demo_mb:.0f}MB") 
            print(f"  📊 Demo Overhead: {memory_usage['demo_overhead']:.0f}MB")
            
            # Check if within HF Spaces limits (16GB)
            if demo_mb < 2000:  # 2GB threshold for safety
                print("  ✅ Memory usage acceptable for HF Spaces")
            else:
                print("  ⚠️  High memory usage - consider optimization")
            
            return memory_usage
            
        except Exception as e:
            print(f"❌ Memory testing failed: {e}")
            return {}
    
    def test_gradio_interface(self) -> bool:
        """Test Gradio interface creation"""
        print("\n🎨 Testing Gradio Interface...")
        
        try:
            from app import create_interface
            
            # Create interface (don't launch)
            interface = create_interface()
            
            if interface is not None:
                print("  ✅ Interface created successfully")
                
                # Check if queue is supported
                try:
                    interface.queue(max_size=5)
                    print("  ✅ Queue support working")
                except:
                    print("  ⚠️  Queue support limited")
                
                return True
            else:
                print("  ❌ Interface creation failed")
                return False
                
        except Exception as e:
            print(f"❌ Interface testing failed: {e}")
            return False
    
    def benchmark_full_pipeline(self) -> Dict[str, float]:
        """Benchmark complete processing pipeline"""
        print("\n🏁 Benchmarking Full Pipeline...")
        
        try:
            from app import VADDemo
            
            demo = VADDemo()
            test_audio = self.test_audio['speech_sim']
            
            # Simulate audio stream format
            audio_input = (16000, test_audio)  # (sample_rate, data)
            
            # Benchmark complete pipeline
            times = []
            for i in range(3):
                start = time.time()
                
                try:
                    result = demo.process_audio_stream(
                        audio_input, 
                        'Silero-VAD', 
                        'E-PANNs', 
                        0.5
                    )
                    
                    end = time.time()
                    times.append(end - start)
                    
                    print(f"  🔄 Run {i+1}: {end-start:.3f}s")
                    
                except Exception as e:
                    print(f"  ❌ Run {i+1} failed: {e}")
                    times.append(float('inf'))
            
            avg_time = np.mean([t for t in times if t != float('inf')])
            
            if avg_time < 1.0:
                print(f"  ✅ Pipeline average: {avg_time:.3f}s (excellent)")
            elif avg_time < 2.0:
                print(f"  ✅ Pipeline average: {avg_time:.3f}s (good)")
            else:
                print(f"  ⚠️  Pipeline average: {avg_time:.3f}s (slow)")
            
            return {'avg_pipeline_time': avg_time, 'all_times': times}
            
        except Exception as e:
            print(f"❌ Pipeline benchmarking failed: {e}")
            return {}

# ===== OPTIMIZATION UTILITIES =====

class VADOptimizer:
    """Optimization utilities for VAD demo"""
    
    def __init__(self):
        pass
    
    def optimize_torch_settings(self):
        """Optimize PyTorch for CPU inference"""
        print("🔧 Optimizing PyTorch Settings...")
        
        try:
            import torch
            
            # Set CPU threads for optimal performance
            cpu_count = psutil.cpu_count(logical=False)
            torch.set_num_threads(min(cpu_count, 4))  # Don't exceed 4 threads
            
            # Disable gradient computation globally
            torch.set_grad_enabled(False)
            
            # Use optimized CPU operations
            if hasattr(torch.backends, 'mkldnn'):
                torch.backends.mkldnn.enabled = True
                print("  ✅ MKL-DNN enabled")
            
            print(f"  ✅ CPU threads set to: {torch.get_num_threads()}")
            print("  ✅ Gradients disabled globally")
            
        except Exception as e:
            print(f"❌ PyTorch optimization failed: {e}")
    
    def create_optimized_requirements(self):
        """Create optimized requirements.txt"""
        print("📦 Creating Optimized Requirements...")
        
        optimized_requirements = """# Core dependencies - CPU optimized
gradio>=4.0.0,<5.0.0
numpy>=1.21.0,<1.25.0
torch>=2.0.0,<2.1.0
torchaudio>=2.0.0,<2.1.0

# Audio processing - optimized versions
librosa>=0.10.0,<0.11.0
soundfile>=0.12.1,<0.13.0
scipy>=1.9.0,<1.12.0

# Visualization - stable version
plotly>=5.15.0,<5.17.0

# Machine learning - pinned versions
transformers>=4.30.0,<4.35.0
datasets>=2.12.0,<2.15.0

# Optional dependencies with fallbacks
webrtcvad>=2.0.10; sys_platform != "darwin"
scikit-learn>=1.1.0,<1.4.0

# System utilities
psutil>=5.9.0
matplotlib>=3.5.0,<3.8.0

# Memory optimization
pympler>=0.9; python_version >= "3.8"
"""
        
        try:
            with open('requirements_optimized.txt', 'w') as f:
                f.write(optimized_requirements)
            print("  ✅ Optimized requirements.txt created")
            
            # Also create packages.txt for system dependencies
            system_packages = """ffmpeg
libsndfile1
libasound2-dev
portaudio19-dev
"""
            
            with open('packages_optimized.txt', 'w') as f:
                f.write(system_packages)
            print("  ✅ System packages.txt created")
            
        except Exception as e:
            print(f"❌ Requirements optimization failed: {e}")
    
    def create_deployment_config(self):
        """Create optimized deployment configuration"""
        print("⚙️  Creating Deployment Config...")
        
        # Create .gitattributes for Git LFS
        gitattributes = """*.pkl filter=lfs diff=lfs merge=lfs -text
*.bin filter=lfs diff=lfs merge=lfs -text
*.safetensors filter=lfs diff=lfs merge=lfs -text
*.onnx filter=lfs diff=lfs merge=lfs -text
*.h5 filter=lfs diff=lfs merge=lfs -text
"""
        
        try:
            with open('.gitattributes', 'w') as f:
                f.write(gitattributes)
            print("  ✅ .gitattributes created")
            
            # Create Dockerfile for local testing (optional)
            dockerfile = """FROM python:3.10-slim

WORKDIR /app

# System dependencies
RUN apt-get update && apt-get install -y \\
    ffmpeg \\
    libsndfile1 \\
    && rm -rf /var/lib/apt/lists/*

# Python dependencies
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# Copy application
COPY . .

# Expose port
EXPOSE 7860

# Run application
CMD ["python", "app.py"]
"""
            
            with open('Dockerfile', 'w') as f:
                f.write(dockerfile)
            print("  ✅ Dockerfile created for local testing")
            
        except Exception as e:
            print(f"❌ Deployment config failed: {e}")

# ===== MAIN TESTING INTERFACE =====

def run_comprehensive_test():
    """Run all tests and optimizations"""
    print("🧪 VAD Demo - Comprehensive Testing Suite")
    print("=" * 50)
    
    tester = VADTester()
    optimizer = VADOptimizer()
    
    # Optimization first
    print("\n🔧 OPTIMIZATION PHASE")
    optimizer.optimize_torch_settings()
    optimizer.create_optimized_requirements()
    optimizer.create_deployment_config()
    
    # Testing phase
    print("\n🧪 TESTING PHASE")
    
    # Test 1: Dependencies
    deps_ok = tester.test_dependencies()
    if not deps_ok:
        print("\n❌ Critical: Fix dependencies before proceeding")
        return False
    
    # Test 2: Audio generation
    audio_ok = tester.test_audio_generation()
    if not audio_ok:
        print("\n❌ Critical: Audio processing not working")
        return False
    
    # Test 3: Model loading
    model_results = tester.test_model_loading()
    working_models = sum(model_results.values())
    print(f"\n📊 Models Working: {working_models}/5")
    
    if working_models == 0:
        print("❌ Critical: No models working")
        return False
    elif working_models < 3:
        print("⚠️  Warning: Limited models available")
    
    # Test 4: Model inference
    inference_results = tester.test_model_inference(model_results)
    realtime_models = sum(1 for t in inference_results.values() if t < 4.0)
    print(f"\n📊 Real-time Models: {realtime_models}/{len(inference_results)}")
    
    # Test 5: Memory usage
    memory_results = tester.test_memory_usage()
    if memory_results:
        overhead = memory_results.get('demo_overhead', 0)
        if overhead > 1000:  # 1GB
            print("⚠️  Warning: High memory usage")
    
    # Test 6: Interface creation
    interface_ok = tester.test_gradio_interface()
    if not interface_ok:
        print("❌ Critical: Gradio interface not working")
        return False
    
    # Test 7: Full pipeline
    pipeline_results = tester.benchmark_full_pipeline()
    avg_time = pipeline_results.get('avg_pipeline_time', float('inf'))
    
    # Final assessment
    print("\n" + "=" * 50)
    print("📋 FINAL ASSESSMENT")
    print("=" * 50)
    
    if deps_ok and audio_ok and interface_ok and working_models >= 2:
        if avg_time < 1.0 and realtime_models >= 2:
            print("🎉 EXCELLENT - Ready for WASPAA demo!")
            print("✅ All systems optimal")
        elif avg_time < 2.0 and realtime_models >= 1:
            print("✅ GOOD - Demo ready with minor optimizations")
            print("💡 Consider further model optimization")
        else:
            print("⚠️  ACCEPTABLE - Demo functional but slow")
            print("💡 Consider upgrading to GPU Spaces for presentation")
    else:
        print("❌ NOT READY - Critical issues need fixing")
        return False
    
    # Performance summary
    print(f"\n📊 Performance Summary:")
    print(f"   • Working Models: {working_models}/5")
    print(f"   • Real-time Models: {realtime_models}")
    print(f"   • Average Pipeline: {avg_time:.3f}s")
    if memory_results:
        print(f"   • Memory Overhead: {memory_results.get('demo_overhead', 0):.0f}MB")
    
    # Recommendations
    print(f"\n💡 Recommendations:")
    if working_models < 5:
        print("   • Check model loading errors above")
    if realtime_models < 3:
        print("   • Consider model optimization or GPU upgrade")
    if avg_time > 1.0:
        print("   • Optimize audio processing pipeline")
    
    print("\n🚀 Next Steps:")
    print("   1. Fix any critical issues above")
    print("   2. Use optimized files: requirements_optimized.txt")
    print("   3. Deploy to Hugging Face Spaces")
    print("   4. Test live demo URL before WASPAA")
    
    return True

def run_quick_test():
    """Run quick essential tests only"""
    print("⚡ VAD Demo - Quick Test")
    print("=" * 30)
    
    tester = VADTester()
    
    # Essential tests only
    deps_ok = tester.test_dependencies()
    audio_ok = tester.test_audio_generation()
    model_results = tester.test_model_loading()
    
    working_models = sum(model_results.values())
    
    if deps_ok and audio_ok and working_models >= 2:
        print("\n✅ QUICK TEST PASSED")
        print(f"Ready for deployment with {working_models} models")
        return True
    else:
        print("\n❌ QUICK TEST FAILED")
        print("Run --test-all for detailed diagnosis")
        return False

def main():
    parser = argparse.ArgumentParser(description='VAD Demo Testing & Optimization')
    parser.add_argument('--test-all', action='store_true', 
                       help='Run comprehensive test suite')
    parser.add_argument('--quick-test', action='store_true',
                       help='Run quick essential tests')
    parser.add_argument('--optimize', action='store_true',
                       help='Create optimized configuration files')
    parser.add_argument('--benchmark', action='store_true',
                       help='Run performance benchmarks only')
    
    args = parser.parse_args()
    
    if args.test_all:
        success = run_comprehensive_test()
        sys.exit(0 if success else 1)
    elif args.quick_test:
        success = run_quick_test()
        sys.exit(0 if success else 1)
    elif args.optimize:
        optimizer = VADOptimizer()
        optimizer.optimize_torch_settings()
        optimizer.create_optimized_requirements()
        optimizer.create_deployment_config()
        print("✅ Optimization complete")
    elif args.benchmark:
        tester = VADTester()
        tester.test_audio_generation()
        model_results = tester.test_model_loading()
        inference_results = tester.test_model_inference(model_results)
        pipeline_results = tester.benchmark_full_pipeline()
        print("📊 Benchmark complete")
    else:
        print("Usage: python test_and_optimize.py [--test-all|--quick-test|--optimize|--benchmark]")
        print("\nFor WASPAA demo preparation, run:")
        print("  python test_and_optimize.py --test-all")

if __name__ == "__main__":
    main()