import Foundation
import MLX
import MLXEmbedders

/// Runs dummy forward passes at every (batch, length) bucket combination to
/// pre-compile the MLX graph. MLX recompiles on shape changes, so we eat the
/// cost once per shape up front instead of on the user's first request.
enum WarmUp {
    static func warm(
        container: MLXEmbedders.ModelContainer,
        maxBatch: Int,
        maxLen: Int
    ) async {
        let lengths = Buckets.lengths.filter { $0 <= maxLen }
        let batches = Buckets.batches.filter { $0 <= maxBatch }

        await container.perform { model, tokenizer, pooling in
            let padId: Int32 = {
                if let eos = tokenizer.eosTokenId { return Int32(eos) }
                return 0
            }()
            for b in batches {
                for l in lengths {
                    let ids = MLXArray([Int32](repeating: padId, count: b * l), [b, l])
                    let mask = MLXArray([Int32](repeating: 1, count: b * l), [b, l])
                    let out = model(ids, positionIds: nil, tokenTypeIds: nil, attentionMask: mask)
                    let pooled = pooling(out, mask: mask, normalize: true)
                    eval(pooled)
                }
            }
        }
    }
}