import XCTest
@testable import MarkBase

final class ExpertKernelParametersTest: XCTestCase {
    
    func testExpertKernelParameters() throws {
        print("\n═══════════════════════════════════════")
        print("      Expert Kernel Parameters Check")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-a4b-it-4bit"
        
        print("Step 1: Load model...")
        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 32)
        print("  ✓ Model loaded")
        
        print("\nStep 2: Check expert parameters...")
        let layer0 = model.layers[0]
        
        guard let eGate = layer0.expertGate, let eUp = layer0.expertUp else {
            XCTFail("Experts not present")
            return
        }
        
        print("  ExpertGate parameters:")
        print("    numExperts: \(eGate.numExperts)")
        print("    expertOutDim: \(eGate.expertOutDim)")
        print("    expertInDim: \(eGate.expertInDim)")
        print("    bits: \(eGate.bits)")
        print("    weightStride: \(eGate.weightStride) bytes")
        print("    scalesStride: \(eGate.scalesStride) bytes")
        print("    weight buffer size: \(eGate.weight.length) bytes")
        print("    scales buffer size: \(eGate.scales.length) bytes")
        
        print("\n  ExpertUp parameters:")
        print("    numExperts: \(eUp.numExperts)")
        print("    expertOutDim: \(eUp.expertOutDim)")
        print("    expertInDim: \(eUp.expertInDim)")
        print("    bits: \(eUp.bits)")
        print("    weightStride: \(eUp.weightStride) bytes")
        print("    scalesStride: \(eUp.scalesStride) bytes")
        
        // Calculate expected sizes
        let expectedGateWeightStride = eGate.expertOutDim * (eGate.expertInDim * eGate.bits / 32) * 4
        let expectedGateScalesStride = eGate.expertOutDim * (eGate.expertInDim / 64) * 4
        
        print("\n  Stride verification:")
        print("    Gate weightStride expected: \(expectedGateWeightStride) bytes")
        print("    Gate weightStride actual: \(eGate.weightStride) bytes")
        print("    Match: \(eGate.weightStride == expectedGateWeightStride ? "✓" : "❌")")
        
        print("    Gate scalesStride expected: \(expectedGateScalesStride) bytes")
        print("    Gate scalesStride actual: \(eGate.scalesStride) bytes")
        print("    Match: \(eGate.scalesStride == expectedGateScalesStride ? "✓" : "❌")")
        
        print("\n═══════════════════════════════════════")
        print("✓ Expert parameters checked")
        print("═══════════════════════════════════════\n")
    }
    
    func testExpertBufferSizes() throws {
        print("\n═══════════════════════════════════════")
        print("      Expert Buffer Size Check")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-a4b-it-4bit"
        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 32)
        
        let layer0 = model.layers[0]
        guard let eGate = layer0.expertGate, let eUp = layer0.expertUp else {
            XCTFail("Experts not present")
            return
        }
        
        let hs = model.hiddenSize
        let moeIntermediate = eGate.expertOutDim
        
        print("Expected buffer sizes:")
        print("  Input (hiddenSize): \(hs) floats = \(hs * 4) bytes")
        print("  Gate+Up output: \(2 * moeIntermediate) floats = \(2 * moeIntermediate * 4) bytes")
        print("  Gate output only: \(moeIntermediate) floats = \(moeIntermediate * 4) bytes")
        print("  Up output only: \(moeIntermediate) floats = \(moeIntermediate * 4) bytes")
        
        print("\nCreating buffers...")
        
        // Create buffers with correct sizes
        let inputBuffer = engine.device.makeBuffer(length: hs * 4)!
        let gateUpOutputBuffer = engine.device.makeBuffer(length: 2 * moeIntermediate * 4)!
        let gateOnlyOutputBuffer = engine.device.makeBuffer(length: moeIntermediate * 4)!
        
        print("  ✓ Input buffer: \(inputBuffer.length) bytes")
        print("  ✓ Gate+Up buffer: \(gateUpOutputBuffer.length) bytes")
        print("  ✓ Gate only buffer: \(gateOnlyOutputBuffer.length) bytes")
        
        // Verify sizes match expectations
        XCTAssertEqual(inputBuffer.length, hs * 4, "Input buffer size should be hiddenSize * 4")
        XCTAssertEqual(gateUpOutputBuffer.length, 2 * moeIntermediate * 4, "Gate+Up buffer should be 2 * moeIntermediate * 4")
        XCTAssertEqual(gateOnlyOutputBuffer.length, moeIntermediate * 4, "Gate buffer should be moeIntermediate * 4")
        
        print("\n═══════════════════════════════════════")
        print("✓ Buffer sizes verified correct")
        print("═══════════════════════════════════════\n")
    }
}