"""
CW-Net: Concept-Wrapper Network
论文: Nature, 2026

核心思想:
1. 保留预训练黑盒模型的特征提取层
2. 替换最终决策层为概念分类器+新奖励层
3. 概念分类器输出人类可理解标签
4. 新奖励层基于概念做决策

适用: 任意预训练DNN的即插即用改造
"""

import torch
import torch.nn as nn
from typing import List

class ConceptClassifier(nn.Module):
    """
    概念分类器
    
    将黑盒模型的嵌入zi映射到人类可理解概念
    如: 靠近停止车辆/接近自行车/跟车过近
    """
    def __init__(self, embed_dim: int = 256,
                 concepts: List[str] = None):
        super().__init__()
        
        if concepts is None:
            concepts = [
                'ASV',      # Approaching Stopped Vehicle
                'CLOSE',    # Close to another vehicle
                'BIKE',     # Close to cyclist
                'LANE_CHG', # Lane change available
                'TURN',     # Approaching intersection
                'PED',      # Pedestrian nearby
                'SPEED',    # Speed mismatch
                'CROWD',    # Dense traffic
            ]
        
        self.num_concepts = len(concepts)
        self.concept_names = concepts
        
        # 概念分类头（多标签）
        self.classifier = nn.Sequential(
            nn.Linear(embed_dim, 128),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(128, self.num_concepts)
        )
    
    def forward(self, z: torch.Tensor) -> torch.Tensor:
        """
        Args:
            z: 场景-轨迹嵌入 (B, embed_dim)
        Returns:
            concept_logits: 概念logits (B, num_concepts)
        """
        return self.classifier(z)
    
    def explain(self, z: torch.Tensor) -> dict:
        """生成人类可读解释"""
        logits = self.forward(z)
        probs = torch.sigmoid(logits)
        
        active_concepts = []
        for i, p in enumerate(probs[0]):
            if p > 0.5:
                active_concepts.append({
                    'concept': self.concept_names[i],
                    'probability': p.item()
                })
        
        return {
            'active_concepts': active_concepts,
            'explanation': self._generate_text(active_concepts)
        }
    
    def _generate_text(self, concepts: list) -> str:
        """概念→自然语言解释"""
        if not concepts:
            return "Normal driving condition"
        
        parts = []
        for c in concepts:
            name = c['concept']
            if name == 'ASV':
                parts.append("approaching a stopped vehicle")
            elif name == 'CLOSE':
                parts.append("close to another vehicle")
            elif name == 'BIKE':
                parts.append("a cyclist is nearby")
            elif name == 'LANE_CHG':
                parts.append("lane change is available")
            elif name == 'TURN':
                parts.append("approaching an intersection")
            elif name == 'PED':
                parts.append("a pedestrian is nearby")
        
        return "AV is " + ", ".join(parts) + " → adjusting trajectory"


class CWNet(nn.Module):
    """
    Concept-Wrapper Network
    
    包装任意预训练planner，添加可解释性
    """
    def __init__(self, pretrained_planner: nn.Module,
                 embed_dim: int = 256):
        super().__init__()
        
        # 冻结预训练模型的特征提取层
        self.planner = pretrained_planner
        for param in self.planner.parameters():
            param.requires_grad = False
        
        # 概念分类器（可训练）
        self.concept_classifier = ConceptClassifier(embed_dim)
        
        # 新奖励层（可训练）
        self.new_reward = nn.Sequential(
            nn.Linear(self.concept_classifier.num_concepts, 64),
            nn.ReLU(),
            nn.Linear(64, 1)
        )
    
    def forward(self, scene: torch.Tensor,
                trajectories: torch.Tensor) -> dict:
        """
        Args:
            scene: 场景上下文
            trajectories: 候选轨迹集
        
        Returns:
            reward: 轨迹奖励
            concepts: 概念分类
            explanation: 人类可读解释
        """
        # 使用冻结的planner提取嵌入
        with torch.no_grad():
            h = self.planner.scene_encoder(scene)
            embeddings = []
            for traj in trajectories:
                z = self.planner.encoder(h, traj)
                embeddings.append(z)
            z = torch.stack(embeddings).squeeze(1)
        
        # 概念分类
        concept_logits = self.concept_classifier(z)
        
        # 基于概念的新奖励
        rewards = self.new_reward(torch.sigmoid(concept_logits))
        
        # 解释
        best_idx = rewards.argmax()
        explanation = self.concept_classifier.explain(
            z[best_idx:best_idx+1]
        )
        
        return {
            'rewards': rewards,
            'concept_logits': concept_logits,
            'explanation': explanation
        }


# 测试
if __name__ == "__main__":
    # 模拟概念分类
    classifier = ConceptClassifier()
    z = torch.randn(1, 256)
    result = classifier.explain(z)
    
    print("概念分类结果:")
    for c in result['active_concepts']:
        print(f"  {c['concept']}: {c['probability']:.2f}")
    print(f"\n解释: {result['explanation']}")