from enum import Enum
from typing import List, Tuple, Optional
from dataclasses import dataclass

class RelationOperator(Enum):
    """Scenic DSL中的空间关系算子"""
    AHEAD_OF = "ahead of"
    BEHIND = "behind"
    LEFT_OF = "left of"
    RIGHT_OF = "right of"
    AT = "at"  # 绝对位置（最简单）

class FacingOperator(Enum):
    """朝向关系算子"""
    AWAY_FROM = "away from"
    TOWARD = "toward"
    SAME_DIRECTION = "same direction as"
    OPPOSITE = "opposite to"

@dataclass
class EntityDeclaration:
    """实体声明"""
    entity_id: str           # car1, car2, ...
    entity_type: str         # Car, Pedestrian, ...
    reference_entity: str    # 参照实体
    relation: RelationOperator  # 空间关系
    facing: FacingOperator   # 朝向关系
    distance_range: Tuple[float, float]  # 距离范围

@dataclass
class SceneProgram:
    """场景程序"""
    declarations: List[EntityDeclaration]
    
    def to_scenic_code(self) -> str:
        """生成Scenic代码"""
        lines = ["ego = new Car on 0@0, facing 0 deg"]
        for d in self.declarations:
            line = f"{d.entity_id} = new {d.entity_type} "
            line += f"{d.relation.value} {d.reference_entity} "
            line += f"by Range({d.distance_range[0]}, {d.distance_range[1]}), "
            line += f"facing {d.facing.value} {d.reference_entity}"
            lines.append(line)
        return "\n".join(lines)


class StructuralSearch:
    """
    R2SGEN离散结构搜索
    
    搜索空间：
    1. 实体声明顺序（O(N!)排列）
    2. 参照实体选择（每个实体可引用任意已声明实体）
    3. 关系算子选择（ahead/behind/left/right）
    
    通过两种剪枝策略加速：
    - BFS + Common Prefix Abstraction
    - DFS + Branch-and-Bound
    """
    
    def __init__(self, entities: List[dict]):
        """
        Args:
            entities: 真实场景实体列表
                [{'id': 'car1', 'type': 'Car', 'pos': (3,8), 'heading': 0}, ...]
        """
        self.entities = entities
        self.n = len(entities)
        self.best_program = None
        self.best_cost = float('inf')
        
    def search_bfs(self) -> Optional[SceneProgram]:
        """BFS + Common Prefix Abstraction Pruning"""
        from collections import deque
        
        # 搜索状态：(声明序列, 已用实体, 当前成本)
        initial_state = ([], set(), 0)
        queue = deque([initial_state])
        
        while queue:
            declarations, used, cost = queue.popleft()
            
            # 剪枝：成本已超过最优
            if cost >= self.best_cost:
                continue
            
            # 终止条件：所有实体已声明
            if len(declarations) == self.n:
                if cost < self.best_cost:
                    self.best_cost = cost
                    self.best_program = SceneProgram(declarations)
                continue
            
            # Common Prefix Abstraction 剪枝
            prefix = tuple(d.entity_id for d in declarations)
            if self._is_dominated_prefix(prefix, cost):
                continue
            
            # 扩展：为下一个实体选择参照和关系
            for entity in self.entities:
                if entity['id'] in used:
                    continue
                
                for ref in [d.entity_id for d in declarations] + ['ego']:
                    for rel in RelationOperator:
                        for face in FacingOperator:
                            new_decl = EntityDeclaration(
                                entity_id=entity['id'],
                                entity_type=entity['type'],
                                reference_entity=ref,
                                relation=rel,
                                facing=face,
                                distance_range=(1, 10)
                            )
                            
                            # SMT验证（连续部分）
                            if self._verify_smt(new_decl, declarations):
                                new_cost = cost + self._cost(new_decl)
                                new_used = used | {entity['id']}
                                queue.append(
                                    (declarations + [new_decl], new_used, new_cost)
                                )
            
            # 记录前缀以供剪枝
            self._record_prefix(prefix, cost)
        
        return self.best_program
    
    def search_dfs(self) -> Optional[SceneProgram]:
        """DFS + Branch-and-Bound Pruning"""
        
        def dfs_recursive(declarations, used, cost):
            if cost >= self.best_cost:
                return  # Branch-and-Bound 剪枝
            
            if len(declarations) == self.n:
                if cost < self.best_cost:
                    self.best_cost = cost
                    self.best_program = SceneProgram(declarations)
                return
            
            for entity in self.entities:
                if entity['id'] in used:
                    continue
                
                for ref in [d.entity_id for d in declarations] + ['ego']:
                    for rel in RelationOperator:
                        for face in FacingOperator:
                            new_decl = EntityDeclaration(
                                entity_id=entity['id'],
                                entity_type=entity['type'],
                                reference_entity=ref,
                                relation=rel,
                                facing=face,
                                distance_range=(1, 10)
                            )
                            
                            if self._verify_smt(new_decl, declarations):
                                new_cost = cost + self._cost(new_decl)
                                dfs_recursive(
                                    declarations + [new_decl],
                                    used | {entity['id']},
                                    new_cost
                                )
        
        dfs_recursive([], set(), 0)
        return self.best_program
    
    def _verify_smt(self, new_decl, existing) -> bool:
        """SMT约束验证（简化版）"""
        # 实际使用Z3 SMT solver
        # 这里简化为几何验证
        entity = next(e for e in self.entities 
                      if e['id'] == new_decl.entity_id)
        ref_entity = next(
            (e for e in self.entities if e['id'] == new_decl.reference_entity),
            {'pos': (0, 0), 'heading': 0}  # ego
        )
        
        dx = entity['pos'][0] - ref_entity['pos'][0]
        dy = entity['pos'][1] - ref_entity['pos'][1]
        dist = (dx**2 + dy**2) ** 0.5
        
        # 验证距离范围
        if not (1 <= dist <= 10):
            return False
        
        # 验证关系算子
        if new_decl.relation == RelationOperator.AHEAD_OF:
            if dy <= 0:  # ahead = y方向前方
                return False
        elif new_decl.relation == RelationOperator.BEHIND:
            if dy >= 0:
                return False
        elif new_decl.relation == RelationOperator.LEFT_OF:
            if dx >= 0:
                return False
        elif new_decl.relation == RelationOperator.RIGHT_OF:
            if dx <= 0:
                return False
        
        return True
    
    def _cost(self, decl: EntityDeclaration) -> float:
        """程序成本（更简洁=更低成本）"""
        cost = 0
        # 绝对坐标成本最高
        if decl.relation == RelationOperator.AT:
            cost += 10
        # 关系描述成本较低
        else:
            cost += 1
        # 参照ego比参照其他实体成本低
        if decl.reference_entity == 'ego':
            cost += 0.5
        else:
            cost += 1.0
        return cost
    
    def _is_dominated_prefix(self, prefix, cost):
        """Common Prefix Abstraction 剪枝"""
        # 如果相同前缀已被更低成本探索过，剪枝
        return self._prefix_cache.get(prefix, float('inf')) <= cost
    
    def _record_prefix(self, prefix, cost):
        if not hasattr(self, '_prefix_cache'):
            self._prefix_cache = {}
        if prefix not in self._prefix_cache or cost < self._prefix_cache[prefix]:
            self._prefix_cache[prefix] = cost


# 测试
if __name__ == "__main__":
    print("=" * 60)
    print("R2SGEN 场景程序合成 - 测试")
    print("=" * 60)
    
    # nuScenes示例场景
    entities = [
        {'id': 'car1', 'type': 'Car', 'pos': (3, 8), 'heading': 0},
        {'id': 'car2', 'type': 'Car', 'pos': (3, 5), 'heading': 0},
        {'id': 'car3', 'type': 'Car', 'pos': (0, 5), 'heading': 0},
    ]
    
    # BFS搜索
    search = StructuralSearch(entities)
    
    print("\nBFS搜索...")
    program = search.search_bfs()
    if program:
        print(f"最优成本: {search.best_cost:.2f}")
        print(f"\n生成的Scenic程序:")
        print(program.to_scenic_code())
    
    # DFS搜索
    search2 = StructuralSearch(entities)
    print("\nDFS搜索...")
    program2 = search2.search_dfs()
    if program2:
        print(f"最优成本: {search2.best_cost:.2f}")
        print(f"\n生成的Scenic程序:")
        print(program2.to_scenic_code())

from typing import List, Tuple

class SMTConstraintVerifier:
    """
    SMT约束验证器（连续几何部分）
    
    将Scenic DSL的几何语义编码为SMT约束
    使用Z3 solver验证程序一致性
    
    约束类型：
    1. 距离约束：|pos_a - pos_b| in [min, max]
    2. 方向约束：heading_a - heading_b in [angle_min, angle_max]
    3. 可见性约束：ego can see entity
    4. 碰撞约束：entities don't overlap
    """
    
    def __init__(self):
        # 简化版：直接几何验证
        # 实际使用Z3 SMT solver
        pass
    
    def verify_distance(self, pos_a: Tuple[float, float],
                        pos_b: Tuple[float, float],
                        dist_range: Tuple[float, float]) -> bool:
        """验证距离约束"""
        dist = ((pos_a[0]-pos_b[0])**2 + (pos_a[1]-pos_b[1])**2)**0.5
        return dist_range[0] <= dist <= dist_range[1]
    
    def verify_direction(self, heading_a: float, heading_b: float,
                         angle_range: Tuple[float, float]) -> bool:
        """验证方向约束"""
        diff = (heading_a - heading_b + 180) % 360 - 180
        return angle_range[0] <= diff <= angle_range[1]
    
    def verify_relation(self, pos_entity, pos_ref, relation, 
                        dist_range=(1, 10)) -> bool:
        """验证空间关系"""
        dx = pos_entity[0] - pos_ref[0]
        dy = pos_entity[1] - pos_ref[1]
        dist = (dx**2 + dy**2)**0.5
        
        if not self.verify_distance(pos_entity, pos_ref, dist_range):
            return False
        
        if relation == "ahead of":
            return dy > 0  # 前方
        elif relation == "behind":
            return dy < 0  # 后方
        elif relation == "left of":
            return dx < 0  # 左侧
        elif relation == "right of":
            return dx > 0  # 右侧
        return True
    
    def verify_no_collision(self, positions: List[Tuple[float, float]],
                            min_distance: float = 2.0) -> bool:
        """验证无碰撞"""
        for i in range(len(positions)):
            for j in range(i+1, len(positions)):
                dist = ((positions[i][0]-positions[j][0])**2 + 
                        (positions[i][1]-positions[j][1])**2)**0.5
                if dist < min_distance:
                    return False
        return True


# 测试
if __name__ == "__main__":
    verifier = SMTConstraintVerifier()
    
    # 测试关系验证
    print("关系验证测试:")
    print(f"  (3,8) ahead of (0,0): {verifier.verify_relation((3,8), (0,0), 'ahead of')}")
    print(f"  (3,5) right of (0,5): {verifier.verify_relation((3,5), (0,5), 'right of')}")
    print(f"  (0,5) ahead of (0,0): {verifier.verify_relation((0,5), (0,0), 'ahead of')}")

# 座舱场景约束示例
cabin_constraints = {
    "空间约束": {
        "驾驶员位置": "fixed at driver seat",
        "乘客位置": "Range(0.3, 0.8) m from seat back",
        "安全带路径": "across shoulder and lap",
        "手部可达范围": "within arm reach sphere R=0.6m"
    },
    "物理约束": {
        "光照": "Range(10, 10000) lux",
        "材质反射率": "Range(0.1, 0.9)",
        "遮挡关系": "driver body occludes rear seat"
    },
    "行为约束": {
        "头部姿态": "pitch: [-30, 30]°, yaw: [-60, 60]°",
        "眼动模式": "saccade + fixation + smooth pursuit",
        "手势区域": "center console / steering wheel / window"
    },
    "时间约束": {
        "行为持续时间": "fixation: 0.2-2s, saccade: 20-80ms",
        "行为转移概率": "P(phone→driving) = 0.15"
    }
}