import torch
import torch.nn as nn
import torch.ao.quantization as quant
import numpy as np
import time

class DMSModel(nn.Module):
    """
    轻量级DMS模型: 面部关键点+眼开度+头部姿态
    
    输入: 224x224 RGB-IR
    输出: 68关键点(2D) + 眼开度(2) + 头部姿态(3)
    """
    
    def __init__(self):
        super().__init__()
        
        # 共享骨干 (MobileNetV3-Small风格)
        self.features = nn.Sequential(
            # Block 1
            nn.Conv2d(3, 16, 3, stride=2, padding=1),
            nn.BatchNorm2d(16),
            nn.ReLU6(),
            
            # Block 2 (depthwise separable)
            nn.Conv2d(16, 16, 3, groups=16, padding=1),
            nn.Conv2d(16, 32, 1),
            nn.BatchNorm2d(32),
            nn.ReLU6(),
            
            # Block 3
            nn.Conv2d(32, 32, 3, groups=32, stride=2, padding=1),
            nn.Conv2d(32, 64, 1),
            nn.BatchNorm2d(64),
            nn.ReLU6(),
            
            # Block 4
            nn.Conv2d(64, 64, 3, groups=64, padding=1),
            nn.Conv2d(64, 128, 1),
            nn.BatchNorm2d(128),
            nn.ReLU6(),
            
            # Block 5
            nn.Conv2d(128, 128, 3, groups=128, stride=2, padding=1),
            nn.Conv2d(128, 256, 1),
            nn.BatchNorm2d(256),
            nn.ReLU6(),
            
            nn.AdaptiveAvgPool2d((7, 7))
        )
        
        # 关键点回归头
        self.kp_head = nn.Sequential(
            nn.Flatten(),
            nn.Linear(256 * 7 * 7, 512),
            nn.ReLU6(),
            nn.Linear(512, 68 * 2)  # 68个2D关键点
        )
        
        # 眼开度+头部姿态头
        self.aux_head = nn.Sequential(
            nn.Linear(256 * 7 * 7, 128),
            nn.ReLU6(),
            nn.Linear(128, 5)  # 2眼开度 + 3头部姿态
        )
        
        # 量化配置
        self.quant = quant.QuantStub()
        self.dequant = quant.DeQuantStub()
    
    def forward(self, x):
        x = self.quant(x)
        feat = self.features(x)
        flat = feat.flatten(1)
        kp = self.kp_head(flat)
        aux = self.aux_head(flat)
        return self.dequant(kp), self.dequant(aux)


def prepare_qat(model):
    """准备量化感知训练"""
    model.qconfig = quant.get_default_qat_qconfig('qnnpack')
    quant.prepare_qat(model, inplace=True)
    return model


def convert_to_int8(model, calib_loader):
    """
    转换为INT8量化模型
    
    Args:
        model: QAT训练后的模型
        calib_loader: 校准数据加载器
    
    Returns:
        int8_model: INT8量化模型
    """
    model.eval()
    
    # 校准
    with torch.no_grad():
        for data, _ in calib_loader:
            model(data)
    
    # 转换
    int8_model = quant.convert(model)
    return int8_model


def benchmark_model(model, input_shape, device='cpu', n_runs=100):
    """基准测试"""
    model = model.to(device).eval()
    dummy = torch.randn(*input_shape).to(device)
    
    # 预热
    with torch.no_grad():
        for _ in range(10):
            model(dummy)
    
    # 测试
    times = []
    with torch.no_grad():
        for _ in range(n_runs):
            t0 = time.perf_counter()
            model(dummy)
            t1 = time.perf_counter()
            times.append((t1 - t0) * 1000)
    
    return {
        'mean_ms': np.mean(times),
        'p99_ms': np.percentile(times, 99),
        'fps': 1000 / np.mean(times)
    }


if __name__ == "__main__":
    model = DMSModel()
    
    # 模型大小对比
    fp32_size = sum(p.nelement() * 4 for p in model.parameters()) / 1024 / 1024
    print(f"=== 模型信息 ===")
    print(f"FP32 大小: {fp32_size:.2f} MB")
    print(f"参数量: {sum(p.nelement() for p in model.parameters()):,}")
    
    # FP32基准
    fp32_result = benchmark_model(model, (1, 3, 224, 224))
    print(f"\n=== FP32 性能 (CPU) ===")
    print(f"延迟: {fp32_result['mean_ms']:.1f}ms")
    print(f"P99: {fp32_result['p99_ms']:.1f}ms")
    print(f"FPS: {fp32_result['fps']:.0f}")
    
    # 模拟INT8量化
    model_int8 = prepare_qat(model)
    # 实际中需要校准数据
    # model_int8 = convert_to_int8(model_int8, calib_loader)
    
    print(f"\n=== INT8 量化预期效果 ===")
    print(f"模型大小: {fp32_size/4:.2f} MB (4x压缩)")
    print(f"CPU延迟: {fp32_result['mean_ms']/3:.1f}ms (3x加速)")
    print(f"NPU延迟: {fp32_result['mean_ms']/8:.1f}ms (8x加速)")
    print(f"精度损失: 1-3% (QAT)")
    
    print(f"\n=== Snapdragon QCS8255 部署参数 ===")
    print(f"芯片: QCS8255")
    print(f"NPU: Hexagon 26 TOPS (INT8)")
    print(f"模型: {fp32_size/4:.1f}MB (INT8)")
    print(f"延迟: ~{fp32_result['mean_ms']/8:.0f}ms (NPU)")
    print(f"FPS: ~{1000/(fp32_result['mean_ms']/8):.0f} (NPU)")
    print(f"功耗: ~1.5W (DMS pipeline)")
    print(f"量化: QAT → INT8 → QNN → HTP")