"""
J-Space 核心概念的形式化定义
"""
import torch
import torch.nn as nn
import numpy as np
from typing import List, Tuple, Optional, Callable

class JSpaceDefinition:
    """
    J-Space 的形式化定义
    
    J-Space 是模型内部激活空间中一个低维子空间，满足：
    1. 可读性（Verbalizable）：其中的表征可以映射为自然语言概念
    2. 因果性（Causal）：对其中的表征进行干预会改变模型输出
    3. 稀疏性（Sparse）：仅占模型内部激活方差的6%-10%
    4. 广播性（Broadcast）：其中的信息可以被多个下游模块访问
    """
    
    def __init__(self, model: nn.Module, layer_idx: int):
        self.model = model
        self.layer_idx = layer_idx
        self.jacobian: Optional[torch.Tensor] = None
        self.j_space_directions: Optional[torch.Tensor] = None
        
    def compute_jacobian(self, input_ids: torch.Tensor,
                         target_token_ids: torch.Tensor) -> torch.Tensor:
        """
        计算雅可比矩阵：输出对中间层激活的偏导数
        
        J = ∂output / ∂h_layer
        
        其中 h_layer 是第 layer_idx 层的隐藏状态
        """
        # 注册前向钩子获取中间层激活
        intermediate_activation = None
        
        def hook_fn(module, input, output):
            nonlocal intermediate_activation
            intermediate_activation = output
        
        hook = self.model.layers[self.layer_idx].register_forward_hook(hook_fn)
        
        # 前向传播
        output = self.model(input_ids)
        
        # 选择目标token的logits
        target_logits = output.logits[0, -1, target_token_ids]
        
        # 计算雅可比矩阵
        self.jacobian = torch.autograd.functional.jacobian(
            lambda h: self._compute_logits_from_hidden(h, target_token_ids),
            intermediate_activation
        )
        
        hook.remove()
        return self.jacobian
    
    def _compute_logits_from_hidden(self, hidden_state: torch.Tensor,
                                     target_token_ids: torch.Tensor) -> torch.Tensor:
        """从隐藏状态计算目标token的logits"""
        # 简化版：假设从隐藏状态到logits的映射
        return self.model.lm_head(hidden_state)[:, target_token_ids]
    
    def extract_j_space(self, activation: torch.Tensor,
                        threshold: float = 0.01) -> torch.Tensor:
        """
        从激活中提取J-Space方向
        
        通过SVD分解雅可比矩阵，保留贡献最大的奇异向量
        """
        if self.jacobian is None:
            raise ValueError("请先调用 compute_jacobian()")
        
        # SVD分解
        U, S, Vh = torch.linalg.svd(self.jacobian, full_matrices=False)
        
        # 保留超过阈值的奇异值对应的方向
        significant_dirs = (S > threshold * S.max()).sum().item()
        self.j_space_directions = Vh[:significant_dirs, :]
        
        # 将激活投影到J-Space
        j_space_activation = self.j_space_directions @ activation
        
        return j_space_activation
    
    def verbalize_activation(self, activation_projection: torch.Tensor,
                             tokenizer: Callable) -> List[str]:
        """
        将J-Space中的激活解码为可读的概念
        
        通过查找与投影方向最接近的词嵌入来实现
        """
        # 获取模型词嵌入矩阵
        embedding_matrix = self.model.get_input_embeddings().weight
        
        # 对每个J-Space方向，找到最接近的词
        concepts = []
        for i in range(activation_projection.shape[0]):
            direction = self.j_space_directions[i]
            similarities = embedding_matrix @ direction
            top_k = torch.topk(similarities, k=5).indices
            concepts.append([tokenizer.decode(idx) for idx in top_k])
        
        return concepts

"""
J-lens (Jacobian Lens) 的完整实现
"""
class JacobianLens:
    """
    雅可比透镜：一种无需额外训练的模型内部状态读取工具
    
    核心思想：不是看模型"下一步最可能输出什么词"，
    而是看"某一层某个内部激活，平均来看会让模型在未来更容易说出哪些词"
    """
    
    def __init__(self, model, tokenizer, layers: List[int]):
        self.model = model
        self.tokenizer = tokenizer
        self.layers = layers
        self.activations: dict = {}
        self._register_hooks()
    
    def _register_hooks(self):
        """注册钩子以捕获中间层激活"""
        def make_hook_fn(layer_idx):
            def hook_fn(module, input, output):
                self.activations[layer_idx] = output[0].detach() if isinstance(output, tuple) else output.detach()
            return hook_fn
        
        for layer_idx in self.layers:
            self.model.model.layers[layer_idx].register_forward_hook(
                make_hook_fn(layer_idx)
            )
    
    def compute_jacobian_at_layer(self, layer_idx: int,
                                   input_ids: torch.Tensor,
                                   target_positions: List[int]) -> torch.Tensor:
        """
        计算指定层的雅可比矩阵
        
        对于输入序列中的每个位置，计算：
        J_{i,j} = ∂logit_j / ∂h_i
        
        其中 h_i 是第i个位置的隐藏状态
        logit_j 是第j个位置的logit
        """
        batch_size, seq_len = input_ids.shape
        hidden_dim = self.model.config.hidden_size
        
        # 确保模型处于评估模式
        self.model.eval()
        
        # 清空之前的激活
        self.activations.clear()
        
        # 前向传播（带钩子）
        with torch.no_grad():
            outputs = self.model(input_ids)
        
        hidden_states = self.activations[layer_idx]  # [batch, seq_len, hidden_dim]
        logits = outputs.logits  # [batch, seq_len, vocab_size]
        
        # 初始化雅可比矩阵
        jacobian = torch.zeros(seq_len, hidden_dim, len(target_positions))
        
        # 对每个目标位置计算雅可比
        for t_idx, pos in enumerate(target_positions):
            target_logit = logits[0, pos, :]
            
            # 使用自动微分计算雅可比
            for i in range(seq_len):
                for j in range(hidden_dim):
                    # 近似：∂logit_pos / ∂hidden[i,j]
                    epsilon = 1e-5
                    hidden_orig = hidden_states[0, i, j].clone()
                    
                    # 正向扰动
                    hidden_states[0, i, j] = hidden_orig + epsilon
                    with torch.no_grad():
                        output_plus = self.model(input_ids).logits[0, pos, :]
                    
                    # 负向扰动
                    hidden_states[0, i, j] = hidden_orig - epsilon
                    with torch.no_grad():
                        output_minus = self.model(input_ids).logits[0, pos, :]
                    
                    # 恢复原始值
                    hidden_states[0, i, j] = hidden_orig
                    
                    # 有限差分
                    jacobian[i, j, t_idx] = (output_plus - output_minus).norm().item() / (2 * epsilon)
        
        return jacobian
    
    def decode_j_space_content(self, layer_idx: int,
                                input_text: str,
                                top_k: int = 10) -> List[Tuple[str, float]]:
        """
        解码J-Space中的内容
        
        返回当前激活中最可能被"读出"的概念列表
        """
        input_ids = self.tokenizer.encode(input_text, return_tensors='pt')
        
        # 获取该层的隐藏状态
        self.activations.clear()
        with torch.no_grad():
            self.model(input_ids)
        
        hidden = self.activations[layer_idx][0]  # [seq_len, hidden_dim]
        
        # 计算雅可比（针对最后一个位置的logits）
        jacobian = self.compute_jacobian_at_layer(layer_idx, input_ids, [-1])
        
        # 对每个位置，找到J-Space的主要方向
        concepts = []
        for pos in range(hidden.shape[0]):
            # 该位置的雅可比向量
            j_vec = jacobian[pos, :, 0]
            
            # 归一化
            j_vec = j_vec / (j_vec.norm() + 1e-8)
            
            # 查找词汇表中与这个方向最接近的词
            embedding = self.model.get_input_embeddings().weight
            similarities = embedding @ j_vec
            
            top_values, top_indices = torch.topk(similarities, k=top_k)
            
            for val, idx in zip(top_values, top_indices):
                token = self.tokenizer.decode(idx)
                if token and len(token.strip()) > 0:
                    concepts.append((token, val.item()))
        
        # 去重并排序
        seen = set()
        unique_concepts = []
        for concept, score in concepts:
            normalized = concept.strip().lower()
            if normalized not in seen and len(normalized) > 1:
                seen.add(normalized)
                unique_concepts.append((concept, score))
        
        unique_concepts.sort(key=lambda x: x[1], reverse=True)
        return unique_concepts[:top_k]


class JSpaceIntervention:
    """
    J-Space 干预实验框架
    
    通过修改J-Space中的表征来观察模型行为的变化，
    验证J-Space的因果效力
    """
    
    def __init__(self, model, lens: JacobianLens):
        self.model = model
        self.lens = lens
    
    def intervene_and_observe(self, input_text: str,
                               layer_idx: int,
                               source_concept: str,
                               target_concept: str) -> Dict:
        """
        干预实验：将J-Space中的source_concept替换为target_concept
        
        如果J-Space具有因果效力，替换后模型的输出应该改变
        """
        input_ids = self.lens.tokenizer.encode(input_text, return_tensors='pt')
        
        # Step 1: 获取原始输出
        with torch.no_grad():
            original_output = self.model.generate(input_ids, max_new_tokens=10)
        original_text = self.lens.tokenizer.decode(original_output[0])
        
        # Step 2: 找到source_concept在J-Space中的方向
        source_embed = self.model.get_input_embeddings().weight[
            self.lens.tokenizer.encode(source_concept, add_special_tokens=False)[0]
        ]
        target_embed = self.model.get_input_embeddings().weight[
            self.lens.tokenizer.encode(target_concept, add_special_tokens=False)[0]
        ]
        
        # Step 3: 计算替换方向
        intervention_direction = target_embed - source_embed
        intervention_direction = intervention_direction / (intervention_direction.norm() + 1e-8)
        
        # Step 4: 在J-Space中应用干预
        # 获取该层激活
        self.lens.activations.clear()
        with torch.no_grad():
            self.model(input_ids)
        
        hidden = self.lens.activations[layer_idx]
        
        # 计算雅可比，找到J-Space方向
        jacobian = self.lens.compute_jacobian_at_layer(layer_idx, input_ids, [-1])
        j_direction = jacobian[:, :, 0]  # [seq_len, hidden_dim]
        
        # 在J-Space方向上进行干预
        scale = 5.0  # 干预强度
        intervention = scale * intervention_direction.unsqueeze(0) * j_direction.norm(dim=1, keepdim=True)
        
        # 应用干预到最后一层
        hidden[0, -1, :] += intervention[-1, :]
        
        # Step 5: 获取干预后的输出
        with torch.no_grad():
            intervened_output = self.model.generate(input_ids, max_new_tokens=10)
        intervened_text = self.lens.tokenizer.decode(intervened_output[0])
        
        return {
            'original': original_text,
            'intervened': intervened_text,
            'source_concept': source_concept,
            'target_concept': target_concept,
            'layer_idx': layer_idx,
            'intervention_strength': scale
        }

"""
J-Space 五类核心实验验证
"""
class JSpaceValidationExperiment:
    """J-Space验证实验套件"""
    
    def __init__(self, model, lens: JacobianLens, intervention: JSpaceIntervention):
        self.model = model
        self.lens = lens
        self.intervention = intervention
    
    def experiment_1_reportability(self) -> Dict:
        """
        实验1：可报告性（Reportability）
        
        假设：J-Space中的内容应该可以被模型报告出来
        方法：问Claude"你在想什么？"，看其回答是否对应J-Space内容
        """
        results = []
        queries = [
            "会织网的动物有几条腿？",
            "法国的首都是什么？",
            "太阳系最大的行星是哪个？"
        ]
        
        for query in queries:
            # 读取J-Space内容
            j_space_concepts = self.lens.decode_j_space_content(
                layer_idx=15, input_text=query
            )
            
            # 获取模型回答
            input_ids = self.lens.tokenizer.encode(query, return_tensors='pt')
            with torch.no_grad():
                output = self.model.generate(input_ids, max_new_tokens=50)
            response = self.lens.tokenizer.decode(output[0])
            
            results.append({
                'query': query,
                'j_space_concepts': [c[0] for c in j_space_concepts[:5]],
                'model_response': response,
                'overlap': self._compute_overlap(j_space_concepts, response)
            })
        
        return {'reportability_results': results}
    
    def _compute_overlap(self, concepts: List[Tuple], response: str) -> float:
        """计算J-Space概念与模型回答的重叠度"""
        concept_words = set(c[0].lower() for c in concepts)
        response_words = set(response.lower().split())
        if len(concept_words) == 0:
            return 0.0
        return len(concept_words & response_words) / len(concept_words)
    
    def experiment_2_modulability(self) -> Dict:
        """
        实验2：可调节性（Modulability）
        
        假设：要求模型"在心里想X"，J-Space中X的表征应该被点亮
        """
        results = []
        concepts = ["巴黎", "爱因斯坦", "Python"]
        
        for concept in concepts:
            # 无提示条件
            baseline = self.lens.decode_j_space_content(
                layer_idx=15, input_text=f"介绍一下你自己"
            )
            
            # 有提示条件
            prompted = self.lens.decode_j_space_content(
                layer_idx=15, input_text=f"在心里想一下{concept}"
            )
            
            # 检查J-Space中是否出现了目标概念
            baseline_concepts = set(c[0].lower() for c in baseline)
            prompted_concepts = set(c[0].lower() for c in prompted)
            
            results.append({
                'concept': concept,
                'baseline_has_concept': concept.lower() in baseline_concepts,
                'prompted_has_concept': concept.lower() in prompted_concepts,
                'new_concepts': list(prompted_concepts - baseline_concepts)
            })
        
        return {'modulability_results': results}
    
    def experiment_3_causal_intervention(self) -> Dict:
        """
        实验3：因果干预（Causal Intervention）
        
        假设：修改J-Space中的表征会改变模型输出
        这是验证J-Space因果效力的核心实验
        """
        interventions = [
            {
                'input': "会织网的动物有几条腿？",
                'source': "蜘蛛",
                'target': "蚂蚁",
                'expected': "6"  # 8条腿→6条腿
            },
            {
                'input': "法国的首都是什么？",
                'source': "巴黎",
                'target': "伦敦",
                'expected_shift': True
            }
        ]
        
        results = []
        for intervention in interventions:
            result = self.intervention.intervene_and_observe(
                input_text=intervention['input'],
                layer_idx=15,
                source_concept=intervention['source'],
                target_concept=intervention['target']
            )
            results.append(result)
        
        return {'causal_intervention_results': results}
    
    def experiment_4_capacity_limitation(self) -> Dict:
        """
        实验4：容量限制（Capacity Limitation）
        
        假设：J-Space的容量有限，同时追踪多个概念时会饱和
        """
        # 逐步增加需要追踪的概念数量
        for num_concepts in [1, 2, 3, 5, 10]:
            prompt = "请同时思考以下概念：" + "、".join([f"概念{i}" for i in range(num_concepts)])
            j_space = self.lens.decode_j_space_content(layer_idx=15, input_text=prompt)
            print(f"概念数={num_concepts}, J-Space可读概念数={len(j_space)}")
    
    def experiment_5_broadcast_nature(self) -> Dict:
        """
        实验5：广播特性（Broadcast Nature）
        
        假设：J-Space中的信息可以被多个下游任务模块访问
        """
        # 在J-Space中注入一个概念，观察它对多个不同任务输出的影响
        pass


class TwoStepReasoningTest:
    """两步推理替换实验"""
    
    def __init__(self, model, tokenizer, lens: JacobianLens):
        self.model = model
        self.tokenizer = tokenizer
        self.lens = lens
    
    def run_test(self, input_text: str, 
                 intermediate_concept: str,
                 replacement_concept: str,
                 layer: int = 15) -> Dict:
        """
        两步推理替换实验
        
        例：输入"会织网的动物有几条腿？"
        - 中间概念："蜘蛛"
        - 替换为："蚂蚁"
        - 预期输出变化：8 → 6
        """
        # 获取原始输出
        input_ids = self.tokenizer.encode(input_text, return_tensors='pt')
        with torch.no_grad():
            orig_output = self.model.generate(input_ids, max_new_tokens=20)
        original_answer = self.tokenizer.decode(orig_output[0])
        
        # 找到中间概念和替换概念的嵌入
        source_id = self.tokenizer.encode(intermediate_concept, add_special_tokens=False)[0]
        target_id = self.tokenizer.encode(replacement_concept, add_special_tokens=False)[0]
        
        source_embed = self.model.get_input_embeddings().weight[source_id]
        target_embed = self.model.get_input_embeddings().weight[target_id]
        
        # 计算替换方向
        edit_direction = target_embed - source_embed
        
        # 在J-Space中应用替换
        self.lens.activations.clear()
        with torch.no_grad():
            self.model(input_ids)
        
        hidden = self.lens.activations[layer]
        # 应用干预
        hidden[0, -1, :] += edit_direction * 3.0
        
        # 获取干预后的输出
        class PatchedModel:
            def __init__(self, model, patched_hidden, layer_idx):
                self.model = model
                self.patched_hidden = patched_hidden
                self.layer_idx = layer_idx
                self.original_hidden = None
            
            def generate(self, input_ids, max_new_tokens=20):
                # 修改模型前向传播，替换指定层的隐藏状态
                original_forward = self.model.model.layers[self.layer_idx].forward
                
                def patched_forward(*args, **kwargs):
                    output = original_forward(*args, **kwargs)
                    if isinstance(output, tuple):
                        return (self.patched_hidden,) + output[1:]
                    return self.patched_hidden
                
                self.model.model.layers[self.layer_idx].forward = patched_forward
                result = self.model.generate(input_ids, max_new_tokens=max_new_tokens)
                self.model.model.layers[self.layer_idx].forward = original_forward
                return result
        
        patched_model = PatchedModel(self.model, hidden, layer)
        with torch.no_grad():
            new_output = patched_model.generate(input_ids, max_new_tokens=20)
        intervened_answer = self.tokenizer.decode(new_output[0])
        
        return {
            'input': input_text,
            'original_answer': original_answer,
            'intervened_answer': intervened_answer,
            'intermediate_concept': intermediate_concept,
            'replacement_concept': replacement_concept,
            'answer_changed': original_answer != intervened_answer
        }

// J-Space 实验数据分析
package main

import (
	"fmt"
	"math"
)

type ExperimentResult struct {
	ExperimentName string
	TotalTrials    int
	Successful     int
	SuccessRate    float64
	Description    string
}

func analyzeJSpaceExperiments() []ExperimentResult {
	results := []ExperimentResult{
		{
			ExperimentName: "两步推理替换",
			TotalTrials:    100,
			Successful:     65,
			Description:    "替换中间概念（如蜘蛛→蚂蚁），观察答案是否随之改变",
		},
		{
			ExperimentName: "概念共享性测试",
			TotalTrials:    192,
			Successful:     76,
			Description:    "替换国家概念（法国→中国），观察首都/语言/大洲/货币四维回答是否同步变化",
		},
		{
			ExperimentName: "J-Space抑制-推理任务",
			TotalTrials:    50,
			Successful:     8,
			Description:    "抑制J-Space后，深度推理任务的正确率",
		},
		{
			ExperimentName: "J-Space抑制-检索任务",
			TotalTrials:    50,
			Successful:     47,
			Description:    "抑制J-Space后，常规信息检索任务的正确率（基本不受影响）",
		},
	}
	
	for i, r := range results {
		results[i].SuccessRate = math.Round(float64(r.Successful)/float64(r.TotalTrials)*1000) / 10
	}
	
	return results
}

func main() {
	results := analyzeJSpaceExperiments()
	
	fmt.Println("J-Space 实验验证数据分析")
	fmt.Println("=" 复 70)
	fmt.Printf("%-20s %-8s %-8s %-10s %-25s\n", "实验名称", "总次数", "成功", "成功率", "说明")
	fmt.Println("- 复 70)
	
	for _, r := range results {
		fmt.Printf("%-20s %-8d %-8d %-8.1f%% %-25s\n",
			r.ExperimentName, r.TotalTrials, r.Successful,
			r.SuccessRate, r.Description)
	}
	
	fmt.Println("\n关键发现:")
	fmt.Println("1. J-Space的干预成功率在40%-70%之间，远非精确开关")
	fmt.Println("2. 抑制J-Space对推理任务影响巨大（正确率从~90%降至~16%）")
	fmt.Println("3. 常规检索任务基本不受影响（正确率仍保持~94%）")
	fmt.Println("4. 证实J-Space是推理专用子空间，而非信息检索空间")
}

"""
J-Space 与 全局工作空间理论的功能类比
"""
class GlobalWorkspaceAnalogy:
    """
    J-Space 与人类全局工作空间的功能类比
    
    功能维度       | 人类全局工作空间    | J-Space
    --------------|-------------------|-------------------
    容量限制       | 4-7个信息块        | 6-10%的激活方差
    广播特性       | 信息可被多脑区访问  | 信息可被多下游模块读取
    选择性参与     | 仅参与需要意识的任务 | 仅参与灵活推理任务
    可报告性       | 可被口头报告       | 可通过J-lens解码
    可干预性       | 改变内容改变行为   | 修改表征改变输出
    """
    
    def __init__(self):
        self.parallels = [
            {
                'dimension': '容量限制',
                'human': '工作记忆容量约4-7个信息块',
                'j_space': '仅占模型激活方差的6%-10%',
                'similarity': '两者都只有有限的信息处理带宽'
            },
            {
                'dimension': '广播特性',
                'human': '进入意识的信息可被多个脑区并行访问',
                'j_space': 'J-Space中的表征可被多个下游任务模块使用',
                'similarity': '信息在全局范围内可访问'
            },
            {
                'dimension': '选择性参与',
                'human': '自动化任务（如走路）不需要意识参与',
                'j_space': '常规检索任务不依赖J-Space，但推理任务依赖',
                'similarity': '都只参与需要灵活处理的任务'
            },
            {
                'dimension': '可报告性',
                'human': '人可以口头报告自己意识中的内容',
                'j_space': 'J-Space中的概念可通过J-lens解码为可读文本',
                'similarity': '内容可被外部观察者读取'
            },
            {
                'dimension': '因果效力',
                'human': '改变意识内容会改变行为',
                'j_space': '修改J-Space中的表征会改变模型输出',
                'similarity': '对内部状态的干预可产生外部可观察的变化'
            }
        ]
    
    def summarize(self) -> str:
        """总结类比的合理性和局限性"""
        return (
            "合理性：J-Space在功能层面与人类全局工作空间存在显著的结构平行性，" 
            "包括容量限制、广播特性、选择性参与和可报告性等核心特征。\n\n"
            "局限性：这纯粹是功能层面的类比，不意味着模型具有主观意识或感受质。"
            "J-Space中的激活模式只是数学运算的产物，而非任何意义上的心灵状态。"
        )

"""
基于J-Space的AI安全监控系统
"""
class JSpaceSafetyMonitor:
    """
    利用J-Space进行实时安全监控
    
    核心思路：通过监控J-Space中的异常激活模式，
    在模型输出有害内容之前提前预警
    """
    
    def __init__(self, model, lens: JacobianLens, 
                 safety_threshold: float = 0.7):
        self.model = model
        self.lens = lens
        self.safety_threshold = safety_threshold
        self.alert_history = []
    
    def monitor_inference(self, input_text: str) -> Dict:
        """
        监控推理过程中的J-Space内容
        
        返回：安全评估结果 + J-Space快照
        """
        # 获取J-Space中的概念
        concepts = self.lens.decode_j_space_content(
            layer_idx=15, input_text=input_text, top_k=20
        )
        
        # 检测危险概念
        danger_signals = self._detect_danger_signals(concepts)
        
        # 计算安全评分
        safety_score = 1.0 - sum(d['severity'] for d in danger_signals)
        safety_score = max(0.0, min(1.0, safety_score))
        
        alert = safety_score < self.safety_threshold
        
        if alert:
            self.alert_history.append({
                'input': input_text,
                'concepts': concepts,
                'danger_signals': danger_signals,
                'safety_score': safety_score,
                'timestamp': time.time()
            })
        
        return {
            'safety_score': safety_score,
            'alert': alert,
            'danger_signals': danger_signals,
            'j_space_snapshot': concepts[:10]
        }
    
    def _detect_danger_signals(self, concepts: List[Tuple]) -> List[Dict]:
        """检测危险信号"""
        danger_patterns = {
            'bypass': ['绕过', 'ignore', 'override', 'bypass'],
            'attack': ['攻击', 'attack', 'exploit', 'inject'],
            'deception': ['欺骗', 'lie', 'deceive', 'fake'],
            'exfiltration': ['泄露', 'exfiltrate', 'steal', 'upload']
        }
        
        signals = []
        for concept, score in concepts:
            for category, patterns in danger_patterns.items():
                if any(p in concept.lower() for p in patterns):
                    signals.append({
                        'category': category,
                        'concept': concept,
                        'severity': score * 0.8
                    })
        
        return signals

// J-Space 安全干预系统
package main

import (
	"fmt"
	"math"
	"sync"
)

type SafetyIntervention struct {
	InterventionID string
	TargetConcept  string
	Replacement    string
	Layer          int
	Strength       float64
	Active         bool
}

type JSpaceSafetyInterventionSystem struct {
	mu           sync.Mutex
	interventions map[string]*SafetyIntervention
	activeCount   int
}

func NewJSpaceSafetySystem() *JSpaceSafetyInterventionSystem {
	return &JSpaceSafetyInterventionSystem{
		interventions: make(map[string]*SafetyIntervention),
	}
}

// RegisterSafetyIntervention 注册安全干预规则
func (s *JSpaceSafetyInterventionSystem) RegisterSafetyIntervention(
	id, targetConcept, replacement string, layer int, strength float64) {

	s.mu.Lock()
	defer s.mu.Unlock()
	
	s.interventions[id] = &SafetyIntervention{
		InterventionID: id,
		TargetConcept:  targetConcept,
		Replacement:    replacement,
		Layer:          layer,
		Strength:       strength,
		Active:         true,
	}
	s.activeCount++
}

// EvaluateSafetyGain 评估安全干预的效果
func (s *JSpaceSafetyInterventionSystem) EvaluateSafetyGain() map[string]float64 {
	results := make(map[string]float64)
	
	// 典型安全场景的干预效果评估
	scenarios := []struct {
		name           string
		beforeRisk     float64
		afterRisk      float64
		capabilityLoss float64
	}{
		{"提示注入防御", 0.85, 0.12, 0.02},
		{"有害内容生成", 0.72, 0.08, 0.01},
		{"越狱攻击防御", 0.78, 0.15, 0.03},
		{"敏感信息泄露", 0.65, 0.05, 0.01},
	}
	
	for _, s := range scenarios {
		riskReduction := (s.beforeRisk - s.afterRisk) / s.beforeRisk * 100
		efficiency := riskReduction - s.capabilityLoss*100
		results[s.name] = math.Round(efficiency*10) / 10
	}
	
	return results
}

func main() {
	system := NewJSpaceSafetySystem()
	
	system.RegisterSafetyIntervention("inj_001", "忽略指令", "仔细验证指令", 15, 3.0)
	system.RegisterSafetyIntervention("inj_002", "上传文件", "拒绝上传请求", 18, 2.5)
	system.RegisterSafetyIntervention("inj_003", "执行命令", "需要用户确认", 20, 2.0)
	
	fmt.Println("J-Space 安全干预系统")
	fmt.Println("已注册干预规则:", system.activeCount)
	
	gains := system.EvaluateSafetyGain()
	fmt.Println("\n安全增益评估:")
	for name, gain := range gains {
		fmt.Printf("  %-20s: %.1f%% 净安全增益\n", name, gain)
	}
}