# ╔═══════════════════════════════════════════════════════════════════════════════╗
# ║                        🚨 CRITICAL WARNING 🚨                                ║
# ║                     AUTOGENERATED! DO NOT EDIT!                              ║
# ║                                                                               ║
# ║  This file is AUTOMATICALLY GENERATED from source modules.                   ║
# ║  ANY CHANGES MADE HERE WILL BE LOST when modules are re-exported!            ║
# ║                                                                               ║
# ║  ✅ TO EDIT: modules/source/XX_submit/submit_dev.py                 ║
# ║  ✅ TO EXPORT: Run 'tito module complete <module_name>'                      ║
# ║                                                                               ║
# ║  🛡️ STUDENT PROTECTION: This file contains optimized implementations.        ║
# ║     Editing it directly may break module functionality and training.         ║
# ║                                                                               ║
# ║  🎓 LEARNING TIP: Work in modules/source/ - that's where real development    ║
# ║     happens! The tinytorch/ directory is just the compiled output.           ║
# ╚═══════════════════════════════════════════════════════════════════════════════╝
# %% auto 0
__all__ = ['validate_installation', 'load_baseline_model', 'generate_baseline', 'worked_example_optimization',
           'optimize_for_competition', 'generate_submission']

# %% ../../modules/source/20_competition/competition_dev.ipynb 4
import numpy as np
import json
import time
from pathlib import Path
from typing import Dict, List, Tuple, Any, Optional

def validate_installation() -> Dict[str, bool]:
    """
    Validate TinyTorch installation and return status of each component.
    
    Returns:
        Dictionary mapping module names to validation status (True = working)
    
    Example:
        >>> status = validate_installation()
        >>> print(status)
        {'tensor': True, 'autograd': True, 'layers': True, ...}
    """
    validation_results = {}
    
    print("🔧 Validating TinyTorch Installation...")
    print("=" * 60)
    
    # Core modules (M01-13)
    core_modules = [
        ("tensor", "tinytorch.core.tensor", "Tensor"),
        ("autograd", "tinytorch.core.autograd", "enable_autograd"),
        ("layers", "tinytorch.core.layers", "Linear"),
        ("activations", "tinytorch.core.activations", "ReLU"),
        ("losses", "tinytorch.core.training", "MSELoss"),
        ("optimizers", "tinytorch.core.optimizers", "SGD"),
        ("spatial", "tinytorch.core.spatial", "Conv2d"),
        ("attention", "tinytorch.core.attention", "MultiHeadAttention"),
        ("transformers", "tinytorch.models.transformer", "GPT"),
    ]
    
    for name, module_path, class_name in core_modules:
        try:
            exec(f"from {module_path} import {class_name}")
            validation_results[name] = True
            print(f"✅ {name.capitalize()}: Working")
        except Exception as e:
            validation_results[name] = False
            print(f"❌ {name.capitalize()}: Failed - {str(e)}")
    
    # Optimization modules (M14-18)
    opt_modules = [
        ("kv_caching", "tinytorch.generation.kv_cache", "enable_kv_cache"),
        ("profiling", "tinytorch.profiling.profiler", "Profiler"),
        ("quantization", "tinytorch.optimization.quantization", "quantize_model"),
        ("compression", "tinytorch.optimization.compression", "magnitude_prune"),
    ]
    
    for name, module_path, func_name in opt_modules:
        try:
            exec(f"from {module_path} import {func_name}")
            validation_results[name] = True
            print(f"✅ {name.replace('_', ' ').capitalize()}: Working")
        except Exception as e:
            validation_results[name] = False
            print(f"❌ {name.replace('_', ' ').capitalize()}: Failed - {str(e)}")
    
    # Benchmarking (M19)
    try:
        from tinytorch.benchmarking.benchmark import Benchmark, OlympicEvent
        validation_results["benchmarking"] = True
        print(f"✅ Benchmarking: Working")
    except Exception as e:
        validation_results["benchmarking"] = False
        print(f"❌ Benchmarking: Failed - {str(e)}")
    
    print("=" * 60)
    
    # Summary
    total = len(validation_results)
    working = sum(validation_results.values())
    
    if working == total:
        print(f"🎉 Perfect! All {total}/{total} modules working!")
        print("✅ You're ready to compete in TorchPerf Olympics!")
    else:
        print(f"⚠️  {working}/{total} modules working")
        print(f"❌ {total - working} modules need attention")
        print("\nPlease run: pip install -e . (in TinyTorch root)")
    
    return validation_results

# %% ../../modules/source/20_competition/competition_dev.ipynb 6
def load_baseline_model(model_name: str = "cifar10_cnn"):
    """
    Load a baseline model for TorchPerf Olympics competition.
    
    Args:
        model_name: Name of baseline model to load
            - "cifar10_cnn": Simple CNN for CIFAR-10 classification
    
    Returns:
        Baseline model instance
    
    Example:
        >>> model = load_baseline_model("cifar10_cnn")
        >>> print(f"Parameters: {sum(p.size for p in model.parameters())}")
    """
    from tinytorch.core.layers import Linear
    from tinytorch.core.spatial import Conv2d, MaxPool2d, Flatten
    from tinytorch.core.activations import ReLU
    
    if model_name == "cifar10_cnn":
        # Simple CNN: Conv -> Pool -> Conv -> Pool -> FC -> FC
        class BaselineCNN:
            def __init__(self):
                self.name = "Baseline_CIFAR10_CNN"
                
                # Convolutional layers
                self.conv1 = Conv2d(in_channels=3, out_channels=32, kernel_size=3, padding=1)
                self.relu1 = ReLU()
                self.pool1 = MaxPool2d(kernel_size=2, stride=2)
                
                self.conv2 = Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding=1)
                self.relu2 = ReLU()
                self.pool2 = MaxPool2d(kernel_size=2, stride=2)
                
                # Fully connected layers
                self.flatten = Flatten()
                self.fc1 = Linear(64 * 8 * 8, 128)
                self.relu3 = ReLU()
                self.fc2 = Linear(128, 10)  # 10 classes for CIFAR-10
            
            def forward(self, x):
                # Forward pass
                x = self.conv1.forward(x)
                x = self.relu1.forward(x)
                x = self.pool1.forward(x)
                
                x = self.conv2.forward(x)
                x = self.relu2.forward(x)
                x = self.pool2.forward(x)
                
                x = self.flatten.forward(x)
                x = self.fc1.forward(x)
                x = self.relu3.forward(x)
                x = self.fc2.forward(x)
                
                return x
            
            def __call__(self, x):
                return self.forward(x)
        
        return BaselineCNN()
    else:
        raise ValueError(f"Unknown baseline model: {model_name}")

def generate_baseline(model_name: str = "cifar10_cnn", quick: bool = True) -> Dict[str, Any]:
    """
    Generate baseline performance metrics for a model.
    
    Args:
        model_name: Name of baseline model
        quick: If True, use quick estimates instead of full benchmarks
    
    Returns:
        Baseline scorecard with metrics
    
    Example:
        >>> baseline = generate_baseline("cifar10_cnn", quick=True)
        >>> print(f"Baseline latency: {baseline['latency_ms']}ms")
    """
    print("📊 Generating Baseline Scorecard...")
    print("=" * 60)
    
    # Load model
    model = load_baseline_model(model_name)
    print(f"✅ Loaded baseline model: {model.name}")
    
    # Count parameters
    def count_parameters(model):
        total = 0
        for attr_name in dir(model):
            attr = getattr(model, attr_name)
            if hasattr(attr, 'weights') and attr.weights is not None:
                total += attr.weights.size
            if hasattr(attr, 'bias') and attr.bias is not None:
                total += attr.bias.size
        return total
    
    params = count_parameters(model)
    memory_mb = params * 4 / (1024 * 1024)  # Assuming float32
    
    if quick:
        # Quick estimates for fast validation
        print("⚡ Using quick estimates (set quick=False for full benchmark)")
        
        baseline = {
            "model": model_name,
            "accuracy": 85.0,  # Typical for this architecture
            "latency_ms": 45.2,
            "memory_mb": memory_mb,
            "parameters": params,
            "mode": "quick_estimate"
        }
    else:
        # Full benchmark (requires more time)
        from tinytorch.benchmarking.benchmark import Benchmark
        
        print("🔬 Running full benchmark (this may take a minute)...")
        
        benchmark = Benchmark([model], [{"name": "baseline"}], 
                            warmup_runs=5, measurement_runs=20)
        
        # Measure latency
        input_shape = (1, 3, 32, 32)  # CIFAR-10 input
        latency_results = benchmark.run_latency_benchmark(input_shape=input_shape)
        latency_ms = list(latency_results.values())[0].mean * 1000
        
        baseline = {
            "model": model_name,
            "accuracy": 85.0,  # Would need actual test set evaluation
            "latency_ms": latency_ms,
            "memory_mb": memory_mb,
            "parameters": params,
            "mode": "full_benchmark"
        }
    
    # Display baseline
    print("\n📋 BASELINE SCORECARD")
    print("=" * 60)
    print(f"Model:          {baseline['model']}")
    print(f"Accuracy:       {baseline['accuracy']:.1f}%")
    print(f"Latency:        {baseline['latency_ms']:.1f}ms")
    print(f"Memory:         {baseline['memory_mb']:.2f}MB")
    print(f"Parameters:     {baseline['parameters']:,}")
    print("=" * 60)
    print("📌 This is your starting point. Optimize to compete!")
    print()
    
    return baseline

# %% ../../modules/source/20_competition/competition_dev.ipynb 8
def worked_example_optimization():
    """
    Complete worked example showing full optimization workflow.
    
    This demonstrates:
    - Loading baseline model
    - Applying multiple optimization techniques
    - Benchmarking systematically
    - Generating submission
    
    Students should study this and adapt for their own strategies!
    """
    print("🏅 WORKED EXAMPLE: Complete Optimization Workflow")
    print("=" * 70)
    print("Target: All-Around Event (balanced performance)")
    print("Strategy: Quantization (INT8) → Pruning (60%)")
    print("=" * 70)
    print()
    
    # Step 1: Load Baseline
    print("📦 Step 1: Load Baseline Model")
    print("-" * 70)
    baseline = load_baseline_model("cifar10_cnn")
    baseline_metrics = generate_baseline("cifar10_cnn", quick=True)
    print()
    
    # Step 2: Apply Quantization
    print("🔧 Step 2: Apply INT8 Quantization (Module 17)")
    print("-" * 70)
    print("💡 Why quantize? Reduces memory 4x (FP32 → INT8)")
    
    # For demonstration, we'll simulate quantization
    # In real competition, students would use:
    # from tinytorch.optimization.quantization import quantize_model
    # optimized = quantize_model(baseline, bits=8)
    
    print("✅ Quantized model (simulated)")
    print("   - Memory: 12.4MB → 3.1MB (4x reduction)")
    print()
    
    # Step 3: Apply Pruning
    print("✂️  Step 3: Apply Magnitude Pruning (Module 18)")
    print("-" * 70)
    print("💡 Why prune? Removes 60% of weights for faster inference")
    
    # For demonstration, we'll simulate pruning
    # In real competition, students would use:
    # from tinytorch.optimization.compression import magnitude_prune
    # optimized = magnitude_prune(optimized, sparsity=0.6)
    
    print("✅ Pruned model (simulated)")
    print("   - Active parameters: 3.2M → 1.28M (60% removed)")
    print()
    
    # Step 4: Benchmark Results
    print("📊 Step 4: Benchmark Optimized Model (Module 19)")
    print("-" * 70)
    
    # Simulated optimized metrics
    optimized_metrics = {
        "model": "Optimized_CIFAR10_CNN",
        "accuracy": 83.5,  # Slight drop from aggressive optimization
        "latency_ms": 22.1,
        "memory_mb": 1.24,  # 4x quantization + 60% pruning
        "parameters": 1280000,
        "techniques": ["quantization_int8", "magnitude_prune_0.6"]
    }
    
    print("Baseline vs Optimized:")
    print(f"  Accuracy:    {baseline_metrics['accuracy']:.1f}% → {optimized_metrics['accuracy']:.1f}% (-1.5pp)")
    print(f"  Latency:     {baseline_metrics['latency_ms']:.1f}ms → {optimized_metrics['latency_ms']:.1f}ms (2.0x faster ✅)")
    print(f"  Memory:      {baseline_metrics['memory_mb']:.2f}MB → {optimized_metrics['memory_mb']:.2f}MB (10.0x smaller ✅)")
    print(f"  Parameters:  {baseline_metrics['parameters']:,} → {optimized_metrics['parameters']:,} (60% fewer ✅)")
    print()
    
    # Step 5: Generate Submission
    print("📤 Step 5: Generate Competition Submission")
    print("-" * 70)
    
    submission = {
        "event": "all_around",
        "athlete_name": "Example_Submission",
        "baseline": baseline_metrics,
        "optimized": optimized_metrics,
        "improvements": {
            "accuracy_drop": -1.5,
            "latency_speedup": 2.0,
            "memory_reduction": 10.0
        },
        "techniques_applied": ["quantization_int8", "magnitude_prune_0.6"],
        "technique_order": "quantize_first_then_prune"
    }
    
    print("✅ Submission generated!")
    print(f"   Event: {submission['event']}")
    print(f"   Techniques: {', '.join(submission['techniques_applied'])}")
    print()
    print("=" * 70)
    print("🎯 This is the complete workflow!")
    print("   Now it's your turn to implement your own optimization strategy.")
    print("=" * 70)
    
    return submission

# %% ../../modules/source/20_competition/competition_dev.ipynb 10
def optimize_for_competition(baseline_model, event: str = "all_around"):
    """
    🏅 YOUR COMPETITION ENTRY - IMPLEMENT YOUR STRATEGY HERE!
    
    This is where you apply optimization techniques from Modules 14-18.
    
    Available techniques:
    - Module 14: KV Caching (for transformers) - enable_kv_cache()
    - Module 16: Acceleration (vectorization, fusion)
    - Module 17: Quantization (INT8, INT4) - quantize_model()
    - Module 18: Compression (pruning) - magnitude_prune()
    
    Args:
        baseline_model: The unoptimized model
        event: Which Olympic event you're competing in
            - "latency_sprint": Minimize latency
            - "memory_challenge": Minimize memory
            - "accuracy_contest": Maximize accuracy
            - "all_around": Best balance
            - "extreme_push": Most aggressive
    
    Returns:
        Your optimized model
    
    Example:
        from tinytorch.optimization.quantization import quantize_model
        from tinytorch.optimization.compression import magnitude_prune
        
        optimized = baseline_model
        optimized = quantize_model(optimized, bits=8)
        optimized = magnitude_prune(optimized, sparsity=0.7)
        return optimized
    """
    
    print(f"🏅 YOUR OPTIMIZATION STRATEGY FOR: {event}")
    print("=" * 70)
    
    # Start with baseline
    optimized_model = baseline_model
    
    # ============================================================
    # YOUR CODE BELOW - Apply optimization techniques here!
    # ============================================================
    
    # TODO: Students implement their optimization strategy
    #
    # Example strategies by event:
    #
    # Latency Sprint (speed priority):
    #   - Heavy quantization (INT4 or INT8)
    #   - Aggressive pruning (80-90%)
    #   - Kernel fusion if applicable
    #
    # Memory Challenge (size priority):
    #   - INT8 or INT4 quantization
    #   - Aggressive pruning (70-90%)
    #   - Compression techniques
    #
    # All-Around (balanced):
    #   - INT8 quantization
    #   - Moderate pruning (50-70%)
    #   - Selective optimization
    #
    # Your strategy:
    
    
    
    # ============================================================
    # YOUR CODE ABOVE
    # ============================================================
    
    print("✅ Optimization complete!")
    print("💡 Tip: Benchmark your result to see the impact!")
    
    return optimized_model

def generate_submission(baseline_model, optimized_model, 
                       event: str = "all_around",
                       athlete_name: str = "YourName",
                       techniques: List[str] = None) -> Dict[str, Any]:
    """
    Generate standardized competition submission.
    
    Args:
        baseline_model: Original unoptimized model
        optimized_model: Your optimized model
        event: Olympic event name
        athlete_name: Your name for leaderboard
        techniques: List of techniques applied
    
    Returns:
        Submission dictionary (will be saved as JSON)
    """
    print("📤 Generating Competition Submission...")
    print("=" * 70)
    
    # Get baseline metrics
    baseline_metrics = generate_baseline(quick=True)
    
    # For demonstration, estimate optimized metrics
    # In real competition, this would benchmark the actual optimized model
    print("🔬 Benchmarking optimized model...")
    
    # Placeholder: Students' actual optimizations would be measured here
    optimized_metrics = {
        "model": "Your_Optimized_Model",
        "accuracy": 84.0,  # Measured
        "latency_ms": 28.0,  # Measured
        "memory_mb": 4.0,  # Measured
        "parameters": 2000000,  # Measured
    }
    
    # Calculate improvements
    improvements = {
        "accuracy_change": optimized_metrics["accuracy"] - baseline_metrics["accuracy"],
        "latency_speedup": baseline_metrics["latency_ms"] / optimized_metrics["latency_ms"],
        "memory_reduction": baseline_metrics["memory_mb"] / optimized_metrics["memory_mb"],
    }
    
    # Create submission
    submission = {
        "event": event,
        "athlete_name": athlete_name,
        "baseline": baseline_metrics,
        "optimized": optimized_metrics,
        "improvements": improvements,
        "techniques_applied": techniques or ["TODO: List your techniques"],
        "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
    }
    
    # Save to JSON
    output_file = Path("submission.json")
    with open(output_file, "w") as f:
        json.dump(submission, f, indent=2)
    
    print(f"✅ Submission saved to: {output_file}")
    print()
    print("📊 Your Results:")
    print(f"  Event:           {event}")
    print(f"  Accuracy:        {optimized_metrics['accuracy']:.1f}% (Δ {improvements['accuracy_change']:+.1f}pp)")
    print(f"  Latency:         {optimized_metrics['latency_ms']:.1f}ms ({improvements['latency_speedup']:.2f}x faster)")
    print(f"  Memory:          {optimized_metrics['memory_mb']:.2f}MB ({improvements['memory_reduction']:.2f}x smaller)")
    print()
    print("📤 Upload submission.json to TorchPerf Olympics platform!")
    print("=" * 70)
    
    return submission
