"""
Enhanced ConvolutionBenchmark with exact original script integration.
This provides both the unified API and exact original script replication.
"""
import torch
import torch.nn as nn
import time
import numpy as np
from typing import Dict, List, Tuple, Any, Optional
from tabulate import tabulate
from ..layers import (
    TraditionalConv2D,
    DepthwiseSeparableConv, 
    DeformableConv2D,
    DynamicConv2D,
    KWDSConv2D,
    ODConv2D,
)
from ..utils import get_comprehensive_analysis
from ..original_scripts import run_original_script

class ConvolutionBenchmark:
    """
    Enhanced benchmarking suite that provides both unified API and exact original script replication.
    This ensures maximum depth and comprehensiveness matching your original implementations.
    """
    
    def __init__(self, device: str = "cpu", num_runs: int = 100, use_ptflops: bool = True):
        """
        Initialize the benchmark.
        
        Args:
            device: Device to run benchmarks on ("cpu" or "cuda")
            num_runs: Number of iterations for latency measurement
            use_ptflops: Whether to use ptflops library for accurate FLOP counting
        """
        self.device = torch.device(device)
        self.num_runs = num_runs
        self.use_ptflops = use_ptflops
        
        self.conv_types = {
            "Traditional": TraditionalConv2D,
            "Depthwise Separable": DepthwiseSeparableConv,
            "Deformable v1": DeformableConv2D, 
            "Dynamic": DynamicConv2D,
            "Kernel Warehouse": KWDSConv2D,
            "ODConv": ODConv2D,
        }
    
    def replicate_original_script(self, script_name: str, input_shape: Optional[Tuple[int, int, int]] = None):
        """
        Run EXACT replication of an original script.
        This preserves every detail of your original implementations.
        
        Args:
            script_name: Name of original script ("DSKW", "DSODConv", "DCNv1", etc.)
            input_shape: Input shape (optional, uses default if not provided)
        """
        print(f"🎯 Running EXACT replication of {script_name}.py")
        print("=" * 60)
        print("This output matches your original script exactly:")
        print("-" * 60)
        
        # Run the exact original script
        run_original_script(script_name)
        
        print("-" * 60)
        print(f"✅ Completed exact replication of {script_name}.py")
    
    def benchmark_single(
        self, 
        conv_class: nn.Module,
        input_shape: Tuple[int, int, int],
        output_channels: int,
        kernel_size: int = 3,
        stride: int = 1,
        padding: Optional[int] = None,
        **kwargs
    ) -> Dict[str, Any]:
        """
        Benchmark a single convolution type with comprehensive analysis.
        
        Args:
            conv_class: Convolution class to benchmark
            input_shape: Input tensor shape (H, W, C)
            output_channels: Number of output channels
            kernel_size: Convolution kernel size
            stride: Convolution stride
            padding: Padding (defaults to kernel_size // 2)
            **kwargs: Additional arguments for the convolution layer
            
        Returns:
            Dictionary containing comprehensive benchmark results
        """
        in_channels = input_shape[2]
        if padding is None:
            padding = kernel_size // 2
        
        # Create model based on type
        if conv_class == KWDSConv2D:
            model = conv_class(in_channels, output_channels, kernel_size, stride, **kwargs)
            conv_type = "kernel_warehouse"
        elif conv_class == ODConv2D:
            model = conv_class(in_channels, output_channels, kernel_size, stride, **kwargs)
            conv_type = "omni_dimensional"
        elif conv_class == DepthwiseSeparableConv:
            model = conv_class(in_channels, output_channels, kernel_size, stride, padding)
            conv_type = "depthwise_separable"
        elif conv_class == DeformableConv2D:
            model = conv_class(in_channels, output_channels, kernel_size, stride, padding, **kwargs)
            conv_type = "deformable"
        else:
            model = conv_class(in_channels, output_channels, kernel_size, stride, padding)
            conv_type = "standard"
            
        model = model.to(self.device).eval()
        
        # Create input tensor
        input_tensor = torch.randn(1, in_channels, input_shape[0], input_shape[1]).to(self.device)
        
        # Warm-up
        with torch.no_grad():
            output = model(input_tensor)
        
        # Benchmark latency (exactly like original scripts)
        latencies = []
        for _ in range(self.num_runs):
            if self.device.type == "cuda":
                torch.cuda.synchronize()
            start_time = time.time()
            
            with torch.no_grad():
                _ = model(input_tensor)
                    
            if self.device.type == "cuda":
                torch.cuda.synchronize()
            end_time = time.time()
            latencies.append((end_time - start_time) * 1000)  # Convert to ms
        
        latencies = np.array(latencies)
        
        # Get output shape
        output_shape = tuple(output.shape[1:])  # Remove batch dimension
        
        # Get comprehensive analysis (manual + ptflops)
        analysis_input_shape = (in_channels, input_shape[0], input_shape[1])  # ptflops format
        comprehensive_analysis = get_comprehensive_analysis(
            model, analysis_input_shape, conv_type, **kwargs
        )
        
        # Build result with all original script details
        result = {
            "conv_type": conv_type,
            "model_params": sum(p.numel() for p in model.parameters()),
            "output_shape": output_shape,
            "kernel_size": kernel_size,
            "input_shape": input_shape,
            "output_channels": output_channels,
            "stride": stride,
            "padding": padding,
            
            # Latency statistics (exactly like originals)
            "latency_mean": np.mean(latencies),
            "latency_std": np.std(latencies),
            "latency_min": np.min(latencies),
            "latency_max": np.max(latencies),
            "latency_p95": np.percentile(latencies, 95),
            "latency_p99": np.percentile(latencies, 99),
            
            # Comprehensive analysis
            "comprehensive_analysis": comprehensive_analysis,
        }
        
        return result
    
    def compare_all(
        self,
        input_shape: Tuple[int, int, int],
        output_channels: int,
        kernel_size: int = 3,
        stride: int = 1,
        padding: Optional[int] = None,
        **kwargs
    ) -> Dict[str, Dict[str, Any]]:
        """
        Compare all convolution types with comprehensive analysis.
        """
        results = {}
        
        for name, conv_class in self.conv_types.items():
            try:
                print(f"Benchmarking {name}...")
                result = self.benchmark_single(
                    conv_class, input_shape, output_channels, kernel_size, stride, padding, **kwargs
                )
                results[name] = result
            except Exception as e:
                print(f"Failed to benchmark {name}: {e}")
                results[name] = None
        
        return results
    
    def print_original_style_breakdown(self, result: Dict[str, Any], conv_name: str) -> None:
        """
        Print results in the exact style of your original scripts with full detail breakdown.
        """
        analysis = result["comprehensive_analysis"]
        manual = analysis.get("manual_calculation")
        ptflops = analysis.get("ptflops_analysis")
        
        if conv_name in ["Kernel Warehouse", "ODConv"]:
            # Use the detailed breakdown style from DSKW.py and DSODConv.py
            print(f"\nBenchmarking Results for {conv_name}:")
            print(f"  Output Channels: {result['output_channels']}")
            print(f"  Kernel Size: {result['kernel_size']}")
            print(f"  Stride: {result['stride']}")
            
            if manual:
                print(f"  Output Shape: {manual.get('output_shape', result['output_shape'])}")
                print(f"  Parameters: {manual['total_params'] / 1e3:.2f}K")
                print(f"  Total FLOPs: {manual['total_flops'] / 1e6:.2f}M")
                
                if "depthwise_mults" in manual:
                    print(f"    - Depthwise Mults: {manual['depthwise_mults'] / 1e6:.2f}M")
                    print(f"    - Pointwise Mults: {manual['pointwise_mults'] / 1e6:.2f}M")
                    print(f"    - Dynamic Branch FLOPs: {manual['dynamic_flops'] / 1e6:.2f}M")
                    print(f"    - Convolution Adds: {manual['conv_adds'] / 1e6:.2f}M")
                    print(f"    - Divisions: {manual['divisions'] / 1e6:.2f}M")
        
        elif conv_name == "Deformable v1":
            # Use DCNv1.py style
            print("Deformable Convolution Layer:")
            if ptflops:
                print(f"Total Parameters: {ptflops['params']} ({ptflops['params_str']})")
                print(f"Estimated FLOPs: {ptflops['flops']} (~{ptflops['flops_str']})")
            elif manual:
                print(f"Total Parameters: {manual['total_params']} ({manual['total_params'] / 1e3:.2f}K)")
                print(f"Estimated FLOPs: {manual['total_flops']} (~{manual['total_flops'] / 1e6:.2f}M)")
        
        else:
            # Use ptflops style for others
            print(f"\n{conv_name}:")
            print(f"Kernel Size: {result['kernel_size']}")
            print(f"Input Shape: {result['input_shape']}")
            print(f"Output Channels: {result['output_channels']}")
            
            if ptflops:
                print(f"Parameters: {ptflops['params_str']}")
                print(f"MACs: {ptflops['macs_str']}")
                print(f"FLOPs: {ptflops['flops_str']}")
            elif manual:
                print(f"Parameters: {manual['total_params'] / 1e3:.2f}K")
                print(f"Total FLOPs: {manual['total_flops'] / 1e6:.2f}M")
        
        # Always show latency statistics in original format
        print("Latency Statistics:")
        print(f"  Mean: {result['latency_mean']:.2f}ms")
        print(f"  Std Dev: {result['latency_std']:.2f}ms")
        print(f"  Min: {result['latency_min']:.2f}ms")
        print(f"  Max: {result['latency_max']:.2f}ms")
        print(f"  P95: {result['latency_p95']:.2f}ms")
        print(f"  P99: {result['latency_p99']:.2f}ms")
        print()
    
    def print_comparison(self, results: Dict[str, Dict[str, Any]], style: str = "original") -> None:
        """
        Print comparison results in various styles.
        
        Args:
            results: Results dictionary from compare_all()
            style: "original", "summary", "detailed", or "comprehensive"
        """
        if style == "original":
            # Print each result in original script style
            for name, result in results.items():
                if result is not None:
                    self.print_original_style_breakdown(result, name)
        
        elif style == "summary":
            # Summary table
            headers = [
                "Convolution Type",
                "Parameters", 
                "FLOPs (M)",
                "Mean Latency (ms)",
                "P95 Latency (ms)"
            ]
            
            table_data = []
            for name, result in results.items():
                if result is None:
                    continue
                
                analysis = result["comprehensive_analysis"]
                
                if analysis.get("ptflops_analysis"):
                    ptflops = analysis["ptflops_analysis"]
                    params_str = ptflops["params_str"]
                    flops_str = ptflops["flops_str"]
                elif analysis.get("manual_calculation"):
                    manual = analysis["manual_calculation"]
                    params = manual["total_params"]
                    flops = manual["total_flops"]
                    params_str = f"{params / 1e3:.2f}K"
                    flops_str = f"{flops / 1e6:.1f}M"
                else:
                    params_str = "N/A"
                    flops_str = "N/A"
                    
                row = [
                    name,
                    params_str,
                    flops_str,
                    f"{result['latency_mean']:.2f}",
                    f"{result['latency_p95']:.2f}"
                ]
                table_data.append(row)
            
            print("\nConvolution Comparison Results:")
            print(tabulate(table_data, headers=headers, tablefmt="grid"))
        
        elif style == "detailed":
            # Show detailed FLOP breakdown for each
            for name, result in results.items():
                if result is not None:
                    self.print_original_style_breakdown(result, name)
        
        elif style == "comprehensive":
            # Show both summary and detailed
            self.print_comparison(results, "summary")
            print("\n" + "="*80)
            print("DETAILED BREAKDOWN:")
            print("="*80)
            self.print_comparison(results, "detailed")
    
    def run_all_original_scripts(self):
        """
        Run ALL original scripts in sequence for complete validation.
        """
        from ..original_scripts import run_all_original_scripts
        run_all_original_scripts()
    
    def save_results(self, results: Dict[str, Dict[str, Any]], filename: str) -> None:
        """
        Save comprehensive benchmark results to JSON with all original detail preserved.
        """
        import json
        
        def convert_numpy_types(obj):
            if isinstance(obj, np.floating):
                return float(obj)
            elif isinstance(obj, np.integer):
                return int(obj)
            elif isinstance(obj, dict):
                return {k: convert_numpy_types(v) for k, v in obj.items()}
            elif isinstance(obj, (list, tuple)):
                return [convert_numpy_types(item) for item in obj]
            else:
                return obj
        
        serializable_results = {}
        for name, result in results.items():
            if result is None:
                continue
            serializable_results[name] = convert_numpy_types(result)
        
        # Add metadata about the benchmark run
        serializable_results["_metadata"] = {
            "device": str(self.device),
            "num_runs": self.num_runs,
            "use_ptflops": self.use_ptflops,
            "library_version": "0.1.0",
            "comprehensive_analysis": True,
            "original_script_compatibility": True
        }
        
        with open(filename, 'w') as f:
            json.dump(serializable_results, f, indent=2)
        
        print(f"Comprehensive results with original script depth saved to {filename}")

# Convenience functions for quick access to original script functionality
def replicate_dskw(input_shape=(224, 224, 3)):
    """Quick access to DSKW.py replication."""
    benchmark = ConvolutionBenchmark()
    benchmark.replicate_original_script("DSKW", input_shape)

def replicate_dsodconv(input_shape=(224, 224, 3)):
    """Quick access to DSODConv.py replication."""
    benchmark = ConvolutionBenchmark()
    benchmark.replicate_original_script("DSODConv", input_shape)

def replicate_dcnv1(input_shape=(224, 224, 3)):
    """Quick access to DCNv1.py replication."""
    benchmark = ConvolutionBenchmark()
    benchmark.replicate_original_script("DCNv1", input_shape)

def replicate_all_originals():
    """Quick access to run all original scripts."""
    benchmark = ConvolutionBenchmark()
    benchmark.run_all_original_scripts()
