"""
Model optimization commands for TorchBridge CLI.

Provides easy-to-use cross-backend preparation of PyTorch models with different
optimization levels and hardware configurations.
"""

import argparse
import logging
import sys
import time
from pathlib import Path

import torch

import torchbridge
from torchbridge.utils.compiler_assistant import CompilerOptimizationAssistant

logger = logging.getLogger(__name__)


class OptimizeCommand:
    """Model optimization command implementation for cross-backend deployment."""

    @staticmethod
    def register(subparsers) -> None:
        """Register the optimize command with argument parser."""
        parser = subparsers.add_parser(
            'optimize',
            help='Optimize PyTorch models for production deployment',
            description='Apply TorchBridge backend-aware transformations to your PyTorch models',
            formatter_class=argparse.RawDescriptionHelpFormatter,
            epilog="""
Optimization Levels:
  basic      - PyTorch native optimizations (cuDNN, cuBLAS)
  jit        - TorchScript JIT compilation
  compile    - torch.compile with Inductor backend
  triton     - Triton kernel optimizations (GPU required)
  production - Full optimization stack for deployment

Examples:
  tb-optimize --model model.pt --level production --output optimized_model.pt
  tb-optimize --model resnet50 --level compile --benchmark --verbose
  tb-optimize --input-shape 1,3,224,224 --level triton --hardware auto
            """
        )

        parser.add_argument(
            '--model',
            type=str,
            required=True,
            help='Path to model file (.pt/.pth) or HuggingFace model name'
        )

        parser.add_argument(
            '--level',
            choices=['basic', 'jit', 'compile', 'triton', 'production'],
            default='compile',
            help='Optimization level (default: compile)'
        )

        parser.add_argument(
            '--output', '-o',
            type=str,
            help='Output path for optimized model (default: auto-generated)'
        )

        parser.add_argument(
            '--input-shape',
            type=str,
            help='Input tensor shape (e.g., "1,3,224,224" for batch,channels,height,width)'
        )

        parser.add_argument(
            '--hardware',
            choices=['auto', 'cpu', 'cuda', 'mps'],
            default='auto',
            help='Target hardware (default: auto-detect)'
        )

        parser.add_argument(
            '--benchmark',
            action='store_true',
            help='Run performance benchmark after optimization'
        )

        parser.add_argument(
            '--validate',
            action='store_true',
            help='Validate optimization correctness'
        )

        parser.add_argument(
            '--verbose', '-v',
            action='store_true',
            help='Enable verbose output'
        )

        parser.add_argument(
            '--trust-source',
            action='store_true',
            help='Allow loading untrusted model files (enables weights_only=False for pickle deserialization)'
        )

    @staticmethod
    def execute(args) -> int:
        """Execute the optimize command."""
        print(" TorchBridge Model Preparation")
        print("=" * 50)

        try:
            # Detect hardware
            device = OptimizeCommand._detect_hardware(args.hardware, args.verbose)

            # Load model
            trust_source = getattr(args, 'trust_source', False)
            model = OptimizeCommand._load_model(args.model, args.verbose, trust_source=trust_source)
            model = model.to(device)

            # Parse input shape
            input_shape = OptimizeCommand._parse_input_shape(args.input_shape, model)

            # Create sample input
            sample_input = torch.randn(input_shape, device=device)

            # Apply optimizations
            optimized_model = OptimizeCommand._apply_optimizations(
                model, args.level, sample_input, args.verbose
            )

            # Validate if requested
            if args.validate:
                OptimizeCommand._validate_optimization(model, optimized_model, sample_input, args.verbose)

            # Benchmark if requested
            if args.benchmark:
                OptimizeCommand._benchmark_models(model, optimized_model, sample_input, args.verbose)

            # Save optimized model
            if args.output:
                OptimizeCommand._save_model(optimized_model, args.output, args.verbose)

            print("\n Optimization completed successfully!")
            return 0

        except Exception as e:
            print(f" Optimization failed: {e}")
            if args.verbose:
                import traceback
                traceback.print_exc()
            return 1

    @staticmethod
    def _detect_hardware(hardware: str, verbose: bool) -> torch.device:
        """Detect target hardware device."""
        if hardware == 'auto':
            if torch.cuda.is_available():
                device = torch.device('cuda')
            elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
                device = torch.device('mps')
            else:
                device = torch.device('cpu')
        else:
            device = torch.device(hardware)

        if verbose:
            print(f"  Target device: {device}")
            if device.type == 'cuda':
                print(f"   GPU: {torch.cuda.get_device_name()}")
                print(f"   Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")

        return device

    @staticmethod
    def _load_model(model_path: str, verbose: bool, *, trust_source: bool = False) -> torch.nn.Module:
        """Load model from file or create example model."""
        if verbose:
            print(f" Loading model: {model_path}")

        # Check if it's a file path
        if Path(model_path).exists():
            try:
                model = torch.load(model_path, map_location='cpu', weights_only=not trust_source)
            except Exception as e:
                if not trust_source and "Weights only load failed" in str(e):
                    raise RuntimeError(
                        f"Cannot safely load '{model_path}': the file contains objects that "
                        f"require pickle deserialization.\n"
                        f"If you trust this file, re-run with --trust-source to allow loading."
                    ) from e
                raise
            if isinstance(model, dict) and 'model' in model:
                model = model['model']
        else:
            # Create example models for common names
            if model_path.lower() == 'resnet50':
                import torchvision.models as models
                model = models.resnet50(pretrained=False)
            elif model_path.lower() == 'transformer':
                # Simple transformer-like model for demonstration
                model = torch.nn.Sequential(
                    torch.nn.Linear(768, 3072),
                    torch.nn.GELU(),
                    torch.nn.Linear(3072, 768),
                    torch.nn.LayerNorm(768)
                )
            else:
                raise ValueError(f"Model not found: {model_path}")

        model.eval()
        if verbose:
            param_count = sum(p.numel() for p in model.parameters())
            print(f"   Parameters: {param_count:,}")

        return model

    @staticmethod
    def _parse_input_shape(input_shape_str: str | None, model: torch.nn.Module) -> tuple:
        """Parse input shape string or infer from model."""
        if input_shape_str:
            shape = tuple(map(int, input_shape_str.split(',')))
            if any(d <= 0 for d in shape):
                raise ValueError(f"All input shape dimensions must be positive, got {shape}")
        else:
            # Try to infer shape from model
            if hasattr(model, 'forward'):
                # Default shapes for common architectures
                shape = (1, 3, 224, 224)  # Batch, Channels, Height, Width
            else:
                shape = (1, 768)  # Default for transformer-like models

        return shape

    @staticmethod
    def _apply_optimizations(model: torch.nn.Module, level: str, sample_input: torch.Tensor, verbose: bool) -> torch.nn.Module:
        """Apply optimizations based on level."""
        if verbose:
            print(f"  Applying '{level}' optimizations...")

        optimized_model = model

        if level == 'basic':
            # Basic PyTorch optimizations
            optimized_model = model.eval()

        elif level == 'jit':
            # TorchScript JIT
            optimized_model = torch.jit.trace(model, sample_input)  # type: ignore[assignment]

        elif level == 'compile':
            # torch.compile optimization
            optimized_model = torch.compile(model, mode='max-autotune')  # type: ignore[assignment]

        elif level == 'triton':
            # Use TorchBridge Triton backend kernels
            if hasattr(torchbridge, 'OptimizedMultiHeadAttention'):
                # Apply attention optimizations if applicable
                optimized_model = model
            else:
                optimized_model = torch.compile(model, mode='max-autotune')  # type: ignore[assignment]

        elif level == 'production':
            # Full production optimization stack
            try:
                # Use TorchBridge HAL-aware compilation assistant
                assistant = CompilerOptimizationAssistant(device=sample_input.device)
                result = assistant.optimize_model(model, interactive=False)
                optimized_model = torch.compile(model, mode='max-autotune')  # type: ignore[assignment]
                if verbose and result.optimization_opportunities:
                    print(f"   Found {len(result.optimization_opportunities)} optimization opportunities")
            except Exception:
                logger.debug("HAL-aware compilation assistant failed, falling back to torch.compile", exc_info=True)
                # Fallback to torch.compile
                optimized_model = torch.compile(model, mode='max-autotune')  # type: ignore[assignment]

        if verbose:
            print("    Optimizations applied")

        return optimized_model

    @staticmethod
    def _validate_optimization(original: torch.nn.Module, optimized: torch.nn.Module,
                             sample_input: torch.Tensor, verbose: bool) -> None:
        """Validate that optimization preserves correctness."""
        if verbose:
            print(" Validating optimization correctness...")

        with torch.no_grad():
            original_output = original(sample_input)
            optimized_output = optimized(sample_input)

            # Check output shapes match
            if original_output.shape != optimized_output.shape:
                raise ValueError(f"Output shape mismatch: {original_output.shape} vs {optimized_output.shape}")

            # Check numerical accuracy
            max_diff = torch.max(torch.abs(original_output - optimized_output)).item()
            relative_error = max_diff / torch.max(torch.abs(original_output)).item()

            if verbose:
                print(f"   Max difference: {max_diff:.2e}")
                print(f"   Relative error: {relative_error:.2e}")

            if relative_error > 0.01:  # 1% tolerance
                print(f"  Warning: High relative error ({relative_error:.2e})")
            else:
                print("    Validation passed")

    @staticmethod
    def _benchmark_models(original: torch.nn.Module, optimized: torch.nn.Module,
                         sample_input: torch.Tensor, verbose: bool) -> None:
        """Benchmark original vs optimized model performance."""
        if verbose:
            print(" Benchmarking performance...")

        def benchmark_model(model, name: str, warmup: int = 10, runs: int = 100):
            model.eval()
            with torch.no_grad():
                # Warmup
                for _ in range(warmup):
                    _ = model(sample_input)

                if sample_input.device.type == 'cuda':
                    torch.cuda.synchronize()

                # Benchmark
                start_time = time.time()
                for _ in range(runs):
                    _ = model(sample_input)

                if sample_input.device.type == 'cuda':
                    torch.cuda.synchronize()

                end_time = time.time()
                avg_time = (end_time - start_time) / runs * 1000  # ms

                if verbose:
                    print(f"   {name}: {avg_time:.2f} ms/inference")

                return avg_time

        original_time = benchmark_model(original, "Original")
        optimized_time = benchmark_model(optimized, "Optimized")

        if optimized_time > 0:
            speedup = original_time / optimized_time
            print(f"    Speedup: {speedup:.2f}x")
        else:
            print("    Speedup: N/A (optimized time too small to measure)")

    @staticmethod
    def _save_model(model: torch.nn.Module, output_path: str, verbose: bool) -> None:
        """Save optimized model to file."""
        if verbose:
            print(f" Saving optimized model to: {output_path}")

        # Create output directory if needed
        Path(output_path).parent.mkdir(parents=True, exist_ok=True)

        # Save model
        torch.save(model, output_path)

        if verbose:
            file_size = Path(output_path).stat().st_size / 1e6  # MB
            print(f"   File size: {file_size:.1f} MB")


def main():
    """Standalone entry point for tb-optimize."""
    parser = argparse.ArgumentParser(
        prog='tb-optimize',
        description='Apply TorchBridge backend-aware transformations to your PyTorch models',
        formatter_class=argparse.RawDescriptionHelpFormatter
    )

    # Add optimize command arguments directly
    parser.add_argument(
        '--model',
        type=str,
        required=True,
        help='Path to model file (.pt/.pth) or HuggingFace model name'
    )
    parser.add_argument(
        '--level',
        choices=['basic', 'jit', 'compile', 'triton', 'production'],
        default='compile',
        help='Optimization level (default: compile)'
    )
    parser.add_argument(
        '--output', '-o',
        type=str,
        help='Output path for optimized model (default: auto-generated)'
    )
    parser.add_argument(
        '--input-shape',
        type=str,
        help='Input tensor shape (e.g., "1,3,224,224")'
    )
    parser.add_argument(
        '--hardware',
        choices=['auto', 'cpu', 'cuda', 'mps'],
        default='auto',
        help='Target hardware (default: auto-detect)'
    )
    parser.add_argument(
        '--benchmark',
        action='store_true',
        help='Run performance benchmark after optimization'
    )
    parser.add_argument(
        '--validate',
        action='store_true',
        help='Validate optimization correctness'
    )
    parser.add_argument(
        '--verbose', '-v',
        action='store_true',
        help='Enable verbose output'
    )
    parser.add_argument(
        '--trust-source',
        action='store_true',
        help='Allow loading untrusted model files (enables weights_only=False for pickle deserialization)'
    )

    args = parser.parse_args()
    try:
        return OptimizeCommand.execute(args)
    except KeyboardInterrupt:
        print("\nOperation cancelled by user")
        return 130
    except Exception as e:
        from torchbridge.cli import _print_error
        return _print_error(e, verbose=getattr(args, 'verbose', False))


if __name__ == '__main__':
    sys.exit(main())
