import XCTest
@testable import MarkBase

final class MoEHiddenStateSimpleTest: XCTestCase {
    
    func testHiddenStateSimple() throws {
        print("\n═══════════════════════════════════════")
        print("      Hidden State Simple Test")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-a4b-it-4bit"
        
        print("Step 1: Load model...")
        fflush(stdout)
        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 32)
        print("  ✓ Model loaded (30 layers, MoE)")
        fflush(stdout)
        
        print("\nStep 2: Run forward pass for token 10979 (Hello)...")
        fflush(stdout)
        
        // Run forward pass for position 0
        let logits0 = try model.forward(tokenId: 10979, position: 0, debug: true)
        print("  ✓ Position 0 logits: max=\(logits0.max() ?? 0), min=\(logits0.min() ?? 0)")
        fflush(stdout)
        
        // Check hidden state after position 0
        let h = model.temps.io
        let hs = model.hiddenSize
        let hidden0 = engine.readFloats(from: h, count: min(20, hs))
        print("  Hidden state after pos 0: h[0:20] = \(hidden0)")
        print("  Hidden max/min: \(hidden0.max() ?? 0), \(hidden0.min() ?? 0)")
        fflush(stdout)
        
        // Run forward pass for position 1
        print("\nStep 3: Run forward pass for position 1...")
        fflush(stdout)
        let logits1 = try model.forward(tokenId: 95837, position: 1, debug: false)
        print("  ✓ Position 1 logits: max=\(logits1.max() ?? 0), min=\(logits1.min() ?? 0)")
        fflush(stdout)
        
        // Check hidden state after position 1
        let hidden1 = engine.readFloats(from: h, count: min(20, hs))
        print("  Hidden state after pos 1: h[0:20] = \(hidden1)")
        print("  Hidden max/min: \(hidden1.max() ?? 0), \(hidden1.min() ?? 0)")
        fflush(stdout)
        
        print("\n═══════════════════════════════════════")
        print("✓ Hidden state test completed")
        print("═══════════════════════════════════════\n")
        fflush(stdout)
    }
}