import torch
import torch.nn as nn
from typing import Optional, Dict, Any
import onnx
import onnxruntime as ort
import os


class ModelOptimizer:
    """模型优化器"""
    def __init__(self, model: nn.Module, device: str = 'cuda'):
        self.model = model
        self.device = device
        self.optimized_model = None
    
    def optimize_for_inference(self, use_jit: bool = True, use_cuda_graph: bool = False):
        """优化模型用于推理"""
        self.model.eval()
        
        # 应用一系列优化
        if use_jit:
            self._jit_compile()
        
        if use_cuda_graph and torch.cuda.is_available():
            self._capture_cuda_graph()
        
        # 应用其他优化
        self._apply_inference_optimizations()
        
        return self.optimized_model or self.model
    
    def _jit_compile(self):
        """使用TorchScript编译模型"""
        try:
            # 创建示例输入
            example_input = torch.randn(1, 4, 64, 64, device=self.device)
            example_timestep = torch.tensor([500], device=self.device)
            example_context = torch.randn(1, 77, 768, device=self.device)
            
            # 脚本编译
            scripted_model = torch.jit.trace(
                self.model,
                (example_input, example_timestep, example_context),
                check_trace=False
            )
            
            self.optimized_model = scripted_model
            print("模型已使用TorchScript编译")
        except Exception as e:
            print(f"TorchScript编译失败: {e}")
    
    def _capture_cuda_graph(self):
        """捕获CUDA图（用于重复推理）"""
        if not torch.cuda.is_available():
            return
        
        # 创建静态输入
        static_input = torch.randn(1, 4, 64, 64, device='cuda', dtype=torch.float16)
        static_timestep = torch.tensor([500], device='cuda')
        static_context = torch.randn(1, 77, 768, device='cuda', dtype=torch.float16)
        
        # 预热
        with torch.no_grad():
            for _ in range(3):
                _ = self.model(static_input, static_timestep, static_context)
        
        # 捕获图
        graph = torch.cuda.CUDAGraph()
        with torch.cuda.graph(graph):
            static_output = self.model(static_input, static_timestep, static_context)
        
        # 创建包装函数
        def graph_executor(input_tensor, timestep, context):
            static_input.copy_(input_tensor)
            static_timestep.copy_(timestep)
            static_context.copy_(context)
            graph.replay()
            return static_output.clone()
        
        self.optimized_model = graph_executor
        print("已捕获CUDA图")
    
    def _apply_inference_optimizations(self):
        """应用推理优化"""
        # 设置为评估模式
        self.model.eval()
        
        # 融合操作（如果可用）
        if hasattr(torch, 'compile'):
            try:
                self.model = torch.compile(self.model, mode='max-autotune')
                print("模型已使用torch.compile优化")
            except:
                pass
        
        # 使用半精度
        if self.device == 'cuda':
            self.model.half()
            print("模型已转换为半精度")
    
    def quantize(self, quantization_mode: str = 'dynamic'):
        """量化模型"""
        if quantization_mode == 'dynamic':
            # 动态量化
            quantized_model = torch.quantization.quantize_dynamic(
                self.model,
                {nn.Linear, nn.Conv2d},
                dtype=torch.qint8
            )
            self.optimized_model = quantized_model
            print("模型已动态量化")
        
        elif quantization_mode == 'static':
            # 静态量化需要校准数据
            print("静态量化需要校准数据，暂未实现")
        
        else:
            raise ValueError(f"未知的量化模式: {quantization_mode}")
    
    def prune(self, pruning_rate: float = 0.2):
        """剪枝模型"""
        from torch.nn.utils import prune
        
        # 对线性层和卷积层进行剪枝
        for name, module in self.model.named_modules():
            if isinstance(module, (nn.Linear, nn.Conv2d)):
                prune.l1_unstructured(module, name='weight', amount=pruning_rate)
                prune.remove(module, 'weight')
        
        print(f"模型已剪枝，剪枝率: {pruning_rate}")
    
    def get_model_size(self) -> Dict[str, float]:
        """获取模型大小"""
        # 计算参数量
        total_params = sum(p.numel() for p in self.model.parameters())
        trainable_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
        
        # 计算模型大小（MB）
        param_size = 0
        for param in self.model.parameters():
            param_size += param.nelement() * param.element_size()
        
        buffer_size = 0
        for buffer in self.model.buffers():
            buffer_size += buffer.nelement() * buffer.element_size()
        
        size_mb = (param_size + buffer_size) / 1024**2
        
        return {
            'total_params': total_params,
            'trainable_params': trainable_params,
            'size_mb': size_mb
        }


class ONNXExporter:
    """ONNX导出器"""
    def __init__(self, model: nn.Module):
        self.model = model
    
    def export(
        self,
        output_path: str,
        input_shape: tuple = (1, 4, 64, 64),
        opset_version: int = 14,
        dynamic_axes: Optional[Dict] = None
    ):
        """导出为ONNX格式"""
        # 设置为评估模式
        self.model.eval()
        
        # 创建示例输入
        dummy_input = torch.randn(*input_shape)
        dummy_timestep = torch.tensor([500])
        dummy_context = torch.randn(1, 77, 768)
        
        # 默认动态轴
        if dynamic_axes is None:
            dynamic_axes = {
                'input': {0: 'batch_size'},
                'timestep': {0: 'batch_size'},
                'context': {0: 'batch_size'},
                'output': {0: 'batch_size'}
            }
        
        # 导出
        torch.onnx.export(
            self.model,
            (dummy_input, dummy_timestep, dummy_context),
            output_path,
            input_names=['input', 'timestep', 'context'],
            output_names=['output'],
            dynamic_axes=dynamic_axes,
            opset_version=opset_version,
            do_constant_folding=True,
            verbose=False
        )
        
        print(f"模型已导出为ONNX: {output_path}")
        
        # 验证ONNX模型
        self._validate_onnx(output_path)
    
    def _validate_onnx(self, onnx_path: str):
        """验证ONNX模型"""
        try:
            onnx_model = onnx.load(onnx_path)
            onnx.checker.check_model(onnx_model)
            print("ONNX模型验证成功")
        except Exception as e:
            print(f"ONNX模型验证失败: {e}")
    
    def optimize_onnx(self, onnx_path: str, optimized_path: str):
        """优化ONNX模型"""
        try:
            # 使用ONNX Runtime优化
            sess_options = ort.SessionOptions()
            sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
            
            # 创建优化会话
            ort_session = ort.InferenceSession(onnx_path, sess_options)
            
            # 保存优化后的模型
            optimized_model = ort_session.get_model()
            onnx.save(optimized_model, optimized_path)
            
            print(f"ONNX模型已优化: {optimized_path}")
        except Exception as e:
            print(f"ONNX优化失败: {e}")


class MemoryEfficientInference:
    """内存高效推理"""
    def __init__(self, model: nn.Module, chunk_size: int = 32):
        self.model = model
        self.chunk_size = chunk_size
        
    def chunked_inference(self, x: torch.Tensor, t: torch.Tensor, context: torch.Tensor) -> torch.Tensor:
        """分块推理，减少内存使用"""
        B, C, H, W = x.shape
        output = torch.zeros_like(x)
        
        # 分块处理
        for i in range(0, H, self.chunk_size):
            for j in range(0, W, self.chunk_size):
                # 提取块
                chunk = x[:, :, i:i+self.chunk_size, j:j+self.chunk_size]
                
                # 推理
                with torch.no_grad():
                    chunk_output = self.model(chunk, t, context)
                
                # 存储结果
                output[:, :, i:i+self.chunk_size, j:j+self.chunk_size] = chunk_output
        
        return output
    
    def tiled_inference(self, x: torch.Tensor, t: torch.Tensor, context: torch.Tensor, tile_size: int = 512) -> torch.Tensor:
        """平铺推理，用于大图像"""
        B, C, H, W = x.shape
        
        # 如果图像不大，直接推理
        if H <= tile_size and W <= tile_size:
            with torch.no_grad():
                return self.model(x, t, context)
        
        # 计算平铺数量
        n_tiles_h = (H + tile_size - 1) // tile_size
        n_tiles_w = (W + tile_size - 1) // tile_size
        
        output = torch.zeros_like(x)
        
        # 处理每个平铺
        for i in range(n_tiles_h):
            for j in range(n_tiles_w):
                # 计算平铺位置
                h_start = i * tile_size
                w_start = j * tile_size
                h_end = min(h_start + tile_size, H)
                w_end = min(w_start + tile_size, W)
                
                # 提取平铺
                tile = x[:, :, h_start:h_end, w_start:w_end]
                
                # 推理
                with torch.no_grad():
                    tile_output = self.model(tile, t, context)
                
                # 存储结果
                output[:, :, h_start:h_end, w_start:w_end] = tile_output
        
        return output


class InferenceBenchmark:
    """推理基准测试"""
    def __init__(self, model: nn.Module, device: str = 'cuda'):
        self.model = model
        self.device = device
        
    def benchmark(
        self,
        input_shape: tuple = (1, 4, 64, 64),
        num_iterations: int = 100,
        warmup_iterations: int = 10
    ) -> Dict[str, float]:
        """运行基准测试"""
        # 准备输入
        x = torch.randn(*input_shape, device=self.device)
        t = torch.tensor([500], device=self.device)
        context = torch.randn(1, 77, 768, device=self.device)
        
        # 预热
        print("预热...")
        with torch.no_grad():
            for _ in range(warmup_iterations):
                _ = self.model(x, t, context)
        
        # 同步
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        
        # 基准测试
        print("运行基准测试...")
        import time
        
        times = []
        
        for i in range(num_iterations):
            start_time = time.time()
            
            with torch.no_grad():
                _ = self.model(x, t, context)
            
            if torch.cuda.is_available():
                torch.cuda.synchronize()
            
            end_time = time.time()
            times.append(end_time - start_time)
        
        # 统计
        times = torch.tensor(times)
        
        stats = {
            'mean_ms': times.mean().item() * 1000,
            'std_ms': times.std().item() * 1000,
            'min_ms': times.min().item() * 1000,
            'max_ms': times.max().item() * 1000,
            'fps': 1 / times.mean().item(),
            'num_iterations': num_iterations
        }
        
        # 打印结果
        print("\n" + "="*50)
        print("推理基准测试结果:")
        print(f"平均推理时间: {stats['mean_ms']:.2f} ms")
        print(f"标准差: {stats['std_ms']:.2f} ms")
        print(f"最小推理时间: {stats['min_ms']:.2f} ms")
        print(f"最大推理时间: {stats['max_ms']:.2f} ms")
        print(f"FPS: {stats['fps']:.2f}")
        print("="*50)
        
        return stats


def optimize_model_for_p4(model: nn.Module) -> nn.Module:
    """为P4优化模型"""
    optimizer = ModelOptimizer(model)
    
    # 获取模型大小
    size_info = optimizer.get_model_size()
    print(f"优化前模型大小: {size_info['size_mb']:.2f} MB")
    
    # 应用优化
    optimized_model = optimizer.optimize_for_inference(
        use_jit=True,
        use_cuda_graph=False  # P4可能不支持
    )
    
    # 量化（可选）
    if size_info['size_mb'] > 500:  # 如果模型大于500MB，进行量化
        optimizer.quantize('dynamic')
    
    # 获取优化后的模型大小
    size_info_after = optimizer.get_model_size()
    print(f"优化后模型大小: {size_info_after['size_mb']:.2f} MB")
    print(f"压缩比: {size_info['size_mb'] / size_info_after['size_mb']:.2f}x")
    
    return optimized_model


def test_optimization():
    """测试优化"""
    import torch.nn as nn
    
    # 创建模拟模型
    class MockModel(nn.Module):
        def __init__(self):
            super().__init__()
            self.conv1 = nn.Conv2d(4, 64, 3, padding=1)
            self.conv2 = nn.Conv2d(64, 4, 3, padding=1)
        
        def forward(self, x, t, context):
            x = self.conv1(x)
            x = nn.functional.relu(x)
            x = self.conv2(x)
            return x
    
    model = MockModel()
    
    # 测试优化器
    optimizer = ModelOptimizer(model)
    optimized_model = optimizer.optimize_for_inference()
    
    # 测试基准测试
    benchmark = InferenceBenchmark(model)
    stats = benchmark.benchmark(num_iterations=10)
    
    # 测试ONNX导出
    exporter = ONNXExporter(model)
    exporter.export('./test_model.onnx', input_shape=(1, 4, 64, 64))
    
    return optimized_model, stats


if __name__ == '__main__':
    optimized_model, stats = test_optimization()