import XCTest
@testable import MarkBase

final class G12B31BTests: XCTestCase {

    let modelDir = "/Users/accusys/MarkBaseEngine/models/gemma-4-31b-it-4bit"
    let model26BDir = "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-a4b-it-4bit"
    let model26BStdDir = "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-standard"

    // MARK: - Config Validation

    func test31BConfigValidation() throws {
        guard FileManager.default.fileExists(atPath: modelDir + "/config.json") else {
            print("⚠️ 31B model not found, skipping")
            return
        }

        let cfg = try ModelConfig.load(from: modelDir)

        print("31B Config:")
        print("  hidden_size: \(cfg.hiddenSize ?? -1)")
        print("  num_layers: \(cfg.numHiddenLayers ?? -1)")
        print("  num_heads: \(cfg.numAttentionHeads ?? -1)")
        print("  num_kv_heads: \(cfg.numKeyValueHeads ?? -1)")
        print("  num_global_kv_heads: \(cfg.numGlobalKeyValueHeads ?? -1)")
        print("  intermediate_size: \(cfg.intermediateSize ?? -1)")
        print("  head_dim: \(cfg.headDim ?? -1)")
        print("  global_head_dim: \(cfg.globalHeadDim ?? -1)")
        print("  vocab_size: \(cfg.vocabSize ?? -1)")
        print("  sliding_window: \(cfg.slidingWindow ?? -1)")
        print("  num_kv_shared: \(cfg.numKvSharedLayers ?? -1)")
        print("  attention_k_eq_v: \(cfg.attentionKEqualsV ?? false)")
        print("  final_logit_softcapping: \(cfg.finalLogitSoftcapping ?? -1)")
        print("  layer_types count: \(cfg.layerTypes?.count ?? -1)")

        XCTAssertEqual(cfg.hiddenSize, 5376)
        XCTAssertEqual(cfg.numHiddenLayers, 60)
        XCTAssertEqual(cfg.numAttentionHeads, 32)
        XCTAssertEqual(cfg.numKeyValueHeads, 16)
        XCTAssertEqual(cfg.numGlobalKeyValueHeads, 4)
        XCTAssertEqual(cfg.intermediateSize, 21504)
        XCTAssertEqual(cfg.headDim, 256)
        XCTAssertEqual(cfg.globalHeadDim, 512)
        XCTAssertEqual(cfg.vocabSize, 262144)
        XCTAssertEqual(cfg.slidingWindow, 1024)
        XCTAssertEqual(cfg.numKvSharedLayers, 0)
        XCTAssertEqual(cfg.attentionKEqualsV, true)
        XCTAssertEqual(cfg.finalLogitSoftcapping, 30.0)
        XCTAssertEqual(cfg.layerTypes?.count, 60)

        print("✓ 31B config validation passed")
    }

    func test31BLayerTypeAssignment() throws {
        guard FileManager.default.fileExists(atPath: modelDir + "/config.json") else {
            print("⚠️ 31B model not found, skipping")
            return
        }

        let cfg = try ModelConfig.load(from: modelDir)

        let pattern = cfg.slidingWindowPattern ?? 6
        let lt = cfg.layerTypes!

        // Verify layer type count
        let fullCount = lt.filter { $0 == "full_attention" }.count
        let slidingCount = lt.filter { $0 == "sliding_attention" }.count
        XCTAssertEqual(fullCount, 10)
        XCTAssertEqual(slidingCount, 50)
        XCTAssertEqual(fullCount + slidingCount, 60)

        // Verify pattern: every 6th layer (index 5,11,17,23,29,35,41,47,53,59) is full
        let fullIndices = lt.enumerated().filter { $0.element == "full_attention" }.map { $0.offset }
        let expectedFullIndices = stride(from: pattern - 1, to: 60, by: pattern).map { $0 }
        XCTAssertEqual(fullIndices, expectedFullIndices)

        print("✓ Layer types: \(fullCount) full, \(slidingCount) sliding")
        print("  Full indices: \(fullIndices)")
    }

    // MARK: - Model Loading

    func test31BModelInit() throws {
        guard FileManager.default.fileExists(atPath: modelDir + "/config.json") else {
            print("⚠️ 31B model not found, skipping")
            return
        }

        let engine = try MarkBaseEngine(autoCompile: true)

        let start = Date()
        let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 128)
        let loadTime = Date().timeIntervalSince(start)

        print("31B Model loaded in \(String(format: "%.1f", loadTime))s")
        print("  Layers: \(model.numHiddenLayers)")
        print("  Hidden: \(model.hiddenSize)")
        print("  Vocab: \(model.vocabSize)")
        print("  KV shared: \(model.numKvShared)")
        print("  Full layers: \(model.layerTypesIsFull.filter { $0 }.count)")
        print("  Sliding layers: \(model.layerTypesIsFull.filter { !$0 }.count)")
        print("  Embed scale: \(model.embedScale)")
        print("  Final logit softcapping: \(model.finalLogitSoftcapping ?? 0)")

        XCTAssertEqual(model.numHiddenLayers, 60)
        XCTAssertEqual(model.hiddenSize, 5376)
        XCTAssertEqual(model.vocabSize, 262144)
        XCTAssertEqual(model.numKvShared, 0)
        XCTAssertEqual(model.finalLogitSoftcapping, 30.0)

        // Verify layer-level config
        for (i, layer) in model.layers.enumerated() {
            let isFull = model.layerTypesIsFull[i]
            if isFull {
                // Full attention: vProj should be nil, kEqualsV should be true, global head dim=512
                XCTAssertNil(layer.vProj, "Layer \(i): full attention should have nil vProj")
                XCTAssertEqual(layer.config.headDim, 512, "Layer \(i): full attention head dim should be 512")
                XCTAssertEqual(layer.config.nKvHeads, 4, "Layer \(i): full attention kv heads should be 4")
                XCTAssertEqual(layer.config.nHeads, 32, "Layer \(i): full attention heads should be 32")
            } else {
                // Sliding attention: vProj should exist, head dim=256
                XCTAssertNotNil(layer.vProj, "Layer \(i): sliding attention should have vProj")
                XCTAssertEqual(layer.config.headDim, 256, "Layer \(i): sliding attention head dim should be 256")
                XCTAssertEqual(layer.config.nKvHeads, 16, "Layer \(i): sliding attention kv heads should be 16")
                XCTAssertEqual(layer.config.nHeads, 32, "Layer \(i): sliding attention heads should be 32")
            }
        }

        print("✓ 31B model init and layer config validation passed")
    }

    // MARK: - Forward Pass Correctness

    func test31BForwardDeterminism() throws {
        guard FileManager.default.fileExists(atPath: modelDir + "/config.json") else {
            print("⚠️ 31B model not found, skipping")
            return
        }

        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 128)

        // Run forward pass twice with same input
        let logits1 = try model.forward(tokenId: 2, position: 0)
        let logits2 = try model.forward(tokenId: 2, position: 0)

        // Compare - should be bit-identical
        var diffCount = 0
        for i in 0..<min(logits1.count, 10000) {
            if logits1[i] != logits2[i] { diffCount += 1 }
        }
        XCTAssertEqual(diffCount, 0, "Forward pass should be deterministic")

        print("✓ 31B forward pass is deterministic (0 diffs in first 10000 logits)")
        print("  Max logit: \(logits1.max() ?? -999)")
    }

    func test31BMultiPositionForward() throws {
        guard FileManager.default.fileExists(atPath: modelDir + "/config.json") else {
            print("⚠️ 31B model not found, skipping")
            return
        }

        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 128)

        // Forward at different positions with different tokens
        let positions = [0, 1, 2, 5, 10]
        var allLogits: [[Float]] = []

        for pos in positions {
            let logits = try model.forward(tokenId: 2, position: pos)
            allLogits.append(logits)

            let sorted = logits.enumerated().sorted { $0.element > $1.element }
            let top3 = sorted.prefix(3).map { "\($0.offset):\(String(format: "%.2f", $0.element))" }
            print("  Position \(pos): max=\(String(format: "%.2f", logits.max() ?? 0)), top3=[\(top3.joined(separator: ", "))]")
        }

        // Different positions should produce different logits
        for i in 1..<allLogits.count {
            let first = allLogits[0]
            let current = allLogits[i]
            var same = true
            for j in 0..<min(first.count, 100) {
                if first[j] != current[j] { same = false; break }
            }
            XCTAssertFalse(same, "Position \(positions[i]) should differ from position 0")
        }

        print("✓ 31B multi-position forward: all positions produce different outputs")
    }

    func test31BLogitDistribution() throws {
        guard FileManager.default.fileExists(atPath: modelDir + "/config.json") else {
            print("⚠️ 31B model not found, skipping")
            return
        }

        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 128)

        let logits = try model.forward(tokenId: 2, position: 0)

        let maxVal = logits.max()!
        let minVal = logits.min()!
        let softcap: Float = 30.0

        print("  Logit range: [\(minVal), \(maxVal)]")
        print("  Softcap: ±\(softcap)")

        // With tanh softcapping: output = cap * tanh(logits/cap)
        // This means output is in range [-cap, cap]
        XCTAssertLessThanOrEqual(maxVal, softcap + 1e-4)
        XCTAssertGreaterThanOrEqual(minVal, -softcap - 1e-4)

        // Logits should not all be the same
        let maxAbs = max(abs(maxVal), abs(minVal))
        XCTAssertGreaterThan(maxAbs, 0.1, "Logits should have meaningful distribution")

        // Check top tokens change meaningfully
        let sorted = logits.enumerated().sorted { $0.element > $1.element }
        let top5 = sorted.prefix(5)
        print("  Top 5 tokens: \(top5.map { "\($0.offset):\(String(format: "%.2f", $0.element))" }.joined(separator: ", "))")

        print("✓ 31B logit distribution valid (softcapped at ±\(softcap))")
    }

    // MARK: - KV Cache & k_eq_v

    func test31BkEqVVerification() throws {
        guard FileManager.default.fileExists(atPath: modelDir + "/config.json") else {
            print("⚠️ 31B model not found, skipping")
            return
        }

        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 128)

        // The 31B model has attention_k_eq_v=true, meaning full attention layers
        // have no v_proj and use K_proj output as V.
        // However, K gets RMSNorm applied (k_norm) while V does NOT (no v_norm),
        // so K ≠ V in the KV cache. This is expected behavior.

        let fullLayerIdx = model.layerTypesIsFull.firstIndex(of: true)!
        print("  Full attention layer: \(fullLayerIdx)")

        // 1. Verify vProj is nil for full layers (k_eq_v code path is taken)
        XCTAssertNil(model.layers[fullLayerIdx].vProj,
                     "Full attention layer should have nil vProj")

        // 2. Verify v_norm is nil (no V normalization weights in 31B)
        let layer = model.layers[fullLayerIdx]
        XCTAssertNil(layer.vNorm, "31B model has no v_norm for any layer")

        // 3. Verify kVProj exists (K projection is always present)
        XCTAssertNotNil(layer.kProj.weight,
                        "K projection must be present")

        // 4. Verify KV cache has non-zero content for both K and V
        _ = try model.forward(tokenId: 2, position: 0)
        let cache = model.kvCaches[fullLayerIdx]
        let kPtr = cache.buffer.contents().assumingMemoryBound(to: Float.self)
        let vBase = cache.valueBaseOffset / MemoryLayout<Float>.stride
        let nKvHeads = 4
        let headDim = 512
        let totalKvFloats = nKvHeads * headDim

        var kNonZero = false
        var vNonZero = false
        for i in 0..<totalKvFloats {
            if abs(kPtr[i]) > 1e-6 { kNonZero = true }
            if abs(kPtr[vBase + i]) > 1e-6 { vNonZero = true }
        }
        XCTAssertTrue(kNonZero, "K cache should have non-zero values")
        XCTAssertTrue(vNonZero, "V cache should have non-zero values")

        print("  K/V cache: both non-zero (\(totalKvFloats) floats)")

        // 5. Also verify sliding layer has vProj
        let slidingLayerIdx = model.layerTypesIsFull.firstIndex(of: false)!
        XCTAssertNotNil(model.layers[slidingLayerIdx].vProj,
                        "Sliding attention should have vProj")

        print("  Sliding layers: vProj present, k_eq_v not applicable")
        print("✓ 31B k_eq_v configuration verified")
    }

    func test31BKVCacheIntegrity() throws {
        guard FileManager.default.fileExists(atPath: modelDir + "/config.json") else {
            print("⚠️ 31B model not found, skipping")
            return
        }

        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 128)

        // Forward at position 0 and 1
        _ = try model.forward(tokenId: 2, position: 0)
        _ = try model.forward(tokenId: 100, position: 1)

        // Check cache lengths
        for i in 0..<model.numHiddenLayers {
            let cache = model.kvCaches[i]
            XCTAssertEqual(cache.currentLength, 2, "Layer \(i) cache should have length 2 after 2 forward passes")
        }

        // Read KV from a sliding layer cache
        let slidingIdx = model.layerTypesIsFull.firstIndex(of: false)!
        let cache = model.kvCaches[slidingIdx]
        let kPtr = cache.buffer.contents().assumingMemoryBound(to: Float.self)
        let vBase = cache.valueBaseOffset / MemoryLayout<Float>.stride
        let nKvHeads = 16
        let headDim = 256
        let floatsPerPos = nKvHeads * headDim

        // Position 0 should have non-zero K,V
        var kNonZero = false
        var vNonZero = false
        for i in 0..<min(floatsPerPos, 100) {
            if abs(kPtr[i]) > 1e-6 { kNonZero = true }
            if abs(kPtr[vBase + i]) > 1e-6 { vNonZero = true }
        }
        XCTAssertTrue(kNonZero, "K cache should contain non-zero values")
        XCTAssertTrue(vNonZero, "V cache should contain non-zero values")

        // Position 1 should be at a different offset
        let pos1_k_offset = 1 * floatsPerPos
        var pos1NonZero = false
        for i in 0..<min(floatsPerPos, 100) {
            if abs(kPtr[pos1_k_offset + i]) > 1e-6 { pos1NonZero = true }
        }
        XCTAssertTrue(pos1NonZero, "Position 1 K cache should contain non-zero values")

        // Positions 0 and 1 should have different K values
        var diffCount = 0
        for i in 0..<floatsPerPos {
            if abs(kPtr[i] - kPtr[pos1_k_offset + i]) > 1e-6 { diffCount += 1 }
        }
        XCTAssertGreaterThan(diffCount, 0, "Different positions should have different K values")

        print("✓ 31B KV cache integrity: \(model.numHiddenLayers) caches, all length=2, non-zero content")
    }

    // MARK: - Text Generation

    func test31BTextGeneration() throws {
        guard FileManager.default.fileExists(atPath: modelDir + "/config.json") else {
            print("⚠️ 31B model not found, skipping")
            return
        }

        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 128)
        let tokenizer = try TokenizerFactory.load(modelDir: modelDir)

        // Generate 5 tokens
        let genConfig = GenerationConfig(maxTokens: 5, temperature: 0.7)
        let generator = StreamingGenerator(model: model, tokenizer: tokenizer, engine: engine)

        let start = Date()
        let response = try generator.generateComplete(prompt: "Hello", config: genConfig)
        let elapsed = Date().timeIntervalSince(start)

        print("  31B generation: '\(response)' (\(String(format: "%.1f", elapsed))s)")
        XCTAssertGreaterThan(response.count, 0, "Should generate text")

        // Speed check
        let tokPerSec = 5.0 / elapsed
        print("  Speed: \(String(format: "%.1f", tokPerSec)) tok/s")
        XCTAssertGreaterThan(tokPerSec, 1.0, "Should generate at least 1 tok/s")
    }

    // MARK: - 26B A4B Tests

    func test26BConfigValidation() throws {
        guard FileManager.default.fileExists(atPath: model26BDir + "/config.json") else {
            print("⚠️ 26B A4B model not found, skipping")
            return
        }

        let cfg = try ModelConfig.load(from: model26BDir)

        print("26B A4B Config:")
        print("  hidden_size: \(cfg.hiddenSize ?? -1)")
        print("  num_layers: \(cfg.numHiddenLayers ?? -1)")
        print("  num_heads: \(cfg.numAttentionHeads ?? -1)")
        print("  num_kv_heads: \(cfg.numKeyValueHeads ?? -1)")
        print("  intermediate_size: \(cfg.intermediateSize ?? -1)")
        print("  head_dim: \(cfg.headDim ?? -1)")
        print("  sliding_window: \(cfg.slidingWindow ?? -1)")
        print("  attention_k_eq_v: \(cfg.attentionKEqualsV ?? false)")
        print("  layer_types count: \(cfg.layerTypes?.count ?? -1)")

        XCTAssertEqual(cfg.hiddenSize, 2816)
        XCTAssertEqual(cfg.numHiddenLayers, 30)
        XCTAssertEqual(cfg.numAttentionHeads, 16)
        XCTAssertEqual(cfg.numKeyValueHeads, 8)
        XCTAssertEqual(cfg.numGlobalKeyValueHeads, 2)
        XCTAssertEqual(cfg.intermediateSize, 2112)
        XCTAssertEqual(cfg.headDim, 256)
        XCTAssertEqual(cfg.slidingWindow, 1024)
        XCTAssertEqual(cfg.attentionKEqualsV, true)
        XCTAssertEqual(cfg.layerTypes?.count, 30)

        let fullCount = cfg.layerTypes!.filter { $0 == "full_attention" }.count
        let slidingCount = cfg.layerTypes!.filter { $0 == "sliding_attention" }.count
        XCTAssertEqual(fullCount, 5)
        XCTAssertEqual(slidingCount, 25)

        print("✓ 26B config validation passed")
    }

    func test26BModelLoading() throws {
        guard FileManager.default.fileExists(atPath: model26BDir + "/config.json") else {
            print("⚠️ 26B A4B model not found, skipping")
            return
        }

        let engine = try MarkBaseEngine(autoCompile: true)

        let start = Date()
        let model = try E4BModel(modelDir: model26BDir, engine: engine, maxContextLength: 128)
        let loadTime = Date().timeIntervalSince(start)

        print("26B Model loaded in \(String(format: "%.1f", loadTime))s")
        print("  Layers: \(model.numHiddenLayers)")
        print("  Hidden: \(model.hiddenSize)")

        XCTAssertEqual(model.numHiddenLayers, 30)
        XCTAssertEqual(model.hiddenSize, 2816)

        // Verify k_eq_v on full attention layers
        for (i, layer) in model.layers.enumerated() {
            let isFull = model.layerTypesIsFull[i]
            if isFull {
                XCTAssertNil(layer.vProj, "Layer \(i): full attention nil vProj")
            } else {
                XCTAssertNotNil(layer.vProj, "Layer \(i): sliding attention has vProj")
            }
        }

        // Forward pass
        let logits = try model.forward(tokenId: 2, position: 0)
        print("  Forward pass: max=\(logits.max() ?? -999)")

        // Verify k_eq_v config (26B has 0 v_norm weights, so K != V in cache)
        let fullLayerIdx = model.layerTypesIsFull.firstIndex(of: true)!
        XCTAssertNil(model.layers[fullLayerIdx].vProj,
                     "Full attention layer should have nil vProj")
        XCTAssertNil(model.layers[fullLayerIdx].vNorm,
                     "26B model has no v_norm for any layer")

        print("✓ 26B model loading and forward pass passed")
    }

    // MARK: - 26B Standard Tests

    func test26BStdConfigValidation() throws {
        guard FileManager.default.fileExists(atPath: model26BStdDir + "/config.json") else {
            print("⚠️ 26B Standard model not found, skipping")
            return
        }

        let cfg = try ModelConfig.load(from: model26BStdDir)

        print("26B Standard Config:")
        print("  hidden_size: \(cfg.hiddenSize ?? -1)")
        print("  num_layers: \(cfg.numHiddenLayers ?? -1)")
        print("  num_heads: \(cfg.numAttentionHeads ?? -1)")
        print("  num_kv_heads: \(cfg.numKeyValueHeads ?? -1)")
        print("  num_global_kv_heads: \(cfg.numGlobalKeyValueHeads ?? -1)")
        print("  intermediate_size: \(cfg.intermediateSize ?? -1)")
        print("  head_dim: \(cfg.headDim ?? -1)")
        print("  vocab_size: \(cfg.vocabSize ?? -1)")
        print("  sliding_window: \(cfg.slidingWindow ?? -1)")
        print("  attention_k_eq_v: \(cfg.attentionKEqualsV ?? false)")
        print("  final_logit_softcapping: \(cfg.finalLogitSoftcapping ?? -1)")
        print("  layer_types count: \(cfg.layerTypes?.count ?? -1)")
        print("  enable_moe_block: \(cfg.enableMoEBlock ?? false)")
        print("  num_experts: \(cfg.numExperts ?? -1)")
        print("  top_k_experts: \(cfg.topKExperts ?? -1)")

        // Flat config, no text_config nesting
        XCTAssertEqual(cfg.hiddenSize, 2816)
        XCTAssertEqual(cfg.numHiddenLayers, 30)
        XCTAssertEqual(cfg.numAttentionHeads, 8)
        XCTAssertEqual(cfg.numKeyValueHeads, 2)
        XCTAssertEqual(cfg.intermediateSize, 2112)
        XCTAssertEqual(cfg.headDim, 256)
        XCTAssertEqual(cfg.vocabSize, 262144)

        // These fields are NOT in the 26B-standard config
        // (need to be inferred or defaulted at runtime)
        XCTAssertNil(cfg.numGlobalKeyValueHeads, "26B Std config has no global_head_dim")
        XCTAssertNil(cfg.layerTypes, "26B Std config has no layer_types")
        XCTAssertNil(cfg.attentionKEqualsV, "26B Std config has no attention_k_eq_v")
        XCTAssertNil(cfg.finalLogitSoftcapping, "26B Std config has no final_logit_softcapping")
        XCTAssertNil(cfg.slidingWindow, "26B Std config has no sliding_window")
        XCTAssertNil(cfg.enableMoEBlock, "26B Std config has no enable_moe_block")
        XCTAssertNil(cfg.numExperts, "26B Std config has no num_experts")
        XCTAssertNil(cfg.topKExperts, "26B Std config has no top_k_experts")

        print("✓ 26B Standard config validation passed")
    }

    func test26BStdModelLoading() throws {
        print("=== test26BStdModelLoading START ===")
        fflush(stdout)
        
        guard FileManager.default.fileExists(atPath: model26BStdDir + "/config.json") else {
            print("⚠️ 26B Standard model not found, skipping")
            return
        }

        print("Creating engine...")
        fflush(stdout)
        let engine = try MarkBaseEngine(autoCompile: true)
        print("Engine created")

        let start = Date()
        print("Loading model...")
        let model = try E4BModel(modelDir: model26BStdDir, engine: engine, maxContextLength: 128)
        let loadTime = Date().timeIntervalSince(start)

        print("26B Standard Model loaded in \(String(format: "%.1f", loadTime))s")
        print("  Layers: \(model.numHiddenLayers)")
        print("  Hidden: \(model.hiddenSize)")

        XCTAssertEqual(model.numHiddenLayers, 30)
        XCTAssertEqual(model.hiddenSize, 2816)

        // Forward pass - with timing
        print("Running forward pass...")
        let forwardStart = Date()
        let logits = try model.forward(tokenId: 2, position: 0, debug: false)
        let forwardTime = Date().timeIntervalSince(forwardStart)
        print("  Forward pass completed in \(String(format: "%.1f", forwardTime))s")
        let maxVal = logits.max()!
        print("  Forward pass result: max=\(maxVal)")

        XCTAssertFalse(maxVal.isNaN, "Logits should not be NaN")
        XCTAssertGreaterThan(maxVal, -1e6, "Logits should be meaningful")

        print("✓ 26B Standard model loading and forward pass passed")
    }
}