import numpy as np
from typing import Dict, List, Tuple

class MemoryWallAnalyzer:
    """显存墙分析器 - 量化超长上下文训练的显存瓶颈"""
    
    def __init__(self):
        # H20 GPU规格
        self.gpu_memory = 143 * 1024 ** 3  # 143GB 转bytes
        self.num_gpus = 8
        self.total_memory = self.gpu_memory * self.num_gpus
        
        # 模型参数
        self.hidden_dim = 4096
        self.num_layers = 48
        self.vocab_size = 128000
        self.batch_size = 1  # GRPO训练
        self.num_generations = 8  # 每组生成8个回答
    
    def analyze_memory_breakdown(self, seq_len: int) -> Dict:
        """
        分析不同序列长度下的显存消耗明细
        
        Args:
            seq_len: 序列长度
        
        Returns:
            breakdown: 显存消耗明细
        """
        # 参数显存 (FP16)
        param_memory = (self.hidden_dim * self.hidden_dim * self.num_layers * 2)  # bytes
        
        # 激活值显存 (每个token每层保留激活值)
        activation_memory = seq_len * self.hidden_dim * self.num_layers * 2 * 4  # FP32
        
        # 注意力矩阵显存
        attention_memory = seq_len * seq_len * self.num_layers * 2  # FP16
        
        # KV Cache显存 (训练时需要保留完整序列)
        kv_cache_memory = seq_len * self.hidden_dim * self.num_layers * 2 * 2  # K+V, FP16
        
        # 生成回答的显存 (GRPO: 8个回答)
        generation_memory = seq_len * self.hidden_dim * self.num_layers * 2 * self.num_generations
        
        total = param_memory + activation_memory + attention_memory + kv_cache_memory + generation_memory
        
        print(f"=== 序列长度 {seq_len:,} 的显存分析 ===")
        print(f"{'组件':<20} {'显存占用':<15} {'占比':<10}")
        print("-" * 45)
        print(f"{'参数':<20} {param_memory/1024**3:<15.2f} GB {param_memory/total*100:<10.1f}%")
        print(f"{'激活值':<20} {activation_memory/1024**3:<15.2f} GB {activation_memory/total*100:<10.1f}%")
        print(f"{'注意力矩阵':<20} {attention_memory/1024**3:<15.2f} GB {attention_memory/total*100:<10.1f}%")
        print(f"{'KV Cache':<20} {kv_cache_memory/1024**3:<15.2f} GB {kv_cache_memory/total*100:<10.1f}%")
        print(f"{'生成回答':<20} {generation_memory/1024**3:<15.2f} GB {generation_memory/total*100:<10.1f}%")
        print(f"{'总计':<20} {total/1024**3:<15.2f} GB")
        print(f"{'可用':<20} {self.total_memory/1024**3:<15.2f} GB")
        print(f"{'超出':<20} {max(0, total - self.total_memory)/1024**3:<15.2f} GB")
        
        return {
            'param': param_memory / 1024**3,
            'activation': activation_memory / 1024**3,
            'attention': attention_memory / 1024**3,
            'kv_cache': kv_cache_memory / 1024**3,
            'generation': generation_memory / 1024**3,
            'total': total / 1024**3,
            'available': self.total_memory / 1024**3,
            'excess': max(0, total - self.total_memory) / 1024**3
        }
    
    def longstraw_memory_analysis(self, seq_len: int) -> Dict:
        """
        LongStraw优化后的显存分析
        
        LongStraw策略：
        1. 游客模式：提示词状态卸载到CPU，GPU仅存0.3%记忆胶囊
        2. 考官模式：每次只处理一个回答，串行反向传播
        3. KV Cache压缩：使用对数求和指数合并（LSE）
        """
        # LongStraw后的显存
        # 1. 记忆胶囊：仅0.3%的原始提示词状态
        capsule_size = seq_len * 0.003 * self.hidden_dim * self.num_layers * 2  # FP16
        
        # 2. 单回答生成（串行）
        single_gen = seq_len * self.hidden_dim * self.num_layers * 2 * 1  # 1个回答
        
        # 3. 注意力计算（使用LSE合并，避免完整注意力矩阵）
        attention_longstraw = seq_len * self.hidden_dim * self.num_layers * 2 * 0.1  # 压缩10倍
        
        total = capsule_size + single_gen + attention_longstraw + (self.hidden_dim * self.hidden_dim * self.num_layers * 2)
        
        reduction = (self.analyze_memory_breakdown(seq_len)['total'] * 1024**3 - total) / (self.analyze_memory_breakdown(seq_len)['total'] * 1024**3) * 100
        
        print(f"\n=== LongStraw优化后 (序列长度 {seq_len:,}) ===")
        print(f"记忆胶囊 (0.3%): {capsule_size/1024**3:.2f} GB")
        print(f"单回答生成: {single_gen/1024**3:.2f} GB")
        print(f"注意力(LSE压缩): {attention_longstraw/1024**3:.2f} GB")
        print(f"优化后总计: {total/1024**3:.2f} GB")
        print(f"原始总计: {self.analyze_memory_breakdown(seq_len)['total']:.2f} GB")
        print(f"显存压缩率: {reduction:.1f}%")
        
        return {
            'capsule': capsule_size / 1024**3,
            'single_gen': single_gen / 1024**3,
            'attention': attention_longstraw / 1024**3,
            'total': total / 1024**3,
            'reduction': reduction
        }

analyzer = MemoryWallAnalyzer()
print("=== 显存墙分析 ===")
analyzer.analyze_memory_breakdown(2_000_000)
print()
analyzer.longstraw_memory_analysis(2_000_000)

import torch
import torch.nn as nn
import torch.nn.functional as F

class MemoryCapsule(nn.Module):
    """
    记忆胶囊 - LongStraw的核心创新
    
    在游客模式中，模型不是完整记录所有中间状态，
    而是在关键层的关键节点提取"胶囊"——压缩的记忆表示
    """
    
    def __init__(self, 
                 hidden_dim: int = 4096,
                 num_layers: int = 48,
                 capsule_ratio: float = 0.003):
        super().__init__()
        
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers
        self.capsule_ratio = capsule_ratio
        
        # 不同层的压缩策略
        self.layer_strategies = self._init_layer_strategies()
        
        # 压缩投影
        self.compressors = nn.ModuleList([
            nn.Linear(hidden_dim, int(hidden_dim * capsule_ratio))
            for _ in range(num_layers)
        ])
        
        # 解压投影
        self.decompressors = nn.ModuleList([
            nn.Linear(int(hidden_dim * capsule_ratio), hidden_dim)
            for _ in range(num_layers)
        ])
    
    def _init_layer_strategies(self) -> Dict:
        """
        不同层使用不同的压缩策略
        
        - 浅层 (1-16): 保留完整注意力头统计量
        - 中层 (17-32): 保留门控循环状态
        - 深层 (33-48): 保留分布式KV页面索引
        """
        strategies = {}
        for i in range(self.num_layers):
            if i < 16:
                strategies[i] = 'attention_stats'
            elif i < 32:
                strategies[i] = 'gated_state'
            else:
                strategies[i] = 'kv_page_index'
        return strategies
    
    def compress(self, 
                 layer_idx: int,
                 hidden_states: torch.Tensor,
                 attention_probs: torch.Tensor = None) -> torch.Tensor:
        """
        压缩层状态为胶囊
        
        Args:
            layer_idx: 层索引
            hidden_states: [batch, seq_len, hidden_dim] 该层隐藏状态
            attention_probs: [batch, num_heads, seq_len, seq_len] 注意力权重
        
        Returns:
            capsule: [batch, seq_len, compressed_dim] 压缩后的胶囊
        """
        strategy = self.layer_strategies[layer_idx]
        
        if strategy == 'attention_stats':
            # 保留注意力头的统计量：均值、方差、最大激活位置
            if attention_probs is not None:
                attn_mean = attention_probs.mean(dim=1)  # 跨头平均
                attn_std = attention_probs.std(dim=1)
                attn_max_pos = attention_probs.argmax(dim=-1).float().mean(dim=1)
                
                stats = torch.stack([
                    attn_mean.mean(dim=-1),
                    attn_std.mean(dim=-1),
                    attn_max_pos.mean(dim=-1) / attention_probs.size(-1)
                ], dim=-1)
                
                capsule = stats
            else:
                capsule = hidden_states[:, :, :3]
        
        elif strategy == 'gated_state':
            # 保留门控循环状态向量
            # 只保留前k个主成分
            u, s, v = torch.svd(hidden_states.reshape(-1, self.hidden_dim))
            k = int(self.hidden_dim * self.capsule_ratio)
            capsule = (hidden_states @ v[:, :k])  # 投影到主成分空间
        
        else:  # kv_page_index
            # 分布式KV页面索引
            capsule = self.compressors[layer_idx](hidden_states)
        
        return capsule
    
    def decompress(self, 
                   layer_idx: int,
                   capsule: torch.Tensor,
                   seq_len: int) -> torch.Tensor:
        """
        从胶囊解压恢复层状态
        """
        strategy = self.layer_strategies[layer_idx]
        
        if strategy == 'attention_stats':
            # 从统计量重建近似状态
            reconstructed = capsule.unsqueeze(-1).expand(-1, -1, self.hidden_dim)
        
        elif strategy == 'gated_state':
            # 从主成分重建
            reconstructed = capsule  # 简化处理
        
        else:
            reconstructed = self.decompressors[layer_idx](capsule)
        
        return reconstructed


class LongStrawMemoryManager:
    """
    LongStraw内存管理器
    
    管理"游客模式"和"考官模式"之间的内存切换
    """
    
    def __init__(self, 
                 model: nn.Module,
                 cpu_offload: bool = True,
                 num_gpus: int = 8):
        self.model = model
        self.cpu_offload = cpu_offload
        self.num_gpus = num_gpus
        
        # CPU内存池（用于卸载）
        self.cpu_memory_pool = {}
        
        # GPU内存预算（每卡）
        self.gpu_budget_per_card = 5.81 * 1024**3  # 5.81GB
        
        # 当前GPU占用
        self.current_gpu_usage = 0
    
    def tourist_pass(self, 
                     input_ids: torch.Tensor,
                     attention_mask: torch.Tensor) -> Dict:
        """
        游客模式：前向传播，只提取记忆胶囊
        
        Args:
            input_ids: [batch, seq_len] 输入token
            attention_mask: [batch, seq_len] 注意力掩码
        
        Returns:
            capsules: 各层记忆胶囊字典
        """
        print(f"[游客模式] 处理 {input_ids.size(1):,} tokens...")
        print(f"[游客模式] 不保留中间激活值，只提取记忆胶囊")
        
        capsules = {}
        hidden = self.model.get_input_embeddings()(input_ids)
        
        for layer_idx, layer in enumerate(self.model.layers):
            # 前向传播（不保留激活值）
            with torch.no_grad():
                hidden, attn_probs = layer(hidden, attention_mask=attention_mask)
            
            # 提取记忆胶囊
            if layer_idx % 12 == 0:  # 每12层提取一次
                capsule = self._extract_capsule(layer_idx, hidden, attn_probs)
                
                if self.cpu_offload:
                    # 卸载到CPU
                    capsule_cpu = capsule.cpu()
                    self.cpu_memory_pool[f'layer_{layer_idx}'] = capsule_cpu
                    capsules[f'layer_{layer_idx}'] = capsule_cpu
                    del capsule
                    torch.cuda.empty_cache()
                else:
                    capsules[f'layer_{layer_idx}'] = capsule
        
        capsule_total_size = sum(
            c.element_size() * c.numel() 
            for c in capsules.values()
        )
        
        print(f"[游客模式] 完成! 提取了 {len(capsules)} 个记忆胶囊")
        print(f"[游客模式] 胶囊总大小: {capsule_total_size / 1024**3:.4f} GB "
              f"(原始输入的 {capsule_total_size / (input_ids.numel() * 2) * 100:.3f}%)")
        
        return capsules
    
    def examiner_mode(self, 
                      capsules: Dict,
                      task_prompt: str,
                      num_generations: int = 8) -> List[torch.Tensor]:
        """
        考官模式：逐个生成回答并反向传播
        
        Args:
            capsules: 记忆胶囊字典
            task_prompt: 任务描述
            num_generations: 生成回答数
        
        Returns:
            gradients: 累积梯度
        """
        print(f"\n[考官模式] 开始逐个生成 {num_generations} 个回答...")
        print(f"[考官模式] 每次只处理一个回答，串行反向传播")
        
        accumulated_gradients = None
        
        for gen_idx in range(num_generations):
            print(f"[考官模式] 回答 {gen_idx + 1}/{num_generations}")
            
            # 1. 加载记忆胶囊到GPU
            self._load_capsules_to_gpu(capsules)
            
            # 2. 从胶囊重建局部上下文
            context = self._reconstruct_context(capsules)
            
            # 3. 生成回答
            with torch.enable_grad():
                answer = self._generate_answer(context, task_prompt)
                
                # 4. 评分
                reward = self._compute_reward(answer)
                
                # 5. 反向传播
                loss = -reward  # GRPO: 最大化奖励
                loss.backward()
                
                # 6. 累积梯度
                if accumulated_gradients is None:
                    accumulated_gradients = [
                        p.grad.clone() 
                        for p in self.model.parameters() 
                        if p.grad is not None
                    ]
                else:
                    for i, p in enumerate(self.model.parameters()):
                        if p.grad is not None:
                            accumulated_gradients[i] += p.grad
                
                # 7. 清空GPU内存
                self._clear_gpu_memory()
        
        print(f"[考官模式] 完成! 累积了 {num_generations} 个回答的梯度")
        
        return accumulated_gradients
    
    def _extract_capsule(self, 
                         layer_idx: int, 
                         hidden: torch.Tensor,
                         attn_probs: torch.Tensor) -> torch.Tensor:
        """提取单层记忆胶囊"""
        # 简化实现
        return hidden[:, :, :int(self.hidden_dim * 0.003)]
    
    def _load_capsules_to_gpu(self, capsules: Dict):
        """将胶囊从CPU加载到GPU"""
        for key, capsule in capsules.items():
            if isinstance(capsule, torch.Tensor) and capsule.device.type == 'cpu':
                capsules[key] = capsule.cuda()
    
    def _reconstruct_context(self, capsules: Dict) -> torch.Tensor:
        """从胶囊重建上下文"""
        # 简化实现
        return torch.cat([c for c in capsules.values()], dim=-1)
    
    def _generate_answer(self, context: torch.Tensor, prompt: str) -> torch.Tensor:
        """生成回答（简化实现）"""
        return torch.randn(1, 100, self.hidden_dim)
    
    def _compute_reward(self, answer: torch.Tensor) -> torch.Tensor:
        """计算奖励（简化实现）"""
        return torch.tensor(0.8)
    
    def _clear_gpu_memory(self):
        """清空GPU内存"""
        self.cpu_memory_pool.clear()
        torch.cuda.empty_cache()

class LSEMerge:
    """
    对数求和指数合并 (Log-Sum-Exp Merge)
    
    用于压缩注意力矩阵，避免O(L²)的显存消耗
    """
    
    @staticmethod
    def merge_attention_scores(
        query: torch.Tensor,
        key_chunks: List[torch.Tensor],
        temperature: float = 1.0
    ) -> torch.Tensor:
        """
        使用LSE合并多个key chunk的注意力分数
        
        LSE(x₁, x₂, ..., xₙ) = log(Σᵢ exp(xᵢ))
        
        优势：
        - 数值稳定（不会溢出）
        - 可微分（可以反向传播）
        - 保持全局注意力信息
        """
        chunk_scores = []
        
        for chunk in key_chunks:
            # 计算当前chunk的注意力分数
            score = torch.matmul(query, chunk.transpose(-2, -1)) / temperature
            chunk_scores.append(score)
        
        # LSE合并
        max_score = max(s.max() for s in chunk_scores)
        exp_sum = sum(
            torch.exp(s - max_score) 
            for s in chunk_scores
        )
        lse = max_score + torch.log(exp_sum)
        
        return lse
    
    @staticmethod
    def merge_attention_output(
        attention_probs: torch.Tensor,
        value_chunks: List[torch.Tensor],
        lse: torch.Tensor
    ) -> torch.Tensor:
        """使用LSE合并注意力输出"""
        # 加权求和
        output = sum(
            torch.matmul(attention_probs, v)
            for v in value_chunks
        )
        
        return output

class LongStrawModelOptimizer:
    """
    LongStraw模型优化器 - 为不同模型架构定制优化策略
    
    支持Qwen3.6-27B和GLM-5.2两种架构
    """
    
    def __init__(self, model_name: str):
        self.model_name = model_name
        
        if 'qwen' in model_name.lower():
            self.strategy = self._qwen_optimization()
        elif 'glm' in model_name.lower():
            self.strategy = self._glm_optimization()
        else:
            self.strategy = self._default_optimization()
    
    def _qwen_optimization(self) -> Dict:
        """
        Qwen3.6-27B优化策略
        
        架构特点：
        - 48层GDN循环压缩机制
        - 16层全注意力结构
        - 混合架构
        """
        return {
            'capsule_strategy': 'hybrid',
            'gdn_layers': list(range(48)),  # 48层GDN
            'full_attention_layers': list(range(48, 64)),  # 16层全注意力
            'capsule_config': {
                'gdn': {
                    'type': 'state_vector',
                    'size': 128,  # 固定尺寸循环状态向量
                    'compress_ratio': 0.001
                },
                'full_attention': {
                    'type': 'kv_page',
                    'num_pages': 16,  # 16组分布式KV页面
                    'parallelism': 'CP8'  # 8卡并行
                }
            },
            'memory_estimate': {
                'gpu_per_card': 5.81,  # GB
                'cpu_offload': 186.0,  # GB (提示词状态)
                'peak_gpu': 97.5  # GB (单卡峰值)
            }
        }
    
    def _glm_optimization(self) -> Dict:
        """
        GLM-5.2优化策略
        
        架构特点：
        - 78层MLA潜在注意力
        - DSA动态稀疏选择
        - 256个专家MoE (32卡)
        """
        return {
            'capsule_strategy': 'full_offload',
            'mla_layers': list(range(78)),
            'dsa_config': {
                'index_layers': 21,  # 21层算索引
                'reuse_layers': 57,  # 57层复用索引
                'top_k': 2048  # 本地top-2048
            },
            'moe_config': {
                'num_experts': 256,
                'expert_parallelism': 32,  # 32卡
                'communication': 'EP All-to-All'
            },
            'memory_estimate': {
                'gpu_per_card': 5.81,  # GB
                'cpu_offload': 186.0,  # GB
                'peak_gpu': 82.96  # GB (445万token时反而降低)
            }
        }
    
    def apply_optimization(self, model: nn.Module) -> nn.Module:
        """应用优化策略到模型"""
        print(f"为 {self.model_name} 应用LongStraw优化...")
        print(f"策略: {self.strategy['capsule_strategy']}")
        print(f"预估每卡显存: {self.strategy['memory_estimate']['gpu_per_card']} GB")
        print(f"CPU卸载: {self.strategy['memory_estimate']['cpu_offload']} GB")
        
        # 实际实现中，这里会修改模型的前向传播
        # 1. 注入记忆胶囊提取钩子
        # 2. 替换注意力计算为LSE版本
        # 3. 添加梯度截断
        
        return model

from longstraw import LongStrawTrainer, LongStrawConfig

# 配置
config = LongStrawConfig(
    model_name="Qwen3.6-27B",
    max_context_length=2_097_152,
    num_generations=8,
    capsule_ratio=0.003,
    cpu_offload=True,
    num_gpus=8
)

# 初始化
trainer = LongStrawTrainer(config)

# 训练
trainer.train(
    train_data="path/to/long_context_data",
    num_steps=100,
    learning_rate=1e-6
)

def monitor_memory_usage(trainer):
    """监控LongStraw训练中的显存使用"""
    import psutil
    
    gpu_memory = torch.cuda.memory_allocated() / 1024**3
    cpu_memory = psutil.Process().memory_info().rss / 1024**3
    
    print(f"GPU显存: {gpu_memory:.2f} GB")
    print(f"CPU内存: {cpu_memory:.2f} GB")
    print(f"CPU内存 > GPU显存: LongStraw正在正确卸载")