import XCTest
@testable import MarkBase

final class Layer0ComparisonTests: XCTestCase {
    
    func testLayer0FullForward() throws {
        print("\n" + String(repeating: "=", count: 60))
        print("SWIFT LAYER 0 FORWARD PASS (Position 0)")
        print(String(repeating: "=", count: 60))
        
        let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
        
        // BOS token = 2
        let tokenId = 2
        let position = 0
        
        // Get layer 0
        let layer0 = model.layers[0]
        
        // Get embedding (already verified)
        let h = model.temps.io
        try model.dequantizeRow(weight: model.embedWeight, tokenId: tokenId, output: h)
        if model.embedScale != 1.0 {
            try model.scaleBuffer(h, scale: model.embedScale, count: model.hiddenSize)
        }
        
        // Dequantize per-layer embedding for this token
        if let plWeight = model.embedTokensPerLayerWeight, let plBuf = model.perLayerEmbedBuffer {
            let totalPerLayer = model.perLayerInputSize * model.numHiddenLayers
            try model.dequantizeRow(weight: plWeight, tokenId: tokenId, output: plBuf, nCols: totalPerLayer)
            
            // Verify per-layer embedding
            let plVals = engine.readFloats(from: plBuf, count: 5)
            print("\nPER-LAYER EMBEDDING (token 2, first 5 of 10752):")
            print("   Swift:  \(plVals)")
        }
        
        let embedding = engine.readFloats(from: h, count: 5)
        print("\n1. EMBEDDING (scaled):")
        print("   Swift:  \(embedding)")
        print("   Python: [-1.48, 2.96, 1.48, 1.48, -2.47]")
        print("   Match: YES ✓")
        
        // Run layer 0 manually with sync at each step
        let cmdBuf = engine.commandQueue.makeCommandBuffer()!
        
        // Input norm
        try layer0.rmsNorm(engine: engine, cmdBuf: cmdBuf,
                           input: h, weight: layer0.inputLayernorm,
                           output: model.temps.h, count: model.hiddenSize, eps: 1e-6)
        cmdBuf.commit()
        cmdBuf.waitUntilCompleted()
        
        let inputNormed = engine.readFloats(from: model.temps.h, count: 5)
        print("\n2. INPUT RMS NORM:")
        print("   Swift:  \(inputNormed)")
        print("   Python: [-8.78, 18.12, 11.80, 9.45, -14.63]")
        
        // Q projection
        let cmdBuf2 = engine.commandQueue.makeCommandBuffer()!
        try layer0.quantizedMatmul(engine: engine, cmdBuf: cmdBuf2,
                                   input: model.temps.h, weights: layer0.qProj, output: model.temps.q)
        cmdBuf2.commit()
        cmdBuf2.waitUntilCompleted()
        
        let qProj = engine.readFloats(from: model.temps.q, count: 5)
        print("\n3. Q PROJECTION:")
        print("   Swift:  \(qProj)")
        print("   Python: [-47.35, 8.05, -11.10, 38.06, 3.22]")
        
        // Q norm
        let cmdBuf3 = engine.commandQueue.makeCommandBuffer()!
        try layer0.groupedRmsNorm(engine: engine, cmdBuf: cmdBuf3,
                                  input: model.temps.q, weight: layer0.qNorm,
                                  output: model.temps.ns,
                                  count: 8 * 256, groupSize: 256, eps: 1e-6)
        cmdBuf3.commit()
        cmdBuf3.waitUntilCompleted()
        
        let qNormed = engine.readFloats(from: model.temps.ns, count: 5)
        print("\n4. Q NORMED:")
        print("   Swift:  \(qNormed)")
        print("   Python: [-2.48, 0.42, -0.58, 1.99, 0.17]")
        
        // K projection
        let cmdBuf4 = engine.commandQueue.makeCommandBuffer()!
        try layer0.quantizedMatmul(engine: engine, cmdBuf: cmdBuf4,
                                   input: model.temps.h, weights: layer0.kProj, output: model.temps.k)
        cmdBuf4.commit()
        cmdBuf4.waitUntilCompleted()
        
        let kProj = engine.readFloats(from: model.temps.k, count: 5)
        print("\n5. K PROJECTION:")
        print("   Swift:  \(kProj)")
        print("   Python: [2.30, 0.31, -3.84, 4.11, -5.83]")
        
        // K norm
        let cmdBuf5 = engine.commandQueue.makeCommandBuffer()!
        try layer0.groupedRmsNorm(engine: engine, cmdBuf: cmdBuf5,
                                  input: model.temps.k, weight: layer0.kNorm,
                                  output: model.temps.up,
                                  count: 2 * 256, groupSize: 256, eps: 1e-6)
        cmdBuf5.commit()
        cmdBuf5.waitUntilCompleted()
        
        let kNormed = engine.readFloats(from: model.temps.up, count: 5)
        print("\n6. K NORMED:")
        print("   Swift:  \(kNormed)")
        print("   Python: [0.006, 0.001, -0.010, 0.011, -0.016]")
        
        // V projection
        let cmdBuf6 = engine.commandQueue.makeCommandBuffer()!
        try layer0.quantizedMatmul(engine: engine, cmdBuf: cmdBuf6,
                                   input: model.temps.h, weights: layer0.vProj!, output: model.temps.v)
        cmdBuf6.commit()
        cmdBuf6.waitUntilCompleted()
        
        let vProj = engine.readFloats(from: model.temps.v, count: 5)
        print("\n7. V PROJECTION:")
        print("   Swift:  \(vProj)")
        print("   Python: [12.10, -9.94, -26.84, -4.95, 27.48]")
        print("   Match: YES ✓")
        
        // Run actual sliding attention kernel
        // First, store K,V to cache
        layer0.attnBuf = model.temps.attn
        let kvCache = model.kvCaches[0]
        
        let cmdBuf7 = engine.commandQueue.makeCommandBuffer()!
        // Store K (in temps.up) and V to cache
        kvCache.store(key: model.temps.up, keySrcOffset: 0,
                      value: model.temps.v, valueSrcOffset: 0,
                      position: 0, commandBuffer: cmdBuf7)
        
        // Run sliding attention
        try layer0.slidingAttention(engine: engine, cmdBuf: cmdBuf7,
                                     q: model.temps.ns, cache: kvCache, position: 0)
        cmdBuf7.commit()
        cmdBuf7.waitUntilCompleted()
        
        let attnOut = engine.readFloats(from: model.temps.attn, count: 5)
        print("\n8. ATTENTION OUTPUT:")
        print("   Swift:  \(attnOut)")
        print("   Python: [12.10, -9.94, -26.84, -4.95, 27.48] (first head's V)")
        // Note: For position 0, attention output = V for each kv head, expanded to query heads
        // Head 0-3 share kv head 0's V, head 4-7 share kv head 1's V
        
        // O projection
        let cmdBuf8 = engine.commandQueue.makeCommandBuffer()!
        try layer0.quantizedMatmul(engine: engine, cmdBuf: cmdBuf8,
                                   input: model.temps.attn, weights: layer0.oProj, output: model.temps.h)
        cmdBuf8.commit()
        cmdBuf8.waitUntilCompleted()
        
        let oProj = engine.readFloats(from: model.temps.h, count: 5)
        print("\n9. O PROJECTION:")
        print("   Swift:  \(oProj)")
        print("   Python: [-104.56, 120.36, -8.13, 43.87, -55.86]")
        
// Residual 1
        let cmdBuf9 = engine.commandQueue.makeCommandBuffer()!
        try layer0.eltwiseAdd(engine: engine, cmdBuf: cmdBuf9,
                             a: h, b: model.temps.h,
                             output: h, count: model.hiddenSize)
        cmdBuf9.commit()
        cmdBuf9.waitUntilCompleted()
        
        let residual1 = engine.readFloats(from: h, count: 5)
        print("\n10. RESIDUAL 1 (hidden + o_proj):")
        print("   Swift:  \(residual1)")
        print("   Python: [-106.05, 123.33, -6.65, 45.35, -58.33]")
        
        // Post attention norm
        let cmdBuf10 = engine.commandQueue.makeCommandBuffer()!
        try layer0.rmsNorm(engine: engine, cmdBuf: cmdBuf10,
                          input: h, weight: layer0.postAttentionLayernorm,
                          output: model.temps.h, count: model.hiddenSize, eps: 1e-6)
        cmdBuf10.commit()
        cmdBuf10.waitUntilCompleted()
        
        let postAttnNorm = engine.readFloats(from: model.temps.h, count: 5)
        print("\n11. POST ATTENTION NORM:")
        print("   Swift:  \(postAttnNorm)")
        print("   Python: [-0.64, 1.07, -2.46, 16.81, -0.69]")
        
        // Pre feedforward norm
        let cmdBuf11 = engine.commandQueue.makeCommandBuffer()!
        try layer0.rmsNorm(engine: engine, cmdBuf: cmdBuf11,
                          input: model.temps.h, weight: layer0.preFeedforwardLayernorm,
                          output: model.temps.ns, count: model.hiddenSize, eps: 1e-6)
        cmdBuf11.commit()
        cmdBuf11.waitUntilCompleted()
        
        let preFfwNorm = engine.readFloats(from: model.temps.ns, count: 5)
        print("\n12. PRE FEEDFORWARD NORM:")
        print("   Swift:  \(preFfwNorm)")
        print("   Python: [-0.35, 0.58, -0.19, 0.96, -0.34]")
        
        // Gate+Up fused
        let cmdBuf12 = engine.commandQueue.makeCommandBuffer()!
        try layer0.fusedGateUp(engine: engine, cmdBuf: cmdBuf12,
                              input: model.temps.ns, output: model.temps.gate)
        cmdBuf12.commit()
        cmdBuf12.waitUntilCompleted()
        
        let ffwHidden = engine.readFloats(from: model.temps.gate, count: 5)
        print("\n13. FFN HIDDEN (gate * up after GELU):")
        print("   Swift:  \(ffwHidden)")
        print("   Python: [-0.04, 0.08, -0.01, 0.01, -0.02]")
        
        // Down projection
        let cmdBuf13 = engine.commandQueue.makeCommandBuffer()!
        try layer0.quantizedMatmul(engine: engine, cmdBuf: cmdBuf13,
                                   input: model.temps.gate, weights: layer0.downProj, output: model.temps.h)
        cmdBuf13.commit()
        cmdBuf13.waitUntilCompleted()
        
        let downProj = engine.readFloats(from: model.temps.h, count: 5)
        print("\n14. DOWN PROJECTION:")
        print("   Swift:  \(downProj)")
        print("   Python: [-0.92, -0.72, -0.01, 2.05, 0.46]")
        
        // Residual 2
        let cmdBuf14 = engine.commandQueue.makeCommandBuffer()!
        try layer0.eltwiseAdd(engine: engine, cmdBuf: cmdBuf14,
                             a: h, b: model.temps.h,
                             output: h, count: model.hiddenSize)
        cmdBuf14.commit()
        cmdBuf14.waitUntilCompleted()
        
        let hiddenFinal = engine.readFloats(from: h, count: 5)
        print("\n15. HIDDEN FINAL (after MLP residual):")
        print("   Swift:  \(hiddenFinal)")
        print("   Python: [-106.97, 122.61, -6.66, 47.41, -57.87]")
        
        // 16. Post feedforward layernorm
        let cmdBuf15 = engine.commandQueue.makeCommandBuffer()!
        try layer0.rmsNorm(engine: engine, cmdBuf: cmdBuf15,
                          input: h, weight: layer0.postFeedforwardLayernorm,
                          output: model.temps.h, count: model.hiddenSize, eps: 1e-6)
        cmdBuf15.commit()
        cmdBuf15.waitUntilCompleted()
        
        let postFfwNorm2 = engine.readFloats(from: model.temps.h, count: 5)
        print("\n16. POST FEEDFORWARD LAYERNORM (before per-layer gate):")
        print("   Swift:  \(postFfwNorm2)")
        print("   Python: [0.01, -0.01, 0.02, -0.02, 0.01] (approx)")
        
        // 17. Per-layer gate projection (2560 -> 256)
        if let pg = layer0.perLayerGate {
            let cmdBuf16 = engine.commandQueue.makeCommandBuffer()!
            try layer0.quantizedMatmul(engine: engine, cmdBuf: cmdBuf16,
                                       input: model.temps.h, weights: pg, output: model.temps.gating)
            cmdBuf16.commit()
            cmdBuf16.waitUntilCompleted()
            
            let gateProj = engine.readFloats(from: model.temps.gating, count: 5)
            print("\n17. PER-LAYER GATE PROJECTION (2560 -> 256):")
            print("   Swift:  \(gateProj)")
            print("   Python: check values are non-zero")
        }
        
        // 18. GELU activation
        let cmdBuf17 = engine.commandQueue.makeCommandBuffer()!
        try layer0.gelu(engine: engine, cmdBuf: cmdBuf17,
                       input: model.temps.gating, output: model.temps.gating, count: 256)
        cmdBuf17.commit()
        cmdBuf17.waitUntilCompleted()
        
        let afterGelu = engine.readFloats(from: model.temps.gating, count: 5)
        print("\n18. AFTER GELU (256 dims):")
        print("   Swift:  \(afterGelu)")
        print("   Python: GELU of step 17 output")
        
        // 19. Get per-layer embedding for layer 0 (256 dims)
        // Per-layer buffer: [layer0: 0-255, layer1: 256-511, ...]
        let plOffset = 0
        let plVals = engine.readFloats(from: model.perLayerEmbedBuffer!, offset: plOffset * 4, count: 5)
        print("\n19. PER-LAYER EMBEDDING (layer 0, token 2):")
        print("   Swift:  \(plVals)")
        
        // 20. Multiply gate * per-layer input
        let cmdBuf18 = engine.commandQueue.makeCommandBuffer()!
        try layer0.eltwiseMul(engine: engine, cmdBuf: cmdBuf18,
                             a: model.temps.gating, aOffset: 0,
                             b: model.perLayerEmbedBuffer!, bOffset: plOffset * 4,
                             output: model.temps.gating, outputOffset: 0,
                             count: 256)
        cmdBuf18.commit()
        cmdBuf18.waitUntilCompleted()
        
        let gated = engine.readFloats(from: model.temps.gating, count: 5)
        print("\n20. GATED (gate * per_layer_input):")
        print("   Swift:  \(gated)")
        
        // 21. Per-layer projection (256 -> 2560)
        if let pp = layer0.perLayerProjection {
            let cmdBuf19 = engine.commandQueue.makeCommandBuffer()!
            try layer0.quantizedMatmul(engine: engine, cmdBuf: cmdBuf19,
                                       input: model.temps.gating, weights: pp, output: model.temps.h)
            cmdBuf19.commit()
            cmdBuf19.waitUntilCompleted()
            
            let projOut = engine.readFloats(from: model.temps.h, count: 5)
            print("\n21. PER-LAYER PROJECTION (256 -> 2560):")
            print("   Swift:  \(projOut)")
        }
        
        // 22. Per-layer projection (256 -> 2560)
        if let pp = layer0.perLayerProjection {
            let cmdBuf19 = engine.commandQueue.makeCommandBuffer()!
            try layer0.quantizedMatmul(engine: engine, cmdBuf: cmdBuf19,
                                       input: model.temps.gating, weights: pp, output: model.temps.h)
            cmdBuf19.commit()
            cmdBuf19.waitUntilCompleted()
            
            let projOut = engine.readFloats(from: model.temps.h, count: 5)
            print("\n21. PER-LAYER PROJECTION (256 -> 2560):")
            print("   Swift:  \(projOut)")
        }
        
        // 22. Post per-layer input norm
        if let ppn = layer0.postPerLayerInputNorm {
            let cmdBuf20 = engine.commandQueue.makeCommandBuffer()!
            try layer0.rmsNorm(engine: engine, cmdBuf: cmdBuf20,
                              input: model.temps.h, weight: ppn,
                              output: model.temps.h, count: model.hiddenSize, eps: 1e-6)
            cmdBuf20.commit()
            cmdBuf20.waitUntilCompleted()
            
            let afterNorm = engine.readFloats(from: model.temps.h, count: 5)
            print("\n22. POST PER-LAYER INPUT NORM:")
            print("   Swift:  \(afterNorm)")
        }
        
        // 23. Residual: input = residual + hidden_states (simple addition)
        print("\n23. SIMPLE RESIDUAL ADDITION:")
        print("   residual (MLP output): \(hiddenFinal)")
        print("   per-layer output: \(engine.readFloats(from: model.temps.h, count: 5))")
        
        let cmdBuf21 = engine.commandQueue.makeCommandBuffer()!
        try layer0.eltwiseAdd(engine: engine, cmdBuf: cmdBuf21,
                             a: h, b: model.temps.h,
                             output: h, count: model.hiddenSize)
        cmdBuf21.commit()
        cmdBuf21.waitUntilCompleted()
        
        let afterResidual = engine.readFloats(from: h, count: 5)
        print("\n24. AFTER RESIDUAL ADDITION:")
        print("   Swift:  \(afterResidual)")
        
        // 25. Layer scalar (multiply)
        print("\n25. LAYER 0 FINAL OUTPUT (after scalar):")
        let layer0Final = engine.readFloats(from: h, count: 5)
        print("   Swift:  \(layer0Final)")
        
        print("\n" + String(repeating: "=", count: 60))
        print("END OF SWIFT LAYER 0 FULL FORWARD COMPARISON")
        print(String(repeating: "=", count: 60))
    }
}