import Foundation

// ─────────────────────────────────────────────────────────────
// Sampler - Token sampling strategies
// ─────────────────────────────────────────────────────────────

public final class Sampler: @unchecked Sendable {
    
    public init() {}
    
    // ─────────────────────────────────────────────────────────────
    // Sample - Main sampling function
    // ─────────────────────────────────────────────────────────────
    
    public func sample(
        logits: [Float],
        temperature: Float = 1.0,
        topK: Int? = nil,
        topP: Float? = nil,
        filterUnusedTokens: Bool = true,
        unusedTokenRange: Range<Int> = 258000..<259000
    ) -> Int {
        var filteredLogits = logits
        
        // Filter out unused tokens if requested
        if filterUnusedTokens {
            for i in unusedTokenRange {
                if i < filteredLogits.count {
                    filteredLogits[i] = -Float.infinity
                }
            }
        }
        
        // Handle temperature=0.0 (greedy sampling)
        if temperature == 0.0 {
            return greedySample(logits: filteredLogits)
        }
        
        // Apply temperature
        var scaledLogits = filteredLogits.map { $0 / temperature }
        
        // Apply Top-k
        if let k = topK {
            scaledLogits = applyTopK(logits: scaledLogits, k: k)
        }
        
        // Apply Top-p (nucleus)
        if let p = topP {
            scaledLogits = applyTopP(logits: scaledLogits, p: p)
        }
        
        // Convert to probabilities
        let probs = softmax(logits: scaledLogits)
        
        // Random sample
        return randomSample(probs: probs)
    }
    
    // ─────────────────────────────────────────────────────────────
    // Greedy Sample - Maximum probability
    // ─────────────────────────────────────────────────────────────
    
    public func greedySample(logits: [Float]) -> Int {
        var maxValue = logits[0]
        var maxIndex = 0
        
        for i in 1..<logits.count {
            if logits[i] > maxValue {
                maxValue = logits[i]
                maxIndex = i
            }
        }
        
        return maxIndex
    }
    
    // ─────────────────────────────────────────────────────────────
    // Top-k Filtering - Keep top k tokens
    // ─────────────────────────────────────────────────────────────
    
    private func applyTopK(logits: [Float], k: Int) -> [Float] {
        // Find threshold for top-k
        let sorted = logits.sorted(by: >)
        let threshold = sorted[min(k - 1, sorted.count - 1)]
        
        // Filter logits
        return logits.map { logit in
            logit >= threshold ? logit : -Float.infinity
        }
    }
    
    // ─────────────────────────────────────────────────────────────
    // Top-p Filtering - Nucleus sampling
    // ─────────────────────────────────────────────────────────────
    
    private func applyTopP(logits: [Float], p: Float) -> [Float] {
        // Convert to probabilities
        let probs = softmax(logits: logits)
        
        // Sort by probability
        let sortedIndices = probs.indices.sorted { probs[$0] > probs[$1] }
        
        // Find cutoff
        var cumulativeProb: Float = 0.0
        var cutoffIndex = 0
        
        for idx in sortedIndices {
            cumulativeProb += probs[idx]
            if cumulativeProb >= p {
                cutoffIndex = idx
                break
            }
        }
        
        // Filter logits
        return logits.indices.map { i in
            probs[i] >= probs[cutoffIndex] ? logits[i] : -Float.infinity
        }
    }
    
    // ─────────────────────────────────────────────────────────────
    // Softmax - Convert logits to probabilities
    // ─────────────────────────────────────────────────────────────
    
    private func softmax(logits: [Float]) -> [Float] {
        // Find max for numerical stability
        let maxLogit = logits.max() ?? 0
        
        // Compute exp
        let exps = logits.map { exp($0 - maxLogit) }
        
        // Normalize
        let sum = exps.reduce(0, +)
        return exps.map { $0 / sum }
    }
    
    // ─────────────────────────────────────────────────────────────
    // Random Sample - Sample from probability distribution
    // ─────────────────────────────────────────────────────────────
    
    private func randomSample(probs: [Float]) -> Int {
        // Generate random number
        let rand = Float.random(in: 0..<1)
        
        // Find corresponding token
        var cumulative: Float = 0.0
        for i in probs.indices {
            cumulative += probs[i]
            if rand < cumulative {
                return i
            }
        }
        
        // Fallback to last token
        return probs.count - 1
    }
}