// Minimal test: only forward pass, no token generation

import Foundation
@testable import MarkBase

let modelPath = "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-a4b-it-4bit"

print("=====================================")
print("26B-A4B MoE Forward Pass Test ONLY")
print("=====================================")

print("\n[1] Loading model...")
let start1 = Date()
let engine = try E4BEngine(autoCompile: true)
print("✓ Engine created")

let model = try E4BModel(modelDir: modelPath, engine: engine, maxContextLength: 128)
let loadTime = Date().timeIntervalSince(start1)
print("✓ Model loaded in \(String(format: "%.3f", loadTime))s")

print("\n[2] Testing single forward pass...")
let tokenizer = try TokenizerFactory.load(modelDir: modelPath)
let prompt = "Hello"
let tokens = tokenizer.encode(text: prompt)
print("  Prompt: \"\(prompt)\" -> tokens: \(tokens)")

// Create input buffer
let embed = model.embedTokens
let inputBuffer = engine.createBuffer(length: model.hiddenSize * 4)
let outputBuffer = engine.createBuffer(length: model.vocabSize * 4)

// Get embedding for first token
let embedData = engine.readFloats(from: embed.weight, offset: tokens[0] * model.hiddenSize, count: model.hiddenSize)
engine.writeFloats(embedData, to: inputBuffer)

print("  Input buffer: \(model.hiddenSize) floats")
print("  Running forward pass through \(model.numHiddenLayers) layers...")

let start2 = Date()

// Run through all layers
for i in 0..<model.numHiddenLayers {
    let layer = model.layers[i]
    let normBuffer = engine.createBuffer(length: model.hiddenSize * 4)
    
    // Apply input layernorm
    try layer.inputLayernorm.process(input: inputBuffer, output: normBuffer, engine: engine)
    
    // Check for NaN in layer norm output
    let normOutput = engine.readFloats(from: normBuffer, count: 10)
    let hasNaN = normOutput.contains { $0.isNaN }
    if hasNaN {
        print("  ❌ NaN detected in layer \(i) after input_layernorm")
        print("     First 10 values: \(normOutput)")
        break
    }
    
    if i % 5 == 0 {
        print("    Layer \(i): OK (max=\(normOutput.max() ?? 0), min=\(normOutput.min() ?? 0))")
    }
    
    // Clean up
    engine.releaseBuffer(normBuffer)
}

let forwardTime = Date().timeIntervalSince(start2)
print("✓ Forward pass completed in \(String(format: "%.3f", forwardTime))s")

print("\n✅ Forward pass test completed!")