"""
Progressive Optimization Examples

This module demonstrates how the same neural network computations can be implemented
with different levels of kernel optimization, showing the computational
equivalence while highlighting performance improvements.
"""

import time

import torch
import torch.nn as nn

# Import our optimized components
from ..core.components.basic_optimized import (
    OptimizedTransformerBlock as BasicTransformerBlock,
)
from ..core.components.jit_optimized import FullyJITTransformerBlock

try:
    from .triton_fused_ops import (
        TritonLayerNorm,  # noqa: F401
        TritonOptimizedTransformerBlock,
        TritonSwiGLU,  # noqa: F401
    )
    TRITON_AVAILABLE = True
except ImportError:
    TRITON_AVAILABLE = False
    print("Triton not available, skipping Triton optimizations")

try:
    import torchbridge_cuda  # noqa: F401
    CUDA_AVAILABLE = True
except ImportError:
    CUDA_AVAILABLE = False
    print("Custom CUDA kernels not available")


class ProgressiveOptimizationDemo:
    """
    Demonstrates the same transformer model implemented with
    different levels of optimization to show computational equivalence
    and performance improvements.
    """

    def __init__(self, device: str = "cuda"):
        self.device = device
        self.models = {}
        self.results = {}

    def create_models(self, dim: int = 512, num_heads: int = 8, vocab_size: int = 10000):
        """Create models at different optimization levels"""

        # Level 1: Basic PyTorch optimized
        self.models['basic'] = TransformerModel(
            vocab_size, dim, num_layers=6, num_heads=num_heads,
            block_type='basic'
        ).to(self.device)

        # Level 2: TorchScript JIT
        self.models['jit'] = TransformerModel(
            vocab_size, dim, num_layers=6, num_heads=num_heads,
            block_type='jit'
        ).to(self.device)

        # Level 3: torch.compile (if available)
        if hasattr(torch, 'compile'):
            self.models['compiled'] = torch.compile(
                self.models['basic'].clone()
            )

        # Level 4: Triton kernels
        if TRITON_AVAILABLE:
            self.models['triton'] = TransformerModel(
                vocab_size, dim, num_layers=6, num_heads=num_heads,
                block_type='triton'
            ).to(self.device)

        # Level 5: Custom CUDA (simulated interface)
        if CUDA_AVAILABLE:
            self.models['cuda'] = CUDAOptimizedTransformer(
                vocab_size, dim, num_layers=6, num_heads=num_heads
            ).to(self.device)

    def benchmark_models(self, input_ids: torch.Tensor, num_iterations: int = 100):
        """Benchmark all available models"""
        print(f"Benchmarking models with input shape: {input_ids.shape}")
        print("="*60)

        for name, model in self.models.items():
            print(f"Benchmarking {name.upper()}...")

            # Warmup
            with torch.no_grad():
                for _ in range(10):
                    _ = model(input_ids)
                if input_ids.is_cuda:
                    torch.cuda.synchronize()

            # Benchmark
            times = []
            with torch.no_grad():
                for _ in range(num_iterations):
                    start_time = time.perf_counter()
                    output = model(input_ids)
                    if input_ids.is_cuda:
                        torch.cuda.synchronize()
                    end_time = time.perf_counter()
                    times.append(end_time - start_time)

            self.results[name] = {
                'mean_time': sum(times) / len(times),
                'min_time': min(times),
                'max_time': max(times),
                'throughput_tokens_per_sec': input_ids.numel() / (sum(times) / len(times)),
                'output_shape': output.shape
            }

    def compare_outputs(self, input_ids: torch.Tensor, tolerance: float = 1e-5):
        """Verify that all models produce computationally equivalent outputs"""
        print("Comparing model outputs for computational equivalence...")
        print("="*60)

        outputs = {}
        with torch.no_grad():
            for name, model in self.models.items():
                outputs[name] = model(input_ids)

        # Compare outputs pairwise
        model_names = list(outputs.keys())
        for i, name1 in enumerate(model_names):
            for name2 in model_names[i+1:]:
                diff = torch.abs(outputs[name1] - outputs[name2])
                max_diff = diff.max().item()
                mean_diff = diff.mean().item()

                status = " PASS" if max_diff < tolerance else " FAIL"
                print(f"{name1:10s} vs {name2:10s}: max_diff={max_diff:.2e}, mean_diff={mean_diff:.2e} [{status}]")

    def print_results(self):
        """Print benchmark results in a nice format"""
        print("="*80)
        print(f"{'Model':<15} {'Mean Time (ms)':<15} {'Min Time (ms)':<15} {'Throughput (tok/s)':<20} {'Speedup':<10}")
        print("-"*80)

        # Calculate speedups relative to basic model
        baseline_time = self.results.get('basic', {}).get('mean_time', 1.0)

        for name, result in sorted(self.results.items()):
            mean_time_ms = result['mean_time'] * 1000
            min_time_ms = result['min_time'] * 1000
            throughput = result['throughput_tokens_per_sec']
            speedup = baseline_time / result['mean_time']

            print(f"{name:<15} {mean_time_ms:<15.2f} {min_time_ms:<15.2f} {throughput:<20.0f} {speedup:<10.2f}x")

    def demonstrate_kernel_concepts(self):
        """
        Educational demonstration of key kernel optimization concepts
        """
        print("\nKernel Optimization Concepts Demonstration:")
        print("="*60)

        # 1. Memory Coalescing Example
        self._demonstrate_memory_coalescing()

        # 2. Kernel Fusion Example
        self._demonstrate_kernel_fusion()

        # 3. Memory Hierarchy Usage
        self._demonstrate_memory_hierarchy()

        # 4. Parallel Reduction
        self._demonstrate_parallel_reduction()

    def _demonstrate_memory_coalescing(self):
        """Show the impact of memory access patterns"""
        print("\n1. Memory Coalescing:")
        print("   - Contiguous access: Better bandwidth utilization")
        print("   - Strided access: Poor bandwidth utilization")

        batch, seq, dim = 32, 512, 768
        x = torch.randn(batch, seq, dim, device=self.device)

        # Good: Contiguous access
        start_time = time.perf_counter()
        for _ in range(100):
            y = x.view(-1, dim)  # Contiguous reshape
            torch.sum(y, dim=1)
        torch.cuda.synchronize() if self.device == "cuda" else None
        contiguous_time = time.perf_counter() - start_time

        # Bad: Strided access
        start_time = time.perf_counter()
        for _ in range(100):
            y = x[:, ::2, :]  # Strided access
            torch.sum(y, dim=1)
        torch.cuda.synchronize() if self.device == "cuda" else None
        strided_time = time.perf_counter() - start_time

        print(f"   Contiguous access: {contiguous_time*1000:.2f}ms")
        print(f"   Strided access:    {strided_time*1000:.2f}ms")
        print(f"   Speedup:          {strided_time/contiguous_time:.2f}x")

    def _demonstrate_kernel_fusion(self):
        """Show the benefits of operation fusion"""
        print("\n2. Kernel Fusion:")
        print("   - Separate kernels: Multiple memory round-trips")
        print("   - Fused kernels: Single memory round-trip")

        x = torch.randn(1024, 1024, device=self.device, requires_grad=True)

        # Unfused operations
        start_time = time.perf_counter()
        for _ in range(100):
            y = torch.relu(x)
            z = torch.sigmoid(y)
            torch.tanh(z)
        torch.cuda.synchronize() if self.device == "cuda" else None
        unfused_time = time.perf_counter() - start_time

        # Fused with torch.compile (if available)
        if hasattr(torch, 'compile'):
            @torch.compile
            def fused_ops(x):
                return torch.tanh(torch.sigmoid(torch.relu(x)))

            start_time = time.perf_counter()
            for _ in range(100):
                fused_ops(x)
            torch.cuda.synchronize() if self.device == "cuda" else None
            fused_time = time.perf_counter() - start_time

            print(f"   Unfused operations: {unfused_time*1000:.2f}ms")
            print(f"   Fused operations:   {fused_time*1000:.2f}ms")
            print(f"   Speedup:           {unfused_time/fused_time:.2f}x")
        else:
            print("   torch.compile not available for fusion demonstration")

    def _demonstrate_memory_hierarchy(self):
        """Show different memory access patterns"""
        print("\n3. Memory Hierarchy Usage:")
        print("   - Global memory: Slow, large capacity")
        print("   - Shared memory: Fast, limited capacity")
        print("   - Registers: Fastest, very limited")

        # This is conceptual - actual shared memory usage requires custom kernels
        sizes = [1024, 2048, 4096, 8192]
        for size in sizes:
            x = torch.randn(size, size, device=self.device)

            start_time = time.perf_counter()
            torch.matmul(x, x)  # This will use optimized BLAS
            torch.cuda.synchronize() if self.device == "cuda" else None
            time_taken = time.perf_counter() - start_time

            gflops = (2 * size**3) / (time_taken * 1e9)  # Matrix multiply FLOPs
            print(f"   Size {size}x{size}: {time_taken*1000:.2f}ms, {gflops:.1f} GFLOPS")

    def _demonstrate_parallel_reduction(self):
        """Show efficient reduction patterns"""
        print("\n4. Parallel Reduction:")
        print("   - Tree reduction: O(log n) steps")
        print("   - Sequential: O(n) steps")

        sizes = [1024, 4096, 16384, 65536]
        for size in sizes:
            x = torch.randn(size, device=self.device)

            # Parallel reduction (optimized by PyTorch/CUDA)
            start_time = time.perf_counter()
            for _ in range(1000):
                torch.sum(x)
            torch.cuda.synchronize() if self.device == "cuda" else None
            parallel_time = time.perf_counter() - start_time

            print(f"   Size {size}: {parallel_time*1000:.2f}ms")


class TransformerModel(nn.Module):
    """
    Transformer model that can use different optimization levels
    """
    def __init__(self, vocab_size: int, dim: int, num_layers: int,
                 num_heads: int, block_type: str = 'basic'):
        super().__init__()
        self.dim = dim
        self.embedding = nn.Embedding(vocab_size, dim)

        # Choose block type
        if block_type == 'basic':
            self.layers = nn.ModuleList([
                BasicTransformerBlock(dim, num_heads)
                for _ in range(num_layers)
            ])
        elif block_type == 'jit':
            self.layers = nn.ModuleList([
                FullyJITTransformerBlock(dim, num_heads)  # noqa: F821
                for _ in range(num_layers)
            ])
        elif block_type == 'triton' and TRITON_AVAILABLE:
            self.layers = nn.ModuleList([
                TritonOptimizedTransformerBlock(dim, num_heads)
                for _ in range(num_layers)
            ])
        else:
            raise ValueError(f"Unknown block type: {block_type}")

        self.norm = nn.LayerNorm(dim)
        self.output_proj = nn.Linear(dim, vocab_size, bias=False)

    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        x = self.embedding(input_ids)

        for layer in self.layers:
            x = layer(x)

        x = self.norm(x)
        return self.output_proj(x)


class CUDAOptimizedTransformer(nn.Module):
    """
    Transformer using custom CUDA kernels (simulated interface)
    """
    def __init__(self, vocab_size: int, dim: int, num_layers: int, num_heads: int):
        super().__init__()
        self.vocab_size = vocab_size
        self.dim = dim
        self.num_layers = num_layers
        self.num_heads = num_heads

        # Placeholder for CUDA-optimized implementation
        # In practice, this would use the custom CUDA kernels
        self.embedding = nn.Embedding(vocab_size, dim)
        self.layers = nn.ModuleList([
            BasicTransformerBlock(dim, num_heads)  # Fallback to basic
            for _ in range(num_layers)
        ])
        self.norm = nn.LayerNorm(dim)
        self.output_proj = nn.Linear(dim, vocab_size, bias=False)

    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        x = self.embedding(input_ids)

        # In a real implementation, this would call custom CUDA kernels
        for layer in self.layers:
            x = layer(x)

        x = self.norm(x)
        return self.output_proj(x)


def run_progressive_optimization_demo():
    """
    Main function to run the progressive optimization demonstration
    """
    print("Progressive Optimization Demonstration")
    print("="*60)

    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Using device: {device}")

    # Create demo instance
    demo = ProgressiveOptimizationDemo(device)

    # Model parameters
    batch_size = 4
    seq_len = 256
    dim = 512
    num_heads = 8
    vocab_size = 10000

    # Create input
    input_ids = torch.randint(0, vocab_size, (batch_size, seq_len), device=device)

    # Create models at different optimization levels
    print("\nCreating models with different optimization levels...")
    demo.create_models(dim, num_heads, vocab_size)
    print(f"Created {len(demo.models)} model variants")

    # Compare outputs for computational equivalence
    demo.compare_outputs(input_ids)

    # Benchmark performance
    print(f"\nBenchmarking with input shape: {input_ids.shape}")
    demo.benchmark_models(input_ids, num_iterations=50)

    # Print results
    demo.print_results()

    # Demonstrate kernel concepts
    demo.demonstrate_kernel_concepts()

    print("\nDemo complete! Key takeaways:")
    print("1. All optimization levels produce computationally equivalent results")
    print("2. Progressive optimization improves performance while maintaining correctness")
    print("3. Understanding kernel patterns helps design better neural network models")
    print("4. Memory access patterns significantly impact performance")


if __name__ == "__main__":
    run_progressive_optimization_demo()
