# python embedd_micro.py --use_int8 Fastest

import argparse
import time
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple

import numpy as np
import torch
from torch import nn
from transformers import AutoModel, BitsAndBytesConfig
from tqdm import tqdm
from contextlib import contextmanager

@dataclass
class BenchmarkConfig:
    model_path: str
    batch_sizes: List[int]
    seq_length: int
    num_runs: int
    use_fp16: bool = True
    use_int4: bool = False
    use_int8: bool = False  # Add this parameter
    use_cuda_graphs: bool = False
    use_flash_attention: bool = False
    use_linear8bitlt: bool = False


class GraphContainer:
    """Container for managing graphs for different batch sizes (CUDA graphs on NVIDIA, regular on others)."""
    
    def __init__(self, model: nn.Module, seq_length: int):
        self.model = model
        self.seq_length = seq_length
        self.graphs: Dict[int, 'GraphWrapper'] = {}
    
    def get_or_create(self, batch_size: int) -> 'GraphWrapper':
        if batch_size not in self.graphs:
            self.graphs[batch_size] = GraphWrapper(
                self.model, batch_size, self.seq_length
            )
        return self.graphs[batch_size]


class GraphWrapper:
    """Wrapper for graph capture and replay (CUDA graphs on NVIDIA, regular on others)."""
    
    def __init__(self, model: nn.Module, batch_size: int, seq_length: int):
        self.model = model
        self.device = self._get_device()
        self.static_input = self._create_random_batch(batch_size, seq_length)
        self.static_attention_mask = torch.ones_like(self.static_input)
        
        # Warm up
        self._warmup()
        
        # Only use CUDA graphs on NVIDIA GPUs
        if torch.cuda.is_available() and hasattr(torch.cuda, 'CUDAGraph'):
            # Capture graph
            self.graph = torch.cuda.CUDAGraph()
            with torch.cuda.graph(self.graph):
                self.static_output = self.model(
                    input_ids=self.static_input,
                    attention_mask=self.static_attention_mask
                )
            self.use_cuda_graph = True
        else:
            # For MPS or CPU, just store the model
            self.use_cuda_graph = False
            self.static_output = None
    
    def _get_device(self) -> str:
        if torch.cuda.is_available():
            return "cuda"
        elif torch.backends.mps.is_available():
            return "mps"
        else:
            return "cpu"
    
    def _create_random_batch(self, batch_size: int, seq_length: int) -> torch.Tensor:
        return torch.randint(
            0, 1000, (batch_size, seq_length), 
            device=self.device, 
            dtype=torch.long
        )
    
    def _warmup(self, num_warmup: int = 3):
        with torch.no_grad():
            for _ in range(num_warmup):
                self.model(
                    input_ids=self.static_input,
                    attention_mask=self.static_attention_mask
                )
    
    def __call__(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
        if self.use_cuda_graph:
            self.static_input.copy_(input_ids)
            self.static_attention_mask.copy_(attention_mask)
            self.graph.replay()
            return self.static_output
        else:
            # For MPS/CPU, just run normally
            return self.model(input_ids=input_ids, attention_mask=attention_mask)


class ModelOptimizer:
    """Applies various optimizations to the model."""
    
    @staticmethod
    def optimize(model: nn.Module, config: BenchmarkConfig) -> nn.Module:
        print("\nApplying model optimizations:")
        
        if model is None:
            raise ValueError("Cannot optimize None model")
        
        # Move to GPU
        if torch.cuda.is_available():
            model = model.cuda()
            device = "cuda"
        elif torch.backends.mps.is_available():
            model = model.to("mps")
            device = "mps"
        else:
            model = model.cpu()
            device = "cpu"
        print(f"- Model moved to {device}")
        
        # FP16
        if config.use_fp16 and not config.use_int4:
            model = model.half()
            # use torch compile
            model = torch.compile(model)
            print("- Using FP16 precision")
        
        # Check if using SDPA (only on CUDA)
        if torch.cuda.is_available() and torch.version.cuda and float(torch.version.cuda[:3]) >= 11.6:
            if hasattr(torch.nn.functional, 'scaled_dot_product_attention'):
                print("- Using PyTorch SDPA (scaled_dot_product_attention)")
            else:
                print("- PyTorch SDPA not available")
        
        # Flash Attention (only on CUDA)
        if config.use_flash_attention and torch.cuda.is_available():
            try:
                from flash_attn.flash_attention import FlashAttention
                print("- Flash Attention 2 available")
                if hasattr(model.config, "attention_mode"):
                    model.config.attention_mode = "flash_attention_2"
                    print("  - Enabled Flash Attention 2 mode")
            except ImportError:
                print("- Flash Attention not available")
        
        # Memory efficient attention (only on CUDA)
        if torch.cuda.is_available():
            try:
                from xformers.ops import memory_efficient_attention
                if hasattr(model, 'enable_xformers_memory_efficient_attention'):
                    model.enable_xformers_memory_efficient_attention()
                    print("- Enabled xformers memory efficient attention")
                else:
                    print("- Model doesn't support xformers")
            except (ImportError, AttributeError):
                print("- Xformers not available")
        
        model.eval()
        print("- Model set to eval mode")
        
        return model


class Timer:
    """Handles accurate GPU timing using GPU events or CPU timing."""
    
    def __init__(self):
        if torch.cuda.is_available():
            self.start_event = torch.cuda.Event(enable_timing=True)
            self.end_event = torch.cuda.Event(enable_timing=True)
            self.use_gpu_timing = True
        elif torch.backends.mps.is_available():
            # MPS doesn't have events, use CPU timing
            self.use_gpu_timing = False
        else:
            # CPU timing
            self.use_gpu_timing = False
    
    @contextmanager
    def timing(self):
        if self.use_gpu_timing:
            self.start_event.record()
            yield
            self.end_event.record()
            self.end_event.synchronize()
        else:
            # Use CPU timing for MPS/CPU
            start_time = time.time()
            yield
            self.cpu_elapsed = time.time() - start_time
    
    def elapsed_time(self) -> float:
        if self.use_gpu_timing:
            return self.start_event.elapsed_time(self.end_event) / 1000  # ms to seconds
        else:
            return self.cpu_elapsed


class Benchmark:
    """Main benchmark runner."""
    
    def __init__(self, config: BenchmarkConfig):
        self.config = config
        try:
            self.model = self._load_model()
            if self.model is None:
                raise ValueError("Model initialization failed - model is None")
            
            # Only use CUDA graphs on NVIDIA GPUs
            if config.use_cuda_graphs and torch.cuda.is_available():
                self.graphs = GraphContainer(self.model, config.seq_length)
            else:
                self.graphs = None
            self.timer = Timer()
        except Exception as e:
            print(f"ERROR in benchmark initialization: {str(e)}")
            raise
    
    def _load_model(self) -> nn.Module:
        print(f"Loading model from {self.config.model_path}...")
        
        try:
            # Int4 quantization using HuggingFace integration
            if self.config.use_int4:
                import bitsandbytes as bnb
                print(f"- bitsandbytes version: {bnb.__version__}")
                
                # 检查是否使用自定义的8bit量化
                if hasattr(self.config, 'use_linear8bitlt') and self.config.use_linear8bitlt:
                    print("- Using custom Linear8bitLt replacement for all linear layers")
                    
                    # 加载原始模型（不使用量化配置）
                    import bitsandbytes as bnb
                    import torch
                    # set default to half
                    torch.set_default_dtype(torch.float16)
                    compute_dtype = torch.float16 if self.config.use_fp16 else torch.float32
                    model = AutoModel.from_pretrained(
                        self.config.model_path,
                        torch_dtype=compute_dtype,
                    )
                    
                    # 定义替换函数
                    def replace_linear_with_linear8bitlt(model):
                        """递归地将模型中的所有nn.Linear层替换为Linear8bitLt"""
                        for name, module in list(model.named_children()):
                            if isinstance(module, nn.Linear):
                                # 获取原始线性层的参数
                                in_features = module.in_features
                                out_features = module.out_features
                                bias = module.bias is not None
                                
                                # 创建8bit线性层
                                # print size
                                print(f"in_features: {in_features}, out_features: {out_features}")
                                new_module = bnb.nn.Linear8bitLt(
                                    in_features, 
                                    out_features, 
                                    bias=bias, 
                                    has_fp16_weights=False
                                )
                                
                                # 复制权重和偏置
                                new_module.weight.data = module.weight.data
                                if bias:
                                    new_module.bias.data = module.bias.data
                                    
                                # 替换模块
                                setattr(model, name, new_module)
                            else:
                                # 递归处理子模块
                                replace_linear_with_linear8bitlt(module)
                        
                        return model
                    
                    # 替换所有线性层
                    model = replace_linear_with_linear8bitlt(model)
                    # add torch compile
                    model = torch.compile(model)
                    
                    # 将模型移到GPU（量化发生在这里）
                    device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
                    model = model.to(device)
                    
                    print("- All linear layers replaced with Linear8bitLt")
                    
                else:
                    # 使用原来的Int4量化方法
                    print("- Using bitsandbytes for Int4 quantization")
                    
                    # Create quantization config
                    
                    compute_dtype = torch.float16 if self.config.use_fp16 else torch.float32
                    quantization_config = BitsAndBytesConfig(
                        load_in_4bit=True,
                        bnb_4bit_compute_dtype=compute_dtype,
                        bnb_4bit_use_double_quant=True,
                        bnb_4bit_quant_type="nf4"
                    )
                    
                    print("- Quantization config:", quantization_config)
                    
                    # Load model directly with quantization config
                    model = AutoModel.from_pretrained(
                        self.config.model_path,
                        quantization_config=quantization_config,
                        torch_dtype=compute_dtype,
                        device_map="auto"  # Let HF decide on device mapping
                    )
                
                # Check if model loaded successfully
                if model is None:
                    raise ValueError("Model loading returned None")
                    
                print(f"- Model type: {type(model)}")
                
                # Apply optimizations directly here
                print("\nApplying model optimizations:")
                
                if hasattr(self.config, 'use_linear8bitlt') and self.config.use_linear8bitlt:
                    print("- Model moved to GPU with Linear8bitLt quantization")
                else:
                    # Skip moving to GPU since device_map="auto" already did that
                    print("- Model already on GPU due to device_map='auto'")
                
                # Skip FP16 conversion since we specified compute_dtype
                print(f"- Using {compute_dtype} for compute dtype")
                
                # Check CUDA and SDPA
                if torch.cuda.is_available() and torch.version.cuda and float(torch.version.cuda[:3]) >= 11.6:
                    if hasattr(torch.nn.functional, 'scaled_dot_product_attention'):
                        print("- Using PyTorch SDPA (scaled_dot_product_attention)")
                    else:
                        print("- PyTorch SDPA not available")
                
                # Try xformers if available (only on CUDA)
                if torch.cuda.is_available():
                    try:
                        from xformers.ops import memory_efficient_attention
                        if hasattr(model, 'enable_xformers_memory_efficient_attention'):
                            model.enable_xformers_memory_efficient_attention()
                            print("- Enabled xformers memory efficient attention")
                        else:
                            print("- Model doesn't support xformers")
                    except (ImportError, AttributeError):
                        print("- Xformers not available")
                
                # Set to eval mode
                model.eval()
                print("- Model set to eval mode")
            # Int8 quantization using HuggingFace integration
            elif self.config.use_int8:
                print("- Using INT8 quantization")
                # For now, just use standard loading with INT8 config
                compute_dtype = torch.float16 if self.config.use_fp16 else torch.float32
                quantization_config = BitsAndBytesConfig(
                    load_in_8bit=True,
                    llm_int8_threshold=6.0,
                    llm_int8_has_fp16_weight=False,
                )
                
                model = AutoModel.from_pretrained(
                    self.config.model_path,
                    quantization_config=quantization_config,
                    torch_dtype=compute_dtype,
                    device_map="auto"
                )
                
                if model is None:
                    raise ValueError("Model loading returned None")
                    
                print(f"- Model type: {type(model)}")
                model.eval()
                print("- Model set to eval mode")
                
            else:
                # Standard loading for FP16/FP32
                model = AutoModel.from_pretrained(self.config.model_path)
                print("- Model loaded in standard precision")
                print(f"- Model type: {type(model)}")
                
                # Apply standard optimizations
                # set default to half
                import torch
                torch.set_default_dtype(torch.bfloat16)
                model = ModelOptimizer.optimize(model, self.config)
                model = model.half()
                # add torch compile
                model = torch.compile(model)
            
            # Final check to ensure model is not None
            if model is None:
                raise ValueError("Model is None after optimization")
                
            print(f"- Final model type: {type(model)}")
            return model
        
        except Exception as e:
            print(f"ERROR loading model: {str(e)}")
            import traceback
            traceback.print_exc()
            raise
    
    def _create_random_batch(self, batch_size: int) -> torch.Tensor:
        device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
        return torch.randint(
            0, 1000,
            (batch_size, self.config.seq_length),
            device=device,
            dtype=torch.long
        )
    
    def _run_inference(
        self,
        input_ids: torch.Tensor,
        graph_wrapper: Optional[GraphWrapper] = None
    ) -> Tuple[float, torch.Tensor]:
        attention_mask = torch.ones_like(input_ids)
        
        with torch.no_grad(), self.timer.timing():
            if graph_wrapper is not None:
                output = graph_wrapper(input_ids, attention_mask)
            else:
                output = self.model(input_ids=input_ids, attention_mask=attention_mask)
        
        return self.timer.elapsed_time(), output
    
    def run(self) -> Dict[int, Dict[str, float]]:
        results = {}
        
        # Reset peak memory stats
        if torch.cuda.is_available():
            torch.cuda.reset_peak_memory_stats()
        elif torch.backends.mps.is_available():
            # MPS doesn't have reset_peak_memory_stats, skip it
            pass
        else:
            print("- No GPU memory stats available")
        
        for batch_size in self.config.batch_sizes:
            print(f"\nTesting batch size: {batch_size}")
            times = []
            
            # Get or create graph for this batch size
            graph_wrapper = (
                self.graphs.get_or_create(batch_size)
                if self.graphs is not None
                else None
            )
            
            # Pre-allocate input tensor
            input_ids = self._create_random_batch(batch_size)
            print(f"Input shape: {input_ids.shape}")
            
            # Run benchmark
            for i in tqdm(range(self.config.num_runs), desc=f"Batch size {batch_size}"):
                try:
                    elapsed_time, output = self._run_inference(input_ids, graph_wrapper)
                    if i == 0:  # Only print on first run
                        print(f"Output shape: {output.last_hidden_state.shape}")
                    times.append(elapsed_time)
                except Exception as e:
                    print(f"Error during inference: {e}")
                    break
            
            if not times:
                print(f"No successful runs for batch size {batch_size}, skipping")
                continue
                
            # Calculate statistics
            avg_time = np.mean(times)
            std_time = np.std(times)
            throughput = batch_size / avg_time
            
            results[batch_size] = {
                "avg_time": avg_time,
                "std_time": std_time,
                "throughput": throughput,
            }
            
            print(f"Avg Time: {avg_time:.4f}s ± {std_time:.4f}s")
            print(f"Throughput: {throughput:.2f} sequences/second")
        
        # Log memory usage
        if torch.cuda.is_available():
            peak_memory_gb = torch.cuda.max_memory_allocated() / (1024 ** 3)
        elif torch.backends.mps.is_available():
            # MPS doesn't have max_memory_allocated, use 0
            peak_memory_gb = 0.0
        else:
            peak_memory_gb = 0.0
            print("- No GPU memory usage available")
        
        if peak_memory_gb > 0:
            print(f"\nPeak GPU memory usage: {peak_memory_gb:.2f} GB")
        else:
            print("\n- GPU memory usage not available")
        
        # Add memory info to results
        for batch_size in results:
            results[batch_size]["peak_memory_gb"] = peak_memory_gb
        
        return results


def main():
    parser = argparse.ArgumentParser(description="Model Inference Benchmark")
    parser.add_argument(
        "--model_path",
        type=str,
        default="facebook/contriever",
        help="Path to the model",
    )
    parser.add_argument(
        "--batch_sizes",
        type=str,
        default="1,2,4,8,16,32",
        help="Comma-separated list of batch sizes",
    )
    parser.add_argument(
        "--seq_length",
        type=int,
        default=256,
        help="Sequence length for input",
    )
    parser.add_argument(
        "--num_runs",
        type=int,
        default=5,
        help="Number of runs for each batch size",
    )
    parser.add_argument(
        "--use_fp16",
        action="store_true",
        help="Enable FP16 inference",
    )
    parser.add_argument(
        "--use_int4",
        action="store_true",
        help="Enable INT4 quantization using bitsandbytes",
    )
    parser.add_argument(
        "--use_int8",
        action="store_true",
        help="Enable INT8 quantization for both activations and weights using bitsandbytes",
    )
    parser.add_argument(
        "--use_cuda_graphs",
        action="store_true",
        help="Enable CUDA Graphs optimization (only on NVIDIA GPUs)",
    )
    parser.add_argument(
        "--use_flash_attention",
        action="store_true",
        help="Enable Flash Attention 2 if available (only on NVIDIA GPUs)",
    )
    parser.add_argument(
        "--use_linear8bitlt",
        action="store_true",
        help="Enable Linear8bitLt quantization for all linear layers",
    )
    
    args = parser.parse_args()
    
    # Print arguments for debugging
    print("\nCommand line arguments:")
    for arg, value in vars(args).items():
        print(f"- {arg}: {value}")
    
    config = BenchmarkConfig(
        model_path=args.model_path,
        batch_sizes=[int(bs) for bs in args.batch_sizes.split(",")],
        seq_length=args.seq_length,
        num_runs=args.num_runs,
        use_fp16=args.use_fp16,
        use_int4=args.use_int4,
        use_int8=args.use_int8,  # Add this line
        use_cuda_graphs=args.use_cuda_graphs,
        use_flash_attention=args.use_flash_attention,
        use_linear8bitlt=args.use_linear8bitlt,
    )
        
    # Print configuration for debugging
    print("\nBenchmark configuration:")
    for field, value in vars(config).items():
        print(f"- {field}: {value}")
    
    try:
        benchmark = Benchmark(config)
        results = benchmark.run()
        
        # Save results to file
        import json
        import os
        
        # Create results directory if it doesn't exist
        os.makedirs("results", exist_ok=True)
        
        # Generate filename based on configuration
        precision_type = "int4" if config.use_int4 else "int8" if config.use_int8 else "fp16" if config.use_fp16 else "fp32"
        model_name = os.path.basename(config.model_path)
        output_file = f"results/benchmark_{model_name}_{precision_type}.json"
        
        # Save results
        with open(output_file, "w") as f:
            json.dump(
                {
                    "config": {k: str(v) if isinstance(v, list) else v for k, v in vars(config).items()},
                    "results": {str(k): v for k, v in results.items()}
                }, 
                f, 
                indent=2
            )
        print(f"Results saved to {output_file}")
        
    except Exception as e:
        print(f"Benchmark failed: {e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()