import numpy as np
from typing import Tuple

class ForwardThermalModel:
    """
    前向热物理模型
    
    描述3D人体热源到2D热阵列平面的映射关系。
    基于热辐射物理方程：
    I(d) = ε * σ * T^4 / (4 * π * d^2)
    
    其中：
    - I(d): 距离d处的辐射强度
    - ε: 发射率（人体≈0.98）
    - σ: Stefan-Boltzmann常数 (5.67e-8 W/m²K⁴)
    - T: 表面温度(K)
    - d: 传感器到热源的距离
    """
    
    STEFAN_BOLTZMANN = 5.67e-8  # W/m²K⁴
    HUMAN_EMISSIVITY = 0.98
    
    def __init__(self, sensor_resolution: Tuple[int, int] = (16, 16),
                 fov: Tuple[float, float] = (60.0, 45.0),
                 ambient_temp: float = 22.0):
        self.resolution = sensor_resolution
        self.fov = np.radians(fov)
        self.ambient_temp = ambient_temp
        self.ambient_k = ambient_temp + 273.15
        
        # 生成传感器像素方向向量
        rows, cols = sensor_resolution
        v_angles = np.linspace(-self.fov[1]/2, self.fov[1]/2, rows)
        h_angles = np.linspace(-self.fov[0]/2, self.fov[0]/2, cols)
        self.pixel_directions = np.array([
            [np.cos(v) * np.sin(h), np.sin(v), np.cos(v) * np.cos(h)]
            for v in v_angles for h in h_angles
        ]).reshape(rows, cols, 3)
    
    def forward(self, heat_sources: np.ndarray, depths: np.ndarray) -> np.ndarray:
        """
        前向模型：从3D热源到2D热阵列读数
        
        Args:
            heat_sources: 人体表面热源点 (N, 3) [x, y, z] 坐标
            depths: 各热源点到传感器距离 (N,)
            
        Returns:
            thermal_image: 2D热阵列读数 (rows, cols) 温度值
        """
        rows, cols = self.resolution
        thermal_image = np.full((rows, cols), self.ambient_temp)
        
        for i, (source, depth) in enumerate(zip(heat_sources, depths)):
            # 计算热源方向
            direction = source / (np.linalg.norm(source) + 1e-8)
            
            # 找到最近的传感器像素
            dots = self.pixel_directions.reshape(-1, 3) @ direction
            pixel_idx = np.argmax(dots)
            r, c = pixel_idx // cols, pixel_idx % cols
            
            # 热辐射强度衰减（距离平方）
            source_temp_k = 34.0 + 273.15  # 人体表面~34°C
            intensity = (self.HUMAN_EMISSIVITY * self.STEFAN_BOLTZMANN * 
                        source_temp_k**4 / (4 * np.pi * max(depth**2, 0.01)))
            
            # 转换为温度增量
            temp_delta = (intensity / (self.STEFAN_BOLTZMANN * self.ambient_k**3 + 1e-8)) * 100
            thermal_image[r, c] = max(thermal_image[r, c], self.ambient_temp + temp_delta)
        
        return thermal_image
    
    def inverse_prior(self, thermal_image: np.ndarray) -> np.ndarray:
        """
        粗略逆模型：从2D热阵列读数估计深度
        
        Args:
            thermal_image: (rows, cols) 温度矩阵
            
        Returns:
            depth_map: (rows, cols) 粗略深度图
        """
        # 基于温度差异估算距离
        temp_diff = thermal_image - self.ambient_temp
        temp_diff = np.maximum(temp_diff, 0.01)
        
        # 逆辐射模型：d ∝ sqrt(1/ΔT)
        depth_map = np.sqrt(12.0 / temp_diff)  # 12.0为校准常数
        depth_map = np.clip(depth_map, 0.3, 5.0)
        
        return depth_map


# 实际测试
if __name__ == "__main__":
    model = ForwardThermalModel(sensor_resolution=(16, 16))
    
    # 模拟人体热源（简化为5个关键点）
    heat_sources = np.array([
        [0.0, 0.2, 1.0],   # 头部
        [-0.2, 0.0, 1.1],  # 左肩
        [0.2, 0.0, 1.1],   # 右肩
        [-0.15, -0.3, 1.2],# 左膝
        [0.15, -0.3, 1.2], # 右膝
    ])
    depths = np.array([1.0, 1.1, 1.1, 1.2, 1.2])
    
    thermal_img = model.forward(heat_sources, depths)
    depth_map = model.inverse_prior(thermal_img)
    
    print(f"热阵列分辨率: {thermal_img.shape}")
    print(f"最大温度: {np.max(thermal_img):.2f}°C")
    print(f"环境温度: {model.ambient_temp}°C")
    print(f"深度估计范围: {depth_map.min():.2f}m ~ {depth_map.max():.2f}m")

import torch
import torch.nn as nn

class MultiPrimitiveEstimator(nn.Module):
    """
    多原语估计模块
    
    从2D热阵列帧序列中联合估计：
    - 深度图 (depth map)
    - 热辐射强度 (thermal intensity)
    - 运动场 (motion field)
    
    采用自监督训练策略：
    1. 前向渲染损失：重建的热图像应匹配原始输入
    2. 时间一致性：相邻帧的深度估计应平滑变化
    3. 稀疏性约束：深度图应在人体区域集中
    """
    
    def __init__(self, input_channels: int = 1, hidden_dim: int = 64):
        super().__init__()
        
        # 编码器：从热阵列帧提取特征
        self.encoder = nn.Sequential(
            nn.Conv2d(input_channels, 32, 3, padding=1),
            nn.ReLU(),
            nn.Conv2d(32, hidden_dim, 3, padding=1),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d((8, 8))
        )
        
        # 深度解码头
        self.depth_head = nn.Sequential(
            nn.ConvTranspose2d(hidden_dim, 32, 3, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(32, 1, 3, padding=1),
            nn.Sigmoid()  # 归一化到 [0, 1]
        )
        
        # 热强度解码头
        self.intensity_head = nn.Sequential(
            nn.ConvTranspose2d(hidden_dim, 32, 3, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(32, 1, 3, padding=1),
            nn.Sigmoid()
        )
        
        # 运动场解码头（从相邻帧差分）
        self.motion_head = nn.Sequential(
            nn.ConvTranspose2d(hidden_dim * 2, 32, 3, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(32, 2, 3, padding=1)  # 2D运动场
        )
    
    def forward(self, thermal_frame: torch.Tensor, 
                prev_frame: torch.Tensor = None) -> dict:
        """
        Args:
            thermal_frame: (B, 1, H, W) 当前热阵列帧
            prev_frame: (B, 1, H, W) 上一帧（可选）
            
        Returns:
            dict containing:
                - depth: (B, 1, H, W) 深度图
                - intensity: (B, 1, H, W) 热辐射强度
                - motion: (B, 2, H, W) 运动场
        """
        feat = self.encoder(thermal_frame)
        
        depth = self.depth_head(feat)
        intensity = self.intensity_head(feat)
        
        result = {'depth': depth, 'intensity': intensity}
        
        if prev_frame is not None:
            prev_feat = self.encoder(prev_frame)
            combined = torch.cat([feat, prev_feat], dim=1)
            motion = self.motion_head(combined)
            result['motion'] = motion
        
        return result


class TAP3DLoss(nn.Module):
    """
    TAP3D 自监督损失函数
    
    L = α * L_render + β * L_temporal + γ * L_sparse
    """
    
    def __init__(self, alpha: float = 1.0, beta: float = 0.3, gamma: float = 0.1):
        super().__init__()
        self.alpha = alpha
        self.beta = beta
        self.gamma = gamma
    
    def forward(self, predictions: dict, targets: dict) -> torch.Tensor:
        # 渲染损失：重建的热图像应匹配原始输入
        l_render = F.mse_loss(predictions['rendered'], targets['thermal'])
        
        # 时间一致性损失
        l_temporal = torch.tensor(0.0, device=l_render.device)
        if 'prev_depth' in predictions:
            l_temporal = F.smooth_l1_loss(
                predictions['depth'], 
                predictions['prev_depth'].detach()
            )
        
        # 稀疏性约束
        l_sparse = torch.norm(predictions['depth'], p=1) / predictions['depth'].numel()
        
        return (self.alpha * l_render + 
                self.beta * l_temporal + 
                self.gamma * l_sparse)


import torch.nn.functional as F

# 测试
if __name__ == "__main__":
    model = MultiPrimitiveEstimator(input_channels=1, hidden_dim=64)
    
    # 模拟热阵列输入 (batch=1, channel=1, 16x16)
    thermal = torch.randn(1, 1, 16, 16) + 22.0  # ~22°C ambient
    prev_thermal = torch.randn(1, 1, 16, 16) + 22.0
    
    output = model(thermal, prev_thermal)
    
    print(f"输入热阵列: {thermal.shape}")
    print(f"深度图: {output['depth'].shape}, 范围: [{output['depth'].min():.3f}, {output['depth'].max():.3f}]")
    print(f"热辐射强度: {output['intensity'].shape}")
    print(f"运动场: {output['motion'].shape}")

class GeometricPerspectiveFusion:
    """
    几何透视融合模块
    
    通过多视角几何约束实现：
    1. 热干扰抑制：利用空间几何一致性过滤环境热源
    2. 多人分离：基于热特征聚类和3D空间分离
    
    核心思想：
    - 人体热源在3D空间中应形成连通的骨架结构
    - 环境热源（阳光、暖风）不符合人体几何先验
    - 多人可通过空间聚类分离
    """
    
    def __init__(self, human_temp_range: tuple = (30.0, 37.0)):
        self.human_temp_range = human_temp_range
    
    def suppress_interference(self, depth_map: np.ndarray, 
                               thermal_image: np.ndarray) -> np.ndarray:
        """
        抑制环境热干扰
        
        Args:
            depth_map: (H, W) 深度图
            thermal_image: (H, W) 原始热阵列读数
            
        Returns:
            filtered_depth: (H, W) 过滤后的深度图
        """
        # 1. 人体温度范围过滤
        human_mask = ((thermal_image >= self.human_temp_range[0]) & 
                      (thermal_image <= self.human_temp_range[1]))
        
        # 2. 深度一致性过滤：人体深度应在合理范围
        depth_mask = (depth_map > 0.3) & (depth_map < 3.0)
        
        # 3. 形态学操作去除孤立热源
        from scipy.ndimage import binary_opening, binary_closing
        
        combined_mask = human_mask & depth_mask
        combined_mask = binary_opening(combined_mask, iterations=1)
        combined_mask = binary_closing(combined_mask, iterations=2)
        
        filtered_depth = depth_map.copy()
        filtered_depth[~combined_mask] = 0.0
        
        return filtered_depth
    
    def separate_persons(self, point_cloud: np.ndarray,
                         min_cluster_size: int = 50) -> list:
        """
        分离多个人体点云
        
        Args:
            point_cloud: (N, 3) 3D点云
            min_cluster_size: 最小聚类大小
            
        Returns:
            clusters: list of (M, 3) arrays, 每个元素为一个人的点云
        """
        from sklearn.cluster import DBSCAN
        
        if len(point_cloud) < min_cluster_size:
            return [point_cloud] if len(point_cloud) > 0 else []
        
        # DBSCAN 聚类
        clustering = DBSCAN(
            eps=0.15,  # 15cm邻域
            min_samples=min_cluster_size // 5
        ).fit(point_cloud)
        
        clusters = []
        for label in set(clustering.labels_):
            if label == -1:  # 噪声
                continue
            mask = clustering.labels_ == label
            clusters.append(point_cloud[mask])
        
        return clusters


# 测试
if __name__ == "__main__":
    fusion = GeometricPerspectiveFusion()
    
    # 模拟深度图和热图像（16x16）
    depth_map = np.random.uniform(0.3, 2.0, (16, 16))
    thermal = np.random.uniform(20, 35, (16, 16))
    thermal[5:10, 5:10] = 34.0  # 人体区域
    
    filtered = fusion.suppress_interference(depth_map, thermal)
    print(f"过滤前有效像素: {(depth_map > 0).sum()}")
    print(f"过滤后有效像素: {(filtered > 0).sum()}")
    
    # 模拟多人点云
    person1 = np.random.randn(100, 3) + np.array([0, 0, 1.0])
    person2 = np.random.randn(100, 3) + np.array([0.5, 0, 1.5])
    point_cloud = np.vstack([person1, person2])
    
    clusters = fusion.separate_persons(point_cloud)
    print(f"检测到 {len(clusters)} 个人体聚类")
    for i, c in enumerate(clusters):
        print(f"  人物 {i+1}: {len(c)} 个点, 中心: {c.mean(axis=0)}")

# 下游任务：跌倒检测示例
class FallDetector:
    """
    基于 TAP3D 点云的跌倒检测
    
    通过分析人体点云的高度变化率和水平面积变化判断跌倒
    """
    
    def __init__(self, fall_threshold: float = 0.3, 
                 time_window: int = 10):
        self.fall_threshold = fall_threshold  # 高度变化阈值(m/s)
        self.time_window = time_window
    
    def detect(self, point_clouds: list) -> dict:
        """
        Args:
            point_clouds: 时间序列点云列表，每个 (N, 3)
            
        Returns:
            result: {'is_fall': bool, 'confidence': float}
        """
        if len(point_clouds) < self.time_window:
            return {'is_fall': False, 'confidence': 0.0}
        
        heights = []
        widths = []
        for pc in point_clouds[-self.time_window:]:
            if len(pc) > 0:
                heights.append(pc[:, 1].max() - pc[:, 1].min())  # 垂直高度
                widths.append(np.sqrt(pc[:, 0].max()**2 + pc[:, 2].max()**2))  # 水平宽度
        
        if len(heights) < 2:
            return {'is_fall': False, 'confidence': 0.0}
        
        # 高度变化率
        height_rate = np.diff(heights) / 0.1  # 假设10fps
        
        # 跌倒特征：高度快速下降 + 水平面积增加
        fall_score = np.mean(np.abs(height_rate) if len(height_rate) > 0 else [0])
        is_fall = fall_score > self.fall_threshold
        
        return {
            'is_fall': is_fall,
            'confidence': min(fall_score / self.fall_threshold, 1.0),
            'height_change_rate': float(fall_score)
        }


# 测试
if __name__ == "__main__":
    detector = FallDetector()
    
    # 模拟正常站立点云序列
    normal_pcs = [np.random.randn(50, 3) * 0.2 + np.array([0, 1.0, 1.0]) 
                  for _ in range(10)]
    
    # 模拟跌倒点云序列（高度逐渐下降）
    fall_pcs = []
    for i in range(10):
        h = 1.0 - i * 0.08  # 高度逐渐下降
        fall_pcs.append(np.random.randn(50, 3) * 0.2 + np.array([0, h, 1.0]))
    
    print("正常活动:", detector.detect(normal_pcs))
    print("跌倒事件:", detector.detect(fall_pcs))

# 座舱 CPD 应用示例
class CabinCPD_Detector:
    """
    基于 TAP3D 的座舱儿童遗留检测
    
    优势：
    - 完全被动式：无辐射、无摄像头（隐私友好）
    - 成本极低：$5-15 热阵列传感器
    - 穿透能力强：可检测毯下/座椅后的儿童
    
    部署方案：
    - 顶棚安装：1个 32×24 热阵列
    - 覆盖范围：前排+后排座椅
    - 功耗：<0.5W（适合停车后持续监测）
    """
    
    def __init__(self, sensor_resolution=(32, 24)):
        self.thermal_model = ForwardThermalModel(sensor_resolution=sensor_resolution)
        self.fusion = GeometricPerspectiveFusion()
        self.fall_detector = FallDetector()
    
    def detect_child_presence(self, thermal_frame: np.ndarray) -> dict:
        """
        检测座舱内是否有儿童遗留
        
        Args:
            thermal_frame: (32, 24) 热阵列读数
            
        Returns:
            result: {'has_child': bool, 'location': tuple, 'confidence': float}
        """
        # 环境温度基线
        ambient = np.median(thermal_frame)
        
        # 人体温度区域
        human_mask = (thermal_frame > ambient + 2.0) & (thermal_frame < 40.0)
        
        # 如果检测到人体热源
        if human_mask.sum() < 5:
            return {'has_child': False, 'location': None, 'confidence': 0.0}
        
        # 估计热源位置
        rows, cols = np.where(human_mask)
        center_r = int(np.mean(rows))
        center_c = int(np.mean(cols))
        
        # 估计尺寸（区分成人/儿童）
        area = human_mask.sum()
        is_child = area < (self.thermal_model.resolution[0] * 
                          self.thermal_model.resolution[1] * 0.05)  # <5%面积
        
        return {
            'has_child': True,
            'location': (center_r, center_c),
            'is_child': is_child,
            'confidence': min(area / 50.0, 1.0),
            'estimated_area_cm2': area * 4  # 粗略估计
        }


# 测试
if __name__ == "__main__":
    detector = CabinCPD_Detector(sensor_resolution=(32, 24))
    
    # 模拟空座舱
    empty_cabin = np.full((32, 24), 25.0)  # 25°C环境温度
    print("空座舱:", detector.detect_child_presence(empty_cabin))
    
    # 模拟有儿童
    child_cabin = np.full((32, 24), 25.0)
    child_cabin[15:20, 10:14] = 34.0  # 儿童热源（小面积）
    print("有儿童:", detector.detect_child_presence(child_cabin))
    
    # 模拟成人
    adult_cabin = np.full((32, 24), 25.0)
    adult_cabin[10:22, 8:16] = 33.0  # 成人热源（大面积）
    print("有成人:", detector.detect_child_presence(adult_cabin))