import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from typing import Tuple, Optional

class LatentSparseAttention(nn.Module):
    """
    LSA: Latent Sparse Attention Module
    将标准O(n²)注意力降至O(n·k)的线性复杂度
    
    Args:
        d_model: 模型隐藏维度
        n_heads: 注意力头数
        local_window_size: 局部窗口大小
        global_sparse_ratio: 全局稀疏采样比例 (默认1%)
        compression_dim: 压缩查询维度
    """
    def __init__(
        self,
        d_model: int = 7168,
        n_heads: int = 64,
        local_window_size: int = 4096,
        global_sparse_ratio: float = 0.01,
        compression_dim: int = 512,
    ):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads
        self.local_window_size = local_window_size
        self.global_sparse_ratio = global_sparse_ratio
        
        # 标准QKV投影
        self.wq = nn.Linear(d_model, d_model, bias=False)
        self.wk = nn.Linear(d_model, d_model, bias=False)
        self.wv = nn.Linear(d_model, d_model, bias=False)
        self.wo = nn.Linear(d_model, d_model, bias=False)
        
        # 压缩查询投影 (Compression Query Path)
        self.compression_query = nn.Linear(d_model, compression_dim, bias=False)
        self.compression_key = nn.Linear(d_model, compression_dim, bias=False)
        
        # 顶层门控：决定每个token走三条路径的权重
        self.router = nn.Linear(d_model, 3)  # global, local, compression
        
    def _global_sparse_attn(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """全局稀疏注意力：仅采样1%的关键Key参与计算"""
        batch_size, seq_len, _ = q.shape
        n_samples = max(1, int(seq_len * self.global_sparse_ratio))
        
        # 基于Attention Dispersion的Key采样策略
        with torch.no_grad():
            # 计算每个位置的注意力分散度
            q_norm = q.norm(dim=-1, p=2)  # [B, L]
            k_norm = k.norm(dim=-1, p=2)  # [B, L]
            
            # key重要性分数 = query-key归一化内积的距离
            importance = torch.matmul(
                q_norm.unsqueeze(-1), 
                k_norm.unsqueeze(1)
            ).squeeze(-1)  # [B, L, L] → [B, L]
            
            # Top-K采样：选取最重要的n_samples个key
            topk_indices = torch.topk(importance, n_samples, dim=-1).indices
            topk_indices, _ = torch.sort(topk_indices, dim=-1)
            
        # Gather selected keys and values
        selected_k = k.gather(
            1, topk_indices.unsqueeze(-1).expand(-1, -1, k.size(-1))
        )
        selected_v = v.gather(
            1, topk_indices.unsqueeze(-1).expand(-1, -1, v.size(-1))
        )
        
        # 标准Attention计算 (仅在采样后的子集上)
        scale = math.sqrt(self.head_dim)
        q_heads = q.view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
        k_heads = selected_k.view(batch_size, n_samples, self.n_heads, self.head_dim).transpose(1, 2)
        v_heads = selected_v.view(batch_size, n_samples, self.n_heads, self.head_dim).transpose(1, 2)
        
        attn_scores = torch.matmul(q_heads, k_heads.transpose(-2, -1)) / scale
        attn_weights = F.softmax(attn_scores, dim=-1)
        context = torch.matmul(attn_weights, v_heads)
        context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, -1)
        
        return context
    
    def _local_window_attn(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
    ) -> torch.Tensor:
        """局部窗口注意力：每个token只看前后window_size/2个位置"""
        batch_size, seq_len, _ = q.shape
        half_window = self.local_window_size // 2
        
        # 滑动窗口padding
        k_padded = F.pad(k, (0, 0, half_window, half_window), mode='replicate')
        v_padded = F.pad(v, (0, 0, half_window, half_window), mode='replicate')
        
        # 使用unfold提取窗口
        k_windows = k_padded.unfold(1, self.local_window_size, 1)  # [B, L, D, W]
        v_windows = v_padded.unfold(1, self.local_window_size, 1)
        
        # 局部注意力计算
        scale = math.sqrt(self.head_dim)
        q_heads = q.view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
        
        # Reshape for window attention
        k_heads = k_windows.permute(0, 2, 3, 1).contiguous()
        k_heads = k_heads.view(batch_size, self.n_heads, self.head_dim, seq_len, self.local_window_size)
        
        attn_local = torch.einsum('bhnl,bhdlw->bhnlw', q_heads, k_heads) / scale
        attn_local_weights = F.softmax(attn_local, dim=-1)
        
        v_heads = v_windows.permute(0, 2, 3, 1).contiguous()
        v_heads = v_heads.view(batch_size, self.n_heads, self.head_dim, seq_len, self.local_window_size)
        
        context = torch.einsum('bhnlw,bhdlw->bhnl', attn_local_weights, v_heads)
        context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, -1)
        
        return context
    
    def _compressed_query_attn(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
    ) -> torch.Tensor:
        """压缩查询路径：在缩小的隐空间中进行全局注意力"""
        q_compressed = self.compression_query(q)
        k_compressed = self.compression_key(k)
        
        scale = math.sqrt(q_compressed.size(-1))
        attn_compressed = torch.matmul(q_compressed, k_compressed.transpose(-2, -1)) / scale
        attn_compressed_weights = F.softmax(attn_compressed, dim=-1)
        
        context = torch.matmul(attn_compressed_weights, v)
        return context
    
    def forward(
        self,
        x: torch.Tensor,
        mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        batch_size, seq_len, _ = x.shape
        
        # QKV投影
        q, k, v = self.wq(x), self.wk(x), self.wv(x)
        
        # 三条路径
        global_context = self._global_sparse_attn(q, k, v, mask)
        local_context = self._local_window_attn(q, k, v)
        compressed_context = self._compressed_query_attn(q, k, v)
        
        # 门控融合：动态学习每条路径的权重
        gate_logits = self.router(x)  # [B, L, 3]
        gate_weights = F.softmax(gate_logits, dim=-1).unsqueeze(-1)  # [B, L, 3, 1]
        
        # 加权融合
        stacked_context = torch.stack(
            [global_context, local_context, compressed_context], dim=2
        )  # [B, L, 3, D]
        
        output = (gate_weights * stacked_context).sum(dim=2)
        output = self.wo(output)
        
        return output

package moe

import (
	"context"
	"math"
	"sync"

	"gorgonia.org/tensor"
)

// ZeroComputeRouter 零计算专家路由器
// 根据Token复杂度决定是否跳过专家计算
type ZeroComputeRouter struct {
	// 复杂度评分网络参数
	scoreWeight tensor.Tensor // [hiddenDim, 1]
	scoreBias   float64
	threshold   float64

	// 复杂度评估指标权重
	entropyWeight      float64 // Token熵权重
	dispersionWeight   float64 // 注意力分散度权重
	uncertaintyWeight  float64 // 跨层不确定性权重

	// 统计信息
	mu             sync.RWMutex
	totalTokens    int64
	zeroComputeCnt int64
	avgComplexity  float64
}

// NewZeroComputeRouter 创建零计算路由器
func NewZeroComputeRouter(hiddenDim int, threshold float64) *ZeroComputeRouter {
	w := tensor.New(
		tensor.WithShape(hiddenDim, 1),
		tensor.WithRandom(tensor.NormalDistribution(0, 0.02)),
	)

	return &ZeroComputeRouter{
		scoreWeight:      *w,
		scoreBias:        -0.5, // 初始偏置偏向"不跳过"
		threshold:        threshold,
		entropyWeight:    0.35,
		dispersionWeight: 0.40,
		uncertaintyWeight: 0.25,
	}
}

// TokenComplexity 单个Token的复杂度指标
type TokenComplexity struct {
	Entropy       float64 // Token级别信息熵
	Dispersion    float64 // 注意力分散度
	Uncertainty   float64 // 跨层不确定性
	CombinedScore float64 // 综合分数
}

// ComputeComplexity 计算单个Token的复杂度
func (r *ZeroComputeRouter) ComputeComplexity(
	hiddenState []float64,
	attnDistribution []float64,
	layerOutputs [][]float64,
) TokenComplexity {
	// 1. Token Entropy: 衡量hidden state的信息量
	entropy := computeTokenEntropy(hiddenState)

	// 2. Attention Dispersion: 注意力分布越分散越复杂
	dispersion := computeAttentionDispersion(attnDistribution)

	// 3. Layer-wise Uncertainty: 跨层输出方差
	uncertainty := computeCrossLayerUncertainty(layerOutputs)

	// 4. 加权综合分数
	combinedScore := r.entropyWeight*entropy +
		r.dispersionWeight*dispersion +
		r.uncertaintyWeight*uncertainty

	return TokenComplexity{
		Entropy:       entropy,
		Dispersion:    dispersion,
		Uncertainty:   uncertainty,
		CombinedScore: combinedScore,
	}
}

// ShouldSkip 判断当前Token是否应跳过专家计算
func (r *ZeroComputeRouter) ShouldSkip(complexity TokenComplexity) bool {
	skip := complexity.CombinedScore < r.threshold

	r.mu.Lock()
	r.totalTokens++
	if skip {
		r.zeroComputeCnt++
	}
	r.avgComplexity = r.avgComplexity*0.99 + complexity.CombinedScore*0.01
	r.mu.Unlock()

	return skip
}

// ZeroComputeOutput 零计算输出生成器
// 对于跳过计算的Token，通过轻量线性变换生成近似输出
func (r *ZeroComputeRouter) ZeroComputeOutput(
	hiddenState []float64,
	prevLayerOutput []float64,
) []float64 {
	output := make([]float64, len(hiddenState))

	// 轻量级映射：仅用残差连接 + 简单仿射变换
	// 避免全量专家计算的昂贵开销
	for i := range output {
		output[i] = prevLayerOutput[i]*0.8 + hiddenState[i]*0.2
	}

	return output
}

// ComputeStats 返回零计算路由器的运行统计
func (r *ZeroComputeRouter) ComputeStats() (skipRate float64, avgComplexity float64) {
	r.mu.RLock()
	defer r.mu.RUnlock()

	if r.totalTokens == 0 {
		return 0, 0
	}
	skipRate = float64(r.zeroComputeCnt) / float64(r.totalTokens)
	return skipRate, r.avgComplexity
}

// computeTokenEntropy 计算Token的近似信息熵
func computeTokenEntropy(hiddenState []float64) float64 {
	var sum float64
	for _, v := range hiddenState {
		p := 1.0 / (1.0 + math.Exp(-v)) // sigmoid
		if p > 1e-10 {
			sum -= p * math.Log(p)
		}
		if 1-p > 1e-10 {
			sum -= (1 - p) * math.Log(1-p)
		}
	}
	return sum / float64(len(hiddenState))
}

// computeAttentionDispersion 计算注意力分布的分散度
func computeAttentionDispersion(distribution []float64) float64 {
	var sum float64
	for _, p := range distribution {
		if p > 1e-10 {
			sum -= p * math.Log(p)
		}
	}
	// 归一化到[0,1]
	n := float64(len(distribution))
	return sum / math.Log(n)
}

// computeCrossLayerUncertainty 计算跨层输出不确定性
func computeCrossLayerUncertainty(layerOutputs [][]float64) float64 {
	if len(layerOutputs) < 2 {
		return 0
	}

	nLayers := len(layerOutputs)
	dim := len(layerOutputs[0])

	// 计算各层的均值
	mean := make([]float64, dim)
	for _, out := range layerOutputs {
		for j, v := range out {
			mean[j] += v / float64(nLayers)
		}
	}

	// 计算方差
	var variance float64
	for _, out := range layerOutputs {
		for j, v := range out {
			diff := v - mean[j]
			variance += diff * diff / float64(nLayers*dim)
		}
	}

	return math.Sqrt(variance)
}

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from typing import Dict, List, Tuple, Optional

class MOPDGate(nn.Module):
    """
    MOPD: Multi-expert Orchestration via Parallel Dispatch
    
    核心创新：
    1. 三层门控：Type Gate → Group Gate → Expert Gate
    2. 任务感知嵌入：根据输入特征动态调整分组权重
    3. 分组负载均衡：确保每组专家都得到充分训练
    """
    
    def __init__(
        self,
        d_model: int,
        num_experts: int = 256,
        num_groups: int = 3,
        top_k: int = 8,
        group_config: Optional[Dict[str, int]] = None,
    ):
        super().__init__()
        
        self.num_experts = num_experts
        self.num_groups = num_groups
        self.top_k = top_k
        
        # 分组配置: Agent=128, Reasoning=64, Interaction=64
        if group_config is None:
            group_config = {"agent": 128, "reasoning": 64, "interaction": 64}
        self.group_config = group_config
        
        # 各组的专家索引偏移
        self.group_offsets = {}
        offset = 0
        for name, count in group_config.items():
            self.group_offsets[name] = (offset, offset + count)
            offset += count
        
        # 第一层：类型门控 (Type Gate)
        # 决定当前Token应该主要使用哪个专家组的能力
        self.type_gate = nn.Sequential(
            nn.Linear(d_model, d_model // 2),
            nn.GELU(),
            nn.Linear(d_model // 2, num_groups),
        )
        
        # 第二层：分组门控 (Group Gate) — 每组独立的专家选择器
        self.group_gates = nn.ModuleDict()
        for name, count in group_config.items():
            self.group_gates[name] = nn.Sequential(
                nn.Linear(d_model, d_model // 2),
                nn.GELU(),
                nn.Linear(d_model // 2, count),
            )
        
        # 任务类型嵌入（可学习）
        self.task_embedding = nn.Embedding(10, d_model // 4)  # 10种任务类型
        
        # 共享门控编码器
        self.shared_encoder = nn.Sequential(
            nn.Linear(d_model + d_model // 4, d_model),
            nn.GELU(),
            nn.Dropout(0.1),
        )
        
        # 负载均衡辅助损失权重
        self.load_balancing_coef = 0.01
        
    def _infer_task_type(self, hidden_state: torch.Tensor) -> torch.Tensor:
        """
        从hidden state推断任务类型嵌入
        通过分析激活模式、注意力分布等隐式特征推断任务类型
        """
        # 简单实现：用均值池化+线性投影推断任务类型
        pooled = hidden_state.mean(dim=1)  # [B, D]
        
        # 推断任务类型ID (0-9)
        task_logits = self.shared_encoder[:2](pooled)  # 复用部分编码器
        task_ids = task_logits.argmax(dim=-1)
        
        return self.task_embedding(task_ids)
    
    def forward(
        self,
        hidden_state: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor, Dict[str, float]]:
        """
        Args:
            hidden_state: [batch_size, seq_len, d_model] 或 [batch_size, d_model]
        Returns:
            expert_weights: [batch_size, top_k] 每个专家权重
            expert_indices: [batch_size, top_k] 选中专家的索引
            aux_losses: 辅助损失字典
        """
        if hidden_state.dim() == 3:
            B, L, D = hidden_state.shape
            hidden_state = hidden_state.reshape(-1, D)  # [B*L, D]
        
        batch_size = hidden_state.size(0)
        
        # 1. 推断任务类型并拼接任务嵌入
        task_embed = self._infer_task_type(hidden_state)
        combined_input = torch.cat([hidden_state, task_embed], dim=-1)
        encoded = self.shared_encoder(combined_input)
        
        # 2. Type Gate: 决定各组权重
        type_logits = self.type_gate(encoded)  # [B, num_groups]
        type_weights = F.softmax(type_logits, dim=-1)  # [B, num_groups]
        
        # 3. Group Gate: 每组独立选出Top-K专家
        all_expert_logits = []
        all_group_weights = []
        
        aux_losses = {}
        total_load_balance_loss = 0.0
        
        for i, (name, count) in enumerate(self.group_config.items()):
            # 该组的门控logits
            group_logits = self.group_gates[name](encoded)  # [B, count]
            
            # 负载均衡：计算该组的专家使用率
            group_probs = F.softmax(group_logits, dim=-1)
            
            # 在该组内选Top-K (按比例分配)
            group_top_k = max(1, self.top_k * count // self.num_experts)
            top_k_logits, top_k_indices = torch.topk(
                group_logits, group_top_k, dim=-1
            )
            
            # 相对权重 (组内归一化)
            top_k_weights = F.softmax(top_k_logits, dim=-1)
            
            # 调整为绝对权重 (乘以type_gate的组权重)
            group_weight = type_weights[:, i:i+1]  # [B, 1]
            adjusted_weights = top_k_weights * group_weight
            
            # 映射到全局专家索引
            offset_start, offset_end = self.group_offsets[name]
            global_indices = top_k_indices + offset_start
            
            all_expert_logits.append((adjusted_weights, global_indices))
            
            # 负载均衡损失: 鼓励均匀使用组内专家
            importance_per_expert = group_probs.mean(dim=0)
            uniform_dist = torch.ones(count, device=hidden_state.device) / count
            load_balance_loss = F.kl_div(
                importance_per_expert.log(), uniform_dist, reduction='sum'
            )
            total_load_balance_loss += load_balance_loss
            
            aux_losses[f"group_{name}_load_balance"] = load_balance_loss.item()
        
        # 4. 合并所有组的专家选择结果
        all_weights = torch.cat([w for w, _ in all_expert_logits], dim=-1)  # [B, total_top_k]
        all_indices = torch.cat([idx for _, idx in all_expert_logits], dim=-1)  # [B, total_top_k]
        
        # 5. 全局Top-K修剪 (保证不超过top_k个专家被激活)
        global_top_k, top_k_global_indices = torch.topk(all_weights, self.top_k, dim=-1)
        gathered_indices = all_indices.gather(1, top_k_global_indices)
        
        # 最终权重归一化
        final_weights = F.softmax(global_top_k, dim=-1)
        
        aux_losses["total_load_balance_loss"] = total_load_balance_loss.item()
        aux_losses["avg_type_weights"] = {
            name: type_weights[:, i].mean().item()
            for i, name in enumerate(self.group_config.keys())
        }
        
        return final_weights, gathered_indices, aux_losses


class MOPDExpertLayer(nn.Module):
    """
    完整的MOPD专家层
    包含门控 + FFN专家 + Zero-Compute集成
    """
    
    def __init__(
        self,
        d_model: int,
        d_ff: int,
        num_experts: int = 256,
        top_k: int = 8,
        use_zero_compute: bool = True,
    ):
        super().__init__()
        
        self.d_model = d_model
        self.d_ff = d_ff
        self.use_zero_compute = use_zero_compute
        
        # MOPD门控
        self.gate = MOPDGate(
            d_model=d_model,
            num_experts=num_experts,
            top_k=top_k,
        )
        
        # FFN专家 (共享权重池)
        self.experts = nn.ModuleList([
            nn.Sequential(
                nn.Linear(d_model, d_ff),
                nn.GELU(),
                nn.Linear(d_ff, d_model),
            )
            for _ in range(num_experts)
        ])
        
        # Zero-Compute Router (可选)
        if use_zero_compute:
            self.zero_router = nn.Linear(d_model, 1)
    
    def forward(
        self,
        x: torch.Tensor,
    ) -> torch.Tensor:
        """
        MOPD专家层前向传播
        """
        if self.use_zero_compute:
            # Zero-Compute决策
            zero_logits = self.zero_router(x)  # [B, L, 1]
            zero_mask = torch.sigmoid(zero_logits) < 0.3  # 低于阈值则跳过
            
            # 如果所有token都跳过，直接返回残差
            if zero_mask.all():
                return x
        
        # MOPD门控选择专家
        weights, indices, _ = self.gate(x)
        
        # 专家计算
        batch_size, seq_len, d_model = x.shape
        x_flat = x.reshape(-1, d_model)  # [B*L, D]
        
        output = torch.zeros_like(x_flat)
        
        # 专家并行计算
        for k in range(indices.size(1)):
            expert_idx = indices[:, k:k+1].squeeze(-1)  # [B*L]
            weight = weights[:, k:k+1]  # [B*L, 1]
            
            # 按专家索引分组
            expert_outputs = []
            for e_idx in range(len(self.experts)):
                mask = (expert_idx == e_idx)
                if mask.any():
                    expert_out = self.experts[e_idx](x_flat[mask])
                    expert_outputs.append((mask, expert_out))
            
            # 加权聚合
            for mask, expert_out in expert_outputs:
                output[mask] += weight[mask] * expert_out
        
        output = output.reshape(batch_size, seq_len, d_model)
        
        # Zero-Compute位置的输出替换为残差
        if self.use_zero_compute:
            zero_mask_expanded = zero_mask.expand(-1, -1, d_model)
            output = torch.where(zero_mask_expanded, x, output)
        
        return output

package train

import (
	"context"
	"fmt"
	"log"
	"math"
	"sync"
	"time"

	"golang.org/x/sync/errgroup"
)

// 监控指标
type NodeHealth struct {
	NodeID        string    `json:"node_id"`
	HeartbeatTime time.Time `json:"heartbeat_time"`
	TemperatureC  float64   `json:"temperature_c"`
	MemUtilPct    float64   `json:"mem_util_pct"`
	CommLatencyMs float64   `json:"comm_latency_ms"`
	ErrorCount    int       `json:"error_count"`
	IsHealthy     bool      `json:"is_healthy"`
}

// ElasticTrainCluster 弹性训练集群管理器
// 支持自动故障检测、任务迁移和弹性扩缩
type ElasticTrainCluster struct {
	mu sync.RWMutex

	nodes          map[string]*NodeHealth
	totalCards     int
	healthyCards   int
	minHealthyPct  float64 // 最低健康比例 (默认85%)

	// 检查点配置
	checkpointInterval time.Duration
	checkpointPath     string
	lastCheckpoint     string

	// 故障恢复
	failureThreshold    int           // 连续失败次数阈值
	recoveryTimeout     time.Duration // 节点恢复超时
	failedNodes         map[string]time.Time

	// 训练状态
	trainingCtx    context.Context
	trainingCancel context.CancelFunc
	globalStep     int64
	tokensProcessed int64

	// 弹性调度
	taskQueue    chan TrainingTask
	activeTasks  map[string]*TrainingTask
}

// TrainingTask 训练子任务
type TrainingTask struct {
	ID            string
	Type          string // "forward", "backward", "allreduce", "checkpoint"
	DataShardID   int
	ModelShardID  int
	ExpertShardID int
	Priority      int
	CreatedAt     time.Time
	AssignedNode  string
}

// NewElasticTrainCluster 创建弹性训练集群
func NewElasticTrainCluster(totalCards int) *ElasticTrainCluster {
	ctx, cancel := context.WithCancel(context.Background())

	return &ElasticTrainCluster{
		nodes:             make(map[string]*NodeHealth),
		totalCards:        totalCards,
		healthyCards:      totalCards,
		minHealthyPct:     0.85,
		checkpointInterval: 10 * time.Minute,
		checkpointPath:     "/checkpoints/longcat2/",
		failureThreshold:   3,
		recoveryTimeout:    5 * time.Minute,
		failedNodes:        make(map[string]time.Time),
		trainingCtx:        ctx,
		trainingCancel:     cancel,
		taskQueue:          make(chan TrainingTask, 10000),
		activeTasks:        make(map[string]*TrainingTask),
	}
}

// StartHealthMonitor 启动健康监控循环
func (c *ElasticTrainCluster) StartHealthMonitor(ctx context.Context) {
	ticker := time.NewTicker(30 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			c.checkAndRecover()
		}
	}
}

// checkAndRecover 检查集群健康状态并触发恢复
func (c *ElasticTrainCluster) checkAndRecover() {
	c.mu.Lock()
	defer c.mu.Unlock()

	healthyCount := 0
	failureCount := 0

	for id, health := range c.nodes {
		if !health.IsHealthy {
			failureCount++
			// 记录失败时间
			if _, exists := c.failedNodes[id]; !exists {
				c.failedNodes[id] = time.Now()
			}
		} else {
			healthyCount++
			// 清除恢复记录
			delete(c.failedNodes, id)
		}
	}

	c.healthyCards = healthyCount
	healthPct := float64(healthyCount) / float64(c.totalCards)

	log.Printf("[Cluster] Health: %d/%d cards (%.1f%%), failures: %d",
		healthyCount, c.totalCards, healthPct*100, failureCount)

	// 健康比例低于阈值，触发弹性缩容
	if healthPct < c.minHealthyPct {
		c.triggerElasticScaleDown()
	}

	// 检查是否有节点可以恢复
	c.tryRecoverNodes()
}

// triggerElasticScaleDown 弹性缩容
// 停止故障节点的任务，迁移到健康节点
func (c *ElasticTrainCluster) triggerElasticScaleDown() {
	log.Printf("[Elastic] Health below threshold (%.1f%% < %.1f%%), triggering scale-down",
		float64(c.healthyCards)/float64(c.totalCards)*100, c.minHealthyPct*100)

	// 1. 冻结故障节点上的任务
	for nodeID := range c.failedNodes {
		c.migrateTasksFromNode(nodeID)
	}

	// 2. 重新计算模型并行策略
	c.rebalanceParallelism()

	// 3. 触发紧急检查点
	c.triggerEmergencyCheckpoint()
}

// migrateTasksFromNode 迁移故障节点上的任务
func (c *ElasticTrainCluster) migrateTasksFromNode(nodeID string) {
	migratedCount := 0
	for taskID, task := range c.activeTasks {
		if task.AssignedNode == nodeID {
			// 重新分配任务到健康节点
			newNode := c.selectHealthyNode()
			if newNode != "" {
				task.AssignedNode = newNode
				migratedCount++
				log.Printf("[Migration] Task %s: %s → %s", taskID, nodeID, newNode)
			}
		}
	}
	log.Printf("[Migration] Migrated %d tasks from node %s", migratedCount, nodeID)
}

// selectHealthyNode 选择一个健康节点
func (c *ElasticTrainCluster) selectHealthyNode() string {
	// 最简单的策略：选择负载最低的健康节点
	var bestNode string
	minLoad := math.MaxFloat64

	for id, health := range c.nodes {
		if health.IsHealthy {
			load := health.MemUtilPct + health.TemperatureC/100.0
			if load < minLoad {
				minLoad = load
				bestNode = id
			}
		}
	}
	return bestNode
}

// tryRecoverNodes 尝试恢复故障节点
func (c *ElasticTrainCluster) tryRecoverNodes() {
	now := time.Now()
	for nodeID, failTime := range c.failedNodes {
		if now.Sub(failTime) > c.recoveryTimeout {
			// 超时，尝试重新加入集群
			log.Printf("[Recovery] Attempting to recover node %s (failed at %v)", nodeID, failTime)

			// 发起到节点的健康检查
			if c.pingNode(nodeID) {
				c.nodes[nodeID].IsHealthy = true
				delete(c.failedNodes, nodeID)
				log.Printf("[Recovery] Node %s recovered successfully", nodeID)
			}
		}
	}
}

// pingNode 检查节点是否可达
func (c *ElasticTrainCluster) pingNode(nodeID string) bool {
	// 模拟ICMP ping + gRPC健康检查
	// 实际实现中会发送gRPC HealthCheck请求
	return true // 简化实现
}

// triggerEmergencyCheckpoint 触发紧急检查点
func (c *ElasticTrainCluster) triggerEmergencyCheckpoint() {
	step := c.globalStep
	checkpointName := fmt.Sprintf("%scheckpoint_emergency_step_%d.tar", c.checkpointPath, step)

	log.Printf("[Checkpoint] Emergency checkpoint at step %d → %s", step, checkpointName)

	// 保存模型权重、优化器状态、数据加载偏移
	c.lastCheckpoint = checkpointName
}

// rebalanceParallelism 重新平衡并行策略
func (c *ElasticTrainCluster) rebalanceParallelism() {
	// 根据当前健康卡数重新计算:
	// 1. DP (Data Parallel) 度数
	// 2. TP (Tensor Parallel) 度数
	// 3. PP (Pipeline Parallel) 度数
	// 4. EP (Expert Parallel) 度数

	availableCards := c.healthyCards
	ppSize := 16               // 固定
	tpSize := 8                // 固定
	epSize := availableCards / (ppSize * tpSize)

	log.Printf("[Rebalance] New strategy: PP=%d, TP=%d, EP=%d, DP=%d (cards=%d)",
		ppSize, tpSize, epSize, availableCards/(ppSize*tpSize), availableCards)
}

"""
LongCat-2.0 Agentic Coding 实战：仓库级代码重构
使用LongCat-2.0 API完成旧版SDK到新版API的自动迁移
"""
import json
import os
from pathlib import Path
from typing import List, Dict, Optional
import requests
import git


class LongCatCodeRefactor:
    """
    基于LongCat-2.0的智能代码重构引擎
    
    工作流程：
    1. 读取旧版代码库结构
    2. 分析新版SDK文档
    3. 规划重构方案
    4. 逐模块执行迁移
    5. 生成测试用例并验证
    """
    
    def __init__(
        self,
        api_key: str,
        repo_path: str,
        model: str = "longcat-2.0",
        base_url: str = "https://api.longcat.ai/v1",
    ):
        self.api_key = api_key
        self.base_url = base_url
        self.model = model
        self.repo_path = Path(repo_path)
        self.repo = git.Repo(repo_path)
        
        # 代码分析结果
        self.code_graph: Dict = {}
        self.dependency_map: Dict = {}
        self.refactor_plan: List[Dict] = []
    
    def _call_longcat(self, messages: List[Dict], max_tokens: int = 32768) -> str:
        """调用LongCat-2.0 API"""
        response = requests.post(
            f"{self.base_url}/chat/completions",
            headers={
                "Authorization": f"Bearer {self.api_key}",
                "Content-Type": "application/json",
            },
            json={
                "model": self.model,
                "messages": messages,
                "max_tokens": max_tokens,
                "temperature": 0.1,
            },
            timeout=300,
        )
        response.raise_for_status()
        return response.json()["choices"][0]["message"]["content"]
    
    def analyze_repo(self) -> Dict:
        """
        第一阶段：全面分析仓库结构
        利用LongCat-2.0的1M上下文窗口，一次性分析整个仓库
        """
        # 收集文件树
        file_tree = []
        for f in self.repo_path.rglob("*.py"):
            rel_path = f.relative_to(self.repo_path)
            if "venv" not in str(rel_path) and "__pycache__" not in str(rel_path):
                file_tree.append({
                    "path": str(rel_path),
                    "size": f.stat().st_size,
                    "content": f.read_text(encoding="utf-8")[:5000],
                })
        
        # 构建分析Prompt
        analysis_prompt = f"""你是一个顶级代码架构师。请分析以下Python项目的代码库，输出：

1. 整体架构与模块划分
2. 核心数据结构与接口定义
3. 依赖关系图（模块间的导入/调用链）
4. 与外部SDK的交互点
5. 潜在的重构风险点

项目文件（共{len(file_tree)}个模块）：
{json.dumps(file_tree[:20], indent=2, ensure_ascii=False)}

请以JSON格式输出分析结果。"""
        
        analysis = self._call_longcat([
            {"role": "system", "content": "你是一位资深代码架构师，精通Python和SDK迁移。"},
            {"role": "user", "content": analysis_prompt},
        ])
        
        self.code_graph = json.loads(analysis)
        return self.code_graph
    
    def generate_refactor_plan(self, new_sdk_doc: str) -> List[Dict]:
        """
        第二阶段：基于新版SDK文档生成重构方案
        利用LongCat-2.0的MOPD架构的Agent+Reasoning专家协作
        """
        plan_prompt = f"""基于以下信息，生成详细的代码重构计划：

旧版代码架构：
{json.dumps(self.code_graph, indent=2, ensure_ascii=False)}

新版SDK文档：
{new_sdk_doc[:30000]}

请输出逐模块的重构方案，包含：
1. 每个模块的修改类型（新增/删除/替换/修改）
2. API映射关系（旧API→新API）
3. 涉及的测试用例
4. 迁移优先级和执行顺序"""
        
        plan_response = self._call_longcat([
            {"role": "system", "content": "你是一位精通SDK迁移的软件工程师。"},
            {"role": "user", "content": plan_prompt},
        ])
        
        self.refactor_plan = json.loads(plan_response)
        return self.refactor_plan
    
    def execute_refactor(self, module_path: str, dry_run: bool = True) -> str:
        """
        第三阶段：执行单模块重构
        dry_run=True时仅预览修改，不写入文件
        """
        full_path = self.repo_path / module_path
        old_code = full_path.read_text(encoding="utf-8")
        
        refactor_prompt = f"""请将以下旧版SDK代码迁移到新版API：

旧代码 ({module_path}):
```python
{old_code}