"""
座舱单目深度估计应用

基于轻量级自监督深度估计
适用于: OOP距离检测, CPD体型估计, 乘员分类

输入: DMS/OMS摄像头 RGB帧
输出: 深度图 (相对距离, mm)
"""

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

class LocalFeatureEnhancement(nn.Module):
    """
    局部特征增强模块 (LFEM)
    
    通过多尺度局部注意力增强细节特征
    用于提升近距离深度估计精度
    """
    
    def __init__(self, in_channels: int, out_channels: int):
        super().__init__()
        # 多尺度卷积
        self.conv3 = nn.Conv2d(in_channels, out_channels // 4, 3, padding=1)
        self.conv5 = nn.Conv2d(in_channels, out_channels // 4, 5, padding=2)
        self.conv7 = nn.Conv2d(in_channels, out_channels // 4, 7, padding=3)
        self.conv1x1 = nn.Conv2d(in_channels, out_channels // 4, 1)
        
        # 融合
        self.fuse = nn.Conv2d(out_channels, out_channels, 1)
        self.bn = nn.BatchNorm2d(out_channels)
        self.act = nn.SiLU()
    
    def forward(self, x):
        f3 = self.conv3(x)
        f5 = self.conv5(x)
        f7 = self.conv7(x)
        f1 = self.conv1x1(x)
        
        out = torch.cat([f3, f5, f7, f1], dim=1)
        out = self.act(self.bn(self.fuse(out)))
        return out


class LightweightDepthNet(nn.Module):
    """
    轻量级单目深度估计网络
    
    设计目标:
    - 参数量 < 5M (可量化到INT8 < 2MB)
    - 推理速度 > 30fps @ 640x480 (Jetson Orin NX)
    - 近距离(0.5-3m)精度 < 5cm (座舱内关键范围)
    
    架构: MobileNetV3 encoder + LFEM + 多尺度decoder
    """
    
    def __init__(self, pretrained_encoder=True):
        super().__init__()
        
        # 编码器: 使用MobileNetV3-Small (轻量)
        # 实际可用 torchvision.models.mobilenet_v3_small(pretrained=pretrained_encoder)
        self.encoder = nn.Sequential(
            nn.Conv2d(3, 16, 3, stride=2, padding=1),  # 320x240
            nn.SiLU(),
            LFEMBlock(16, 16),
            nn.Conv2d(16, 24, 3, stride=2, padding=1),  # 160x120
            nn.SiLU(),
            LFEMBlock(24, 24),
            nn.Conv2d(24, 40, 3, stride=2, padding=1),  # 80x60
            nn.SiLU(),
            LFEMBlock(40, 40),
            nn.Conv2d(40, 48, 3, stride=2, padding=1),  # 40x30
            nn.SiLU(),
            LFEMBlock(48, 48),
            nn.Conv2d(48, 96, 3, stride=2, padding=1),  # 20x15
            nn.SiLU(),
        )
        
        # LFEM增强
        self.lfem = LocalFeatureEnhancement(96, 96)
        
        # 解码器: 逐步上采样
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(96, 48, 4, stride=2, padding=1),  # 40x30
            nn.SiLU(),
            LFEMBlock(48, 48),
            nn.ConvTranspose2d(48, 24, 4, stride=2, padding=1),  # 80x60
            nn.SiLU(),
            LFEMBlock(24, 24),
            nn.ConvTranspose2d(24, 16, 4, stride=2, padding=1),  # 160x120
            nn.SiLU(),
            LFEMBlock(16, 16),
            nn.ConvTranspose2d(16, 1, 4, stride=2, padding=1),  # 320x240
            nn.Sigmoid(),  # 归一化到0-1 (相对深度)
        )
    
    def forward(self, x):
        feat = self.encoder(x)
        feat = self.lfem(feat)
        depth = self.decoder(feat)
        # 插值到输入分辨率
        depth = F.interpolate(depth, size=x.shape[2:], mode='bilinear', align_corners=False)
        return depth.squeeze(1)


class LFEMBlock(nn.Module):
    """残差LFEM块"""
    def __init__(self, channels, out_channels):
        super().__init__()
        self.conv = nn.Conv2d(channels, out_channels, 3, padding=1, groups=channels)  # DW
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.pw = nn.Conv2d(out_channels, out_channels, 1)  # PW
        self.bn2 = nn.BatchNorm2d(out_channels)
        self.act = nn.SiLU()
    
    def forward(self, x):
        residual = x
        out = self.act(self.bn1(self.conv(x)))
        out = self.bn2(self.pw(out))
        if x.shape[1] == out.shape[1]:
            out = out + residual
        return self.act(out)


class CabinDepthEstimator:
    """座舱深度估计应用"""
    
    def __init__(self, model_path=None, device='cuda'):
        self.device = device
        self.model = LightweightDepthNet().to(device)
        if model_path:
            self.model.load_state_dict(torch.load(model_path, map_location=device))
        self.model.eval()
        
        # 座舱深度范围 (cm)
        self.MIN_DEPTH = 30   # 0.3m (距摄像头最近)
        self.MAX_DEPTH = 300   # 3.0m (座舱最远)
    
    def estimate_depth(self, image: np.ndarray) -> np.ndarray:
        """
        估计座舱深度
        
        Args:
            image: RGB图像 (H, W, 3), uint8
        
        Returns:
            depth_map: 深度图 (H, W), float32, 单位cm
        """
        h, w = image.shape[:2]
        
        # 预处理
        img = torch.from_numpy(image).float().permute(2, 0, 1) / 255.0
        img = img.unsqueeze(0).to(self.device)
        
        # 推理
        with torch.no_grad():
            depth_norm = self.model(img)  # 0-1
        
        # 转为实际深度 (cm)
        depth_cm = depth_norm.squeeze().cpu().numpy()
        depth_cm = self.MIN_DEPTH + depth_cm * (self.MAX_DEPTH - self.MIN_DEPTH)
        
        return depth_cm
    
    def detect_oop_distance(self, depth_map: np.ndarray, bbox: tuple) -> dict:
        """
        检测乘员距方向盘距离 (OOP)
        
        Args:
            depth_map: (H, W) depth in cm
            bbox: (x1, y1, x2, y2) 乘员边界框
        
        Returns:
            distance info
        """
        x1, y1, x2, y2 = bbox
        region = depth_map[y1:y2, x1:x2]
        
        # 乘员中心深度
        center_depth = np.median(region)
        
        # 最小深度(最近点)
        min_depth = np.percentile(region, 10)
        
        # OOP判定
        is_oop = min_depth < 40  # 距方向盘<40cm为OOP
        
        return {
            'center_distance_cm': round(float(center_depth), 1),
            'min_distance_cm': round(float(min_depth), 1),
            'is_oop': bool(is_oop),
            'oop_severity': 'high' if min_depth < 30 else 'medium' if min_depth < 40 else 'normal'
        }


if __name__ == "__main__":
    print("=== 座舱单目深度估计 ===\n")
    
    # 模拟推理(无需实际模型)
    model = LightweightDepthNet()
    n_params = sum(p.numel() for p in model.parameters())
    print(f"模型参数量: {n_params/1e6:.2f}M")
    print(f"INT8量化后大小: {n_params/1e6:.1f}MB")
    
    # 模拟输入
    dummy = torch.randn(1, 3, 240, 320)
    with torch.no_grad():
        output = model(dummy)
    print(f"输入: {dummy.shape}")
    print(f"输出: {output.shape}")
    print(f"深度范围: {output.min():.3f} - {output.max():.3f} (归一化)")
    
    print(f"\n座舱深度范围: 30-300cm")
    print(f"应用场景: OOP检测, CPD体型估计, 乘员分类")
    print(f"部署目标: Jetson Orin NX INT8, 预计>30fps @ 640x480")