import numpy as np
from typing import Tuple

class PhysicsBoundedIntegralProcessing:
    """
    PBIP: 物理约束积分处理
    
    核心思想: mmWave人体反射在空间上紧凑
    → 只处理物理占用的区域
    → 用积分查询代替滑动窗口（常数时间）
    
    论文: arXiv 2608.14176
    """
    
    def __init__(self, radar_config):
        self.num_range_bins = radar_config['range_bins']  # 64
        self.num_angle_bins = radar_config['angle_bins']  # 128
        self.num_doppler_bins = radar_config['doppler_bins']  # 16
        
    def compute_integral_image(self, rad_tensor: np.ndarray) -> np.ndarray:
        """
        计算积分图（常数时间区域查询）
        
        Args:
            rad_tensor: (range_bins, angle_bins, doppler_bins)
            
        Returns:
            integral: 积分图，可用于O(1)区域求和
        """
        # 沿每个维度累积求和
        integral = np.cumsum(np.cumsum(np.cumsum(
            rad_tensor.astype(np.float64), axis=0
        ), axis=1), axis=2)
        
        # 前面补零以便查询
        integral = np.pad(integral, ((1,0),(1,0),(1,0)), mode='constant')
        
        return integral
    
    def bounded_query(self, integral: np.ndarray, 
                      r_range: Tuple[int, int],
                      r_angle: Tuple[int, int],
                      r_doppler: Tuple[int, int]) -> float:
        """
        常数时间区域求和（3D积分查询）
        
        Args:
            integral: 积分图
            r_range: (start, end) range bin范围
            r_angle: (start, end) angle bin范围  
            r_doppler: (start, end) doppler bin范围
            
        Returns:
            sum: 该区域的能量总和
        """
        r0, r1 = r_range
        a0, a1 = r_angle
        d0, d1 = r_doppler
        
        # 3D inclusion-exclusion (8 corners)
        S = integral
        
        result = (S[r1,a1,d1] - S[r0,a1,d1] - S[r1,a0,d1] + S[r0,a0,d1]
                 - (S[r1,a1,d0] - S[r0,a1,d0] - S[r1,a0,d0] + S[r0,a0,d0]))
        
        return float(result)
    
    def find_occupied_regions(self, rad_tensor: np.ndarray, 
                               threshold: float = 0.1) -> list:
        """
        找到物理占用区域（人体反射所在区域）
        
        基于物理约束: 人体反射在空间上紧凑
        """
        integral = self.compute_integral_image(rad_tensor)
        
        total_energy = self.bounded_query(
            integral, 
            (0, self.num_range_bins),
            (0, self.num_angle_bins),
            (0, self.num_doppler_bins)
        )
        
        if total_energy == 0:
            return []
        
        # 递归二分找到高能量区域
        regions = self._recursive_bisection(
            integral, 
            (0, self.num_range_bins),
            (0, self.num_angle_bins),
            (0, self.num_doppler_bins),
            threshold * total_energy
        )
        
        return regions
    
    def _recursive_bisection(self, integral, r, a, d, threshold):
        """递归二分寻找紧凑区域"""
        energy = self.bounded_query(integral, r, a, d)
        
        if energy < threshold:
            return []
        
        if (r[1]-r[0] <= 2 and a[1]-a[0] <= 2):
            return [(r, a, d)]
        
        # 沿最大维度二分
        regions = []
        dim_size = [r[1]-r[0], a[1]-a[0], d[1]-d[0]]
        max_dim = np.argmax(dim_size)
        
        if max_dim == 0:  # range
            mid = (r[0] + r[1]) // 2
            regions += self._recursive_bisection(integral, (r[0],mid), a, d, threshold)
            regions += self._recursive_bisection(integral, (mid,r[1]), a, d, threshold)
        elif max_dim == 1:  # angle
            mid = (a[0] + a[1]) // 2
            regions += self._recursive_bisection(integral, r, (a[0],mid), d, threshold)
            regions += self._recursive_bisection(integral, r, (mid,a[1]), d, threshold)
        else:  # doppler
            mid = (d[0] + d[1]) // 2
            regions += self._recursive_bisection(integral, r, a, (d[0],mid), threshold)
            regions += self._recursive_bisection(integral, r, a, (mid,d[1]), threshold)
        
        return regions


class DeadlineAwareOperationProfiles:
    """
    DAOP: 截止期感知操作配置
    
    动态选择满足截止期的最高精度配置
    """
    
    PROFILES = [
        {'name': 'high', 'resolution': 'full', 'wcet_ms': 45, 'accuracy': 'best'},
        {'name': 'medium', 'resolution': 'half', 'wcet_ms': 28, 'accuracy': 'good'},
        {'name': 'low', 'resolution': 'quarter', 'wcet_ms': 15, 'accuracy': 'fair'},
        {'name': 'minimal', 'resolution': 'eighth', 'wcet_ms': 8, 'accuracy': 'basic'},
    ]
    
    def select_profile(self, deadline_ms: float) -> dict:
        """选择满足截止期的最高精度配置"""
        for profile in sorted(self.PROFILES, key=lambda x: -x['wcet_ms']):
            if profile['wcet_ms'] <= deadline_ms:
                return profile
        return self.PROFILES[-1]  # 最快配置


# 测试
if __name__ == "__main__":
    config = {
        'range_bins': 64,
        'angle_bins': 128,
        'doppler_bins': 16
    }
    
    pbip = PhysicsBoundedIntegralProcessing(config)
    
    # 模拟mmWave RAD张量
    rad_tensor = np.zeros((64, 128, 16), dtype=np.float32)
    # 模拟人体在range=20-30, angle=50-70, doppler=5-10
    rad_tensor[20:30, 50:70, 5:10] = np.random.randn(10, 20, 5) * 10
    
    integral = pbip.compute_integral_image(rad_tensor)
    
    # 查询人体区域能量
    energy = pbip.bounded_query(integral, (20, 30), (50, 70), (5, 10))
    total = pbip.bounded_query(integral, (0, 64), (0, 128), (0, 16))
    
    print("=== PRISM PBIP 测试 ===")
    print(f"RAD张量大小: {rad_tensor.shape}")
    print(f"人体区域能量: {energy:.2f}")
    print(f"总能量: {total:.2f}")
    print(f"人体能量占比: {energy/total*100:.1f}%")
    print(f"空bin占比: {(1-energy/total)*100:.1f}% → 可跳过")
    print()
    
    # DAOP测试
    daop = DeadlineAwareOperationProfiles()
    for deadline in [50, 30, 20, 10, 5]:
        profile = daop.select_profile(deadline)
        print(f"截止期{deadline}ms → 选择: {profile['name']} (WCET={profile['wcet_ms']}ms)")