package moe

import (
    "context"
    "fmt"
    "math"
    "sync"
    "time"
    
    "golang.org/x/sync/errgroup"
)

// 专家接口定义
type Expert interface {
    ID() string
    Forward(ctx context.Context, input []float32) ([]float32, error)
    Capacity() int // 当前可用容量
}

// MoE配置
type MoEConfig struct {
    NumExperts       int     // 专家总数
    TopK             int     // 每个token激活的专家数
    ExpertCapacity   int     // 每个专家最大处理token数
    CapacityFactor   float64 // 容量因子，默认1.25
    BalanceCoeff     float64 // 负载均衡系数
    RouterType       string  // 路由类型: "topk", "random", "roundrobin"
}

// 门控网络
type Router struct {
    weights [][]float32 // 门控权重矩阵 [hidden_dim, num_experts]
    bias    []float32   // 偏置项
    config  *MoEConfig
    mu      sync.RWMutex
}

// 创建门控网络
func NewRouter(config *MoEConfig, hiddenDim int) *Router {
    // 初始化权重，使用Xavier初始化
    weights := make([][]float32, hiddenDim)
    scale := float32(math.Sqrt(2.0 / float64(hiddenDim)))
    for i := range weights {
        weights[i] = make([]float32, config.NumExperts)
        for j := range weights[i] {
            weights[i][j] = (float32(math.Rand()) - 0.5) * 2 * scale
        }
    }
    
    return &Router{
        weights: weights,
        bias:    make([]float32, config.NumExperts),
        config:  config,
    }
}

// 路由决策：为每个token选择Top-K专家
func (r *Router) Route(input []float32) ([]int, []float32, error) {
    r.mu.RLock()
    defer r.mu.RUnlock()
    
    // 计算每个专家的得分
    scores := make([]float32, r.config.NumExperts)
    for j := 0; j < r.config.NumExperts; j++ {
        var sum float32
        for i, v := range input {
            sum += v * r.weights[i][j]
        }
        scores[j] = sum + r.bias[j]
    }
    
    // Softmax归一化
    maxScore := float32(-1e9)
    for _, s := range scores {
        if s > maxScore {
            maxScore = s
        }
    }
    
    var sumExp float32
    for i := range scores {
        scores[i] = float32(math.Exp(float64(scores[i] - maxScore)))
        sumExp += scores[i]
    }
    
    if sumExp > 0 {
        for i := range scores {
            scores[i] /= sumExp
        }
    }
    
    // Top-K选择（使用选择排序优化）
    selected := make([]int, 0, r.config.TopK)
    selectedScores := make([]float32, 0, r.config.TopK)
    
    // 复制并排序
    sorted := make([]struct {
        idx   int
        score float32
    }, r.config.NumExperts)
    
    for i, s := range scores {
        sorted[i] = struct {
            idx   int
            score float32
        }{i, s}
    }
    
    // 部分排序，只找Top-K
    for i := 0; i < r.config.TopK; i++ {
        maxIdx := i
        for j := i + 1; j < len(sorted); j++ {
            if sorted[j].score > sorted[maxIdx].score {
                maxIdx = j
            }
        }
        sorted[i], sorted[maxIdx] = sorted[maxIdx], sorted[i]
        selected = append(selected, sorted[i].idx)
        selectedScores = append(selectedScores, sorted[i].score)
    }
    
    return selected, selectedScores, nil
}

// MoE推理引擎
type MoEInference struct {
    config    *MoEConfig
    router    *Router
    experts   []Expert
    stats     *Statistics
    tokenBuf  *sync.Pool // token缓冲区池，减少GC
}

// 统计信息
type Statistics struct {
    mu            sync.Mutex
    totalTokens   int64
    expertLoad    []int64
    routingTime   time.Duration
    forwardTime   time.Duration
}

// 创建MoE推理引擎
func NewMoEInference(config *MoEConfig, experts []Expert) *MoEInference {
    if len(experts) != config.NumExperts {
        panic(fmt.Sprintf("专家数量不匹配: 期望%d, 实际%d", config.NumExperts, len(experts)))
    }
    
    return &MoEInference{
        config:  config,
        router:  NewRouter(config, 768), // 假设hidden_dim=768
        experts: experts,
        stats: &Statistics{
            expertLoad: make([]int64, config.NumExperts),
        },
        tokenBuf: &sync.Pool{
            New: func() interface{} {
                return make([]float32, 0, 1024)
            },
        },
    }
}

// 批量推理入口
func (m *MoEInference) Forward(ctx context.Context, tokens [][]float32) ([][]float32, error) {
    start := time.Now()
    defer func() {
        m.stats.mu.Lock()
        m.stats.routingTime += time.Since(start)
        m.stats.mu.Unlock()
    }()
    
    // 阶段1: 路由决策
    routingResults := make([]struct {
        experts []int
        scores  []float32
    }, len(tokens))
    
    for i, token := range tokens {
        experts, scores, err := m.router.Route(token)
        if err != nil {
            return nil, fmt.Errorf("路由失败 token %d: %w", i, err)
        }
        routingResults[i] = struct {
            experts []int
            scores  []float32
        }{experts, scores}
    }
    
    // 阶段2: 构建专家任务队列
    expertTasks := make([][]struct {
        tokenIdx int
        score    float32
    }, m.config.NumExperts)
    
    for tokenIdx, result := range routingResults {
        for j, expertIdx := range result.experts {
            // 检查专家容量限制
            if len(expertTasks[expertIdx]) < m.config.ExpertCapacity {
                expertTasks[expertIdx] = append(expertTasks[expertIdx], struct {
                    tokenIdx int
                    score    float32
                }{tokenIdx, result.scores[j]})
            }
        }
    }
    
    // 更新负载统计
    m.stats.mu.Lock()
    for i, tasks := range expertTasks {
        m.stats.expertLoad[i] += int64(len(tasks))
    }
    m.stats.totalTokens += int64(len(tokens))
    m.stats.mu.Unlock()
    
    // 阶段3: 并行专家推理
    outputs := make([][]float32, len(tokens))
    var mu sync.Mutex
    
    g, ctx := errgroup.WithContext(ctx)
    g.SetLimit(len(m.experts)) // 限制并发度
    
    for expertIdx, tasks := range expertTasks {
        expertIdx, tasks := expertIdx, tasks
        if len(tasks) == 0 {
            continue
        }
        
        g.Go(func() error {
            expert := m.experts[expertIdx]
            
            // 批量处理该专家的所有token
            batchInput := make([][]float32, len(tasks))
            for i, task := range tasks {
                batchInput[i] = tokens[task.tokenIdx]
            }
            
            // 调用专家前向推理
            batchOutput, err := expert.Forward(ctx, flatten(batchInput))
            if err != nil {
                return fmt.Errorf("专家 %s 推理失败: %w", expert.ID(), err)
            }
            
            // 还原输出形状并加权
            outputDim := len(batchOutput) / len(tasks)
            unflattened := unflatten(batchOutput, len(tasks), outputDim)
            
            mu.Lock()
            for i, task := range tasks {
                // 加权合并：score * expert_output
                weighted := make([]float32, outputDim)
                for j, v := range unflattened[i] {
                    weighted[j] = v * task.score
                }
                outputs[task.tokenIdx] = weighted
            }
            mu.Unlock()
            
            return nil
        })
    }
    
    if err := g.Wait(); err != nil {
        return nil, err
    }
    
    return outputs, nil
}

// 辅助函数：展平二维数组
func flatten(input [][]float32) []float32 {
    total := 0
    for _, v := range input {
        total += len(v)
    }
    result := make([]float32, 0, total)
    for _, v := range input {
        result = append(result, v...)
    }
    return result
}

// 辅助函数：还原二维数组
func unflatten(input []float32, rows, cols int) [][]float32 {
    result := make([][]float32, rows)
    for i := 0; i < rows; i++ {
        result[i] = input[i*cols : (i+1)*cols]
    }
    return result
}

// 获取统计信息
func (m *MoEInference) GetStats() map[string]interface{} {
    m.stats.mu.Lock()
    defer m.stats.mu.Unlock()
    
    stats := make(map[string]interface{})
    stats["total_tokens"] = m.stats.totalTokens
    stats["routing_time_ms"] = m.stats.routingTime.Milliseconds()
    stats["forward_time_ms"] = m.stats.forwardTime.Milliseconds()
    
    // 计算负载均衡指标
    if m.stats.totalTokens > 0 {
        var sum, sumSq float64
        for _, load := range m.stats.expertLoad {
            sum += float64(load)
            sumSq += float64(load) * float64(load)
        }
        mean := sum / float64(len(m.stats.expertLoad))
        variance := sumSq/float64(len(m.stats.expertLoad)) - mean*mean
        stats["load_balance_std"] = math.Sqrt(variance)
    }
    
    return stats
}