import XCTest
import CoreImage
@testable import MarkBase

final class E4BSimpleInferenceTest: XCTestCase {
    
    func testKVCacheDebug() throws {
        // Load model
        let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        print("Loading engine...")
        let engine = try MarkBaseEngine(autoCompile: true)
        print("Loading model...")
        let model = try E4BModel(modelDir: modelDir, engine: engine)
        
        print("\n=== KV Cache Debug Test ===")
        print("Running forward at position 0...")
        
        // Process position 0 and 1 with debug output
        let logits0 = try model.forward(tokenId: 2, position: 0, debug: true)  // BOS
        print("Position 0 logits: max=\(logits0.max() ?? 0), min=\(logits0.min() ?? 0)")
        
        // Check KV cache for layer 0
        let cache0 = model.kvCaches[0]
        print("Layer 0 cache after pos0: currentLength=\(cache0.currentLength)")
        let k0Ptr = cache0.buffer.contents().bindMemory(to: Float.self, capacity: 20)
        let k0vals = [k0Ptr[0], k0Ptr[1], k0Ptr[2], k0Ptr[3], k0Ptr[4], k0Ptr[5], k0Ptr[6], k0Ptr[7], k0Ptr[8], k0Ptr[9]]
        print("K[0, 0:10] = \(k0vals)")
        
        print("---")
        let logits1 = try model.forward(tokenId: 9259, position: 1, debug: true)  // "Hello"
        print("Position 1 logits: max=\(logits1.max() ?? 0), min=\(logits1.min() ?? 0)")
        
        print("Layer 0 cache after pos1: currentLength=\(cache0.currentLength)")
        let k1vals = [k0Ptr[0], k0Ptr[1], k0Ptr[2], k0Ptr[3], k0Ptr[4], k0Ptr[5], k0Ptr[6], k0Ptr[7], k0Ptr[8], k0Ptr[9]]
        let k1posvals = [k0Ptr[256], k0Ptr[257], k0Ptr[258], k0Ptr[259], k0Ptr[260], k0Ptr[261], k0Ptr[262], k0Ptr[263], k0Ptr[264], k0Ptr[265]]
        print("K[0, 0:10] = \(k1vals)")
        print("K[1, 0:10] = \(k1posvals)")
        
        // Check shared layer's source
        if let src = model.kvSourceMap[24] {
            print("Layer 24 (shared) reads from layer \(src)")
            let cacheSrc = model.kvCaches[src]
            print("Layer \(src) cache: currentLength=\(cacheSrc.currentLength)")
            print("Layer 24 cache: currentLength=\(model.kvCaches[24].currentLength)")
        }
        
        // Check all kvSourceMap entries
        print("\nKV source map (shared layers):")
        for layer in 24..<42 {
            let src = model.kvSourceMap[layer] ?? -1
            print("  Layer \(layer) → reads from layer \(src)")
        }
        
        print("---")
        let _ = try model.forward(tokenId: 22470, position: 2, debug: true)  // "Character"
    }
    
    func testLongerPrompt() throws {
        Swift.print("\n=== Longer Prompt Test ===")
        let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
        let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
        let sampler = Sampler()
        
        // Process 10 positions (BOS + 9 random tokens)
        var tokens: [Int] = [2]  // BOS
        for i in 0..<9 {
            let tokenId = 1000 + i  // Some random tokens
            tokens.append(tokenId)
        }
        
        Swift.print("Processing \(tokens.count) tokens...")
        for (i, token) in tokens.enumerated() {
            let _ = try model.forward(tokenId: token, position: i)
        }
        
        // Check cache lengths after processing
        Swift.print("\nCache lengths after processing 10 tokens:")
        for i in [0, 5, 11, 17, 23] {  // Full attention owners
            let cache = model.kvCaches[i]
            Swift.print("  Layer \(i) (full) cache: currentLength=\(cache.currentLength)")
        }
        for i in [1, 10, 22] {  // Sliding attention owners
            let cache = model.kvCaches[i]
            Swift.print("  Layer \(i) (sliding) cache: currentLength=\(cache.currentLength)")
        }
        
        // Get logits at position 9 (last position)
        let logits = try model.forward(tokenId: tokens.last ?? 2, position: tokens.count - 1)
        Swift.print("\nTop 10 logits at position \(tokens.count - 1):")
        let indexed = logits.enumerated().sorted { $0.element > $1.element }.prefix(10)
        for (idx, val) in indexed {
            let tokenStr = tokenizer.decode(tokens: [idx])
            Swift.print("  Token \(idx) ('\(tokenStr)'): \(val)")
        }
    }
    
    func testSimpleTokenGeneration() throws {
        print("\n=== Simple Token Generation Test ===")
        print("Testing E4B model with proper sampling\n")
        
        let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        
        // Load engine and model
        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
        
        // Load tokenizer
        let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
        
        // Create sampler
        let sampler = Sampler()
        
        print("Model loaded, vocab size: \(tokenizer.vocabSize)")
        print("Sampler settings: temperature=0.7, topK=50, topP=0.95")
        
        var lastLogits: [Float]? = nil
        
        // Test 1: Simple BOS → generation
        print("\n--- Test 1: BOS only ---")
        var tokens: [Int] = [2]  // BOS
        // Process BOS and capture logits
        lastLogits = try model.forward(tokenId: 2, position: 0)
        // Sample first generated token from logits at position 0
        if let logits = lastLogits {
            let sampled = sampler.sample(logits: logits, temperature: 0.7, topK: 50, topP: 0.95)
            tokens.append(sampled)
        }
        // Generate continuation
        for _ in 1..<15 {
            guard let lastToken = tokens.last else { break }
            if lastToken == 1 || lastToken == 106 { break }  // EOS
            let logits = try model.forward(tokenId: lastToken, position: tokens.count - 1)
            let sampled = sampler.sample(logits: logits, temperature: 0.7, topK: 50, topP: 0.95)
            tokens.append(sampled)
            if sampled == 1 || sampled == 106 { break }
        }
        print("Generated: '\(tokenizer.decode(tokens: tokens))'")
        
        // Test 2: Hello prompt
        print("\n--- Test 2: 'Hello' prompt ---")
        let model2 = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
        tokens = [2, 9259]  // BOS + "Hello"
        // Process prompt and capture logits
        lastLogits = nil
        for (i, token) in tokens.enumerated() {
            lastLogits = try model2.forward(tokenId: token, position: i)
        }
        // Sample first generated token from logits at position P-1
        if let logits = lastLogits {
            let sampled = sampler.sample(logits: logits, temperature: 0.7, topK: 50, topP: 0.95)
            tokens.append(sampled)
        }
        // Generate continuation
        for _ in 1..<20 {
            guard let lastToken = tokens.last else { break }
            if lastToken == 1 || lastToken == 106 { break }
            let logits = try model2.forward(tokenId: lastToken, position: tokens.count - 1)
            let sampled = sampler.sample(logits: logits, temperature: 0.7, topK: 50, topP: 0.95)
            tokens.append(sampled)
            if sampled == 1 || sampled == 106 { break }
        }
        print("Generated: '\(tokenizer.decode(tokens: tokens))'")
        
        // Test 3: Full prompt format
        print("\n--- Test 3: Full prompt ---")
        let model3 = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
        let prompt = "<start_of_turn>user\nWhat is the capital of France?<end_of_turn>\n<start_of_turn>model\n"
        let promptTokens = tokenizer.encode(text: prompt)
        print("Prompt: \(promptTokens.count) tokens")
        print("Prompt text: '\(prompt)'")
        // Process prompt
        lastLogits = nil
        for (i, token) in promptTokens.enumerated() {
            lastLogits = try model3.forward(tokenId: token, position: i)
        }
        // Generate continuation
        tokens = promptTokens
        // Sample first generated token from logits at position P-1
        if let logits = lastLogits {
            print("Top 10 logits at position \(promptTokens.count - 1):")
            let indexed = logits.enumerated().sorted { $0.element > $1.element }.prefix(10)
            for (idx, val) in indexed {
                let tokenStr = tokenizer.decode(tokens: [idx])
                print("  Token \(idx) ('\(tokenStr)'): \(val)")
            }
        }
    }
    
    func testShortPrompt() throws {
        Swift.print("\n=== Short Prompt Test ===")
        let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
        let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
        
        // Simple prompt: "Hello"
        let prompt = "Hello"
        let promptTokens = tokenizer.encode(text: prompt)
        Swift.print("Prompt: \(promptTokens.count) tokens, tokens: \(promptTokens)")
        
        // Process prompt
        for (i, token) in promptTokens.enumerated() {
            let _ = try model.forward(tokenId: token, position: i)
        }
        
        // Get logits at last position
        let logits = try model.forward(tokenId: promptTokens.last ?? 2, position: promptTokens.count - 1)
        Swift.print("\nTop 10 logits at position \(promptTokens.count - 1):")
        let indexed = logits.enumerated().sorted { $0.element > $1.element }.prefix(10)
        for (idx, val) in indexed {
            let tokenStr = tokenizer.decode(tokens: [idx])
            Swift.print("  Token \(idx) ('\(tokenStr)'): \(val)")
        }
        
        // Decode full output
        var tokens = promptTokens
        let sampler = Sampler()
        
        Swift.print("\n=== Generation trace ===")
let sampled = sampler.sample(logits: logits, temperature: 0.7, topK: 50, topP: 0.95)
        tokens.append(sampled)
        let decoded0 = tokenizer.decode(tokens: [sampled])
        Swift.print("Position \(tokens.count - 1): sampled token \(sampled) ('\(decoded0)')")
        
        for _ in 1..<10 {
            guard let lastToken = tokens.last else { break }
            if lastToken == 1 || lastToken == 106 { break }
            let pos = tokens.count - 1
            let newLogits = try model.forward(tokenId: lastToken, position: pos)
            
            // Show top 5 logits at each position
            let top5 = newLogits.enumerated().sorted { $0.element > $1.element }.prefix(5)
            Swift.print("Position \(pos): top logits:")
            for (idx, val) in top5 {
                let tokenStr = tokenizer.decode(tokens: [idx])
                Swift.print("  Token \(idx) ('\(tokenStr)'): \(val)")
            }
            
            let newSampled = sampler.sample(logits: newLogits, temperature: 0.7, topK: 50, topP: 0.95)
            tokens.append(newSampled)
            let decodedNew = tokenizer.decode(tokens: [newSampled])
            Swift.print("  Sampled: token \(newSampled) ('\(decodedNew)')")
            if newSampled == 1 || newSampled == 106 { break }
        }
        Swift.print("\nFinal output: '\(tokenizer.decode(tokens: tokens))'")
    }
    
    func testPosition1ForwardPass() throws {
        Swift.print("\n=== Position 1 Forward Pass Verification ===")
        let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
        
        for cache in model.kvCaches {
            cache.reset()
        }
        
        // Process position 0 (BOS)
        Swift.print("\nPosition 0: BOS token (2)")
        let logits0 = try model.forward(tokenId: 2, position: 0)
        let top5_0 = logits0.enumerated().sorted { $0.element > $1.element }.prefix(5)
        Swift.print("Top 5 logits:")
        for (idx, val) in top5_0 {
            Swift.print("  Token \(idx): \(val)")
        }
        
        // Check KV cache state after position 0
        Swift.print("\nKV cache state after position 0:")
        for i in 0..<5 {
            let cache = model.kvCaches[i]
            Swift.print("  Layer \(i) cache: currentLength=\(cache.currentLength), maxLength=\(cache.maxLength)")
        }
        
        // Process position 1 (token 568 = '(')
        Swift.print("\nPosition 1: Token 568 ('(')")
        let logits1 = try model.forward(tokenId: 568, position: 1)
        let top5_1 = logits1.enumerated().sorted { $0.element > $1.element }.prefix(5)
        Swift.print("Top 5 logits:")
        for (idx, val) in top5_1 {
            Swift.print("  Token \(idx): \(val)")
        }
        
        // Check KV cache state after position 1
        Swift.print("\nKV cache state after position 1:")
        for i in 0..<5 {
            let cache = model.kvCaches[i]
            Swift.print("  Layer \(i) cache: currentLength=\(cache.currentLength), maxLength=\(cache.maxLength)")
        }
        
        // Process position 2 (greedy choice)
        let greedy1 = logits1.enumerated().max(by: { $0.element < $1.element })!.offset
        Swift.print("\nPosition 2: Token \(greedy1)")
        let logits2 = try model.forward(tokenId: greedy1, position: 2)
        let top5_2 = logits2.enumerated().sorted { $0.element > $1.element }.prefix(5)
        Swift.print("Top 5 logits:")
        for (idx, val) in top5_2 {
            Swift.print("  Token \(idx): \(val)")
        }
        
        // Check if there's a repeating pattern in logits
        Swift.print("\nComparing logits across positions:")
        Swift.print("  Position 0 top token: \(top5_0.first!.offset)")
        Swift.print("  Position 1 top token: \(top5_1.first!.offset)")
        Swift.print("  Position 2 top token: \(top5_2.first!.offset)")
        
        // Check magnitude of hidden states
        Swift.print("\nHidden state magnitudes (checking for explosion):")
        // We can check by looking at the logits values - if they're close to softcapping limit (30), that's normal
        let maxLogit0 = logits0.max()!
        let maxLogit1 = logits1.max()!
        let maxLogit2 = logits2.max()!
        Swift.print("  Position 0 max logit: \(maxLogit0) (softcapped to ~30)")
        Swift.print("  Position 1 max logit: \(maxLogit1)")
        Swift.print("  Position 2 max logit: \(maxLogit2)")
        
        // Check if model enters repeating pattern
        if top5_1.first!.offset == top5_2.first!.offset {
            Swift.print("\n⚠️  WARNING: Same token predicted at position 1 and 2 - potential loop!")
        }
        
        // Check for abnormal patterns
        let logitRange0 = logits0.max()! - logits0.min()!
        let logitRange1 = logits1.max()! - logits1.min()!
        let logitRange2 = logits2.max()! - logits2.min()!
        Swift.print("\nLogit ranges:")
        Swift.print("  Position 0: \(logitRange0)")
        Swift.print("  Position 1: \(logitRange1)")
        Swift.print("  Position 2: \(logitRange2)")
        
        // Very narrow logit range could indicate degenerate state
        if logitRange2 < 1.0 {
            Swift.print("\n⚠️  WARNING: Very narrow logit range at position 2 - degenerate state!")
        }
    }
    
    func testRealImageProperPreprocessing() throws {
        Swift.print("\n=== Real Image with Proper Preprocessing ===")
        
        let imagePath = "/Users/accusys/MarkBase/tests/images/test_image.png"
        
        // Load and preprocess image using CoreImage
        Swift.print("Loading image: \(imagePath)")
        guard let image = CIImage(contentsOf: URL(fileURLWithPath: imagePath)) else {
            Swift.print("ERROR: Cannot load image")
            return
        }
        
        Swift.print("  Original image extent: \(image.extent)")
        
        // Resize to 224x224 (standard ViT input size)
        let targetSize = CGSize(width: 224, height: 224)
        let resizeFilter = CIFilter(name: "CILanczosScaleTransform")!
        resizeFilter.setValue(image, forKey: kCIInputImageKey)
        resizeFilter.setValue(224.0 / image.extent.width, forKey: kCIInputScaleKey)
        resizeFilter.setValue(1.0, forKey: kCIInputAspectRatioKey)
        
        guard let resizedImage = resizeFilter.outputImage else {
            Swift.print("ERROR: Cannot resize image")
            return
        }
        
        Swift.print("  Resized image: \(resizedImage.extent)")
        
        // Convert to pixel data
        let context = CIContext()
        let bitmap = context.createCGImage(resizedImage, from: resizedImage.extent)!
        
        // Extract RGB pixel values
        let dataProvider = bitmap.dataProvider!
        let pixelData = dataProvider.data!
        let ptr = CFDataGetBytePtr(pixelData)!
        let length = CFDataGetLength(pixelData)
        
        Swift.print("  Pixel data length: \(length) bytes")
        Swift.print("  Expected: \(224 * 224 * 4) bytes (RGBA)")
        
        // Convert to normalized floats [0, 1]
        // Extract RGB channels (skip alpha)
        var normalizedPixels = [Float](repeating: 0, count: 224 * 224 * 3)
        
        // Config says standardize=false, so use raw [0,1] pixels
        for i in 0..<224*224 {
            let offset = i * 4  // RGBA
            let r = Float(ptr[offset]) / 255.0
            let g = Float(ptr[offset + 1]) / 255.0
            let b = Float(ptr[offset + 2]) / 255.0
            
            normalizedPixels[i * 3] = r
            normalizedPixels[i * 3 + 1] = g
            normalizedPixels[i * 3 + 2] = b
        }
        
        Swift.print("  ✓ Extracted \(normalizedPixels.count) pixel values (standardize=false)")
        Swift.print("  Sample: r=\(normalizedPixels[0]), g=\(normalizedPixels[1]), b=\(normalizedPixels[2])")
        
        // Create patch embeddings
        let patchSize = 16
        let numPatchesPerRow = 224 / patchSize  // 14
        let numPatches = numPatchesPerRow * numPatchesPerRow  // 196
        let hiddenSize = 768
        
        Swift.print("\nCreating \(numPatches) patch embeddings (16x16 patches):")
        
        // Each patch: 16*16 pixels * 3 channels = 768 floats
        // This matches hiddenSize of vision tower!
        var patchEmbeddings = [Float](repeating: 0, count: numPatches * hiddenSize)
        
        for patchIdx in 0..<numPatches {
            let patchRow = patchIdx / numPatchesPerRow
            let patchCol = patchIdx % numPatchesPerRow
            
            for y in 0..<patchSize {
                for x in 0..<patchSize {
                    // Global pixel coordinates
                    let globalY = patchRow * patchSize + y
                    let globalX = patchCol * patchSize + x
                    let pixelIdx = globalY * 224 + globalX
                    
                    // RGB values
                    let r = normalizedPixels[pixelIdx * 3]
                    let g = normalizedPixels[pixelIdx * 3 + 1]
                    let b = normalizedPixels[pixelIdx * 3 + 2]
                    
                    // Embedding position
                    let embedIdx = patchIdx * hiddenSize + (y * patchSize + x) * 3
                    patchEmbeddings[embedIdx] = r
                    patchEmbeddings[embedIdx + 1] = g
                    patchEmbeddings[embedIdx + 2] = b
                }
            }
        }
        
        Swift.print("  ✓ Created \(patchEmbeddings.count) patch embedding values")
        Swift.print("  First patch sample: \(patchEmbeddings[0..<10]))")
        
        // Load model and run inference
        let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        let engine = try MarkBaseEngine(autoCompile: true)
        let multimodalModel = try MultimodalModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
        let inference = try MultimodalInference(model: multimodalModel)
        let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
        
        Swift.print("\nProcessing through VisionTower...")
        
        // Create buffers
        let visionBuffer = engine.device.makeBuffer(bytes: patchEmbeddings, length: patchEmbeddings.count * 4)!
        let visionOutputBuffer = engine.device.makeBuffer(length: numPatches * 2560 * 4)!
        
        if let tower = multimodalModel.visionTowerFull {
            try tower.forward(patchEmbeddings: visionBuffer, numPatches: numPatches, outputBuffer: visionOutputBuffer)
            Swift.print("  ✓ VisionTower forward complete")
            
            // Pool embeddings
            Swift.print("\nPooling vision features...")
            let visionOutputPtr = visionOutputBuffer.contents().assumingMemoryBound(to: Float.self)
            var pooledEmbedding = [Float](repeating: 0, count: 2560)
            
            for i in 0..<2560 {
                var sum: Float = 0
                for p in 0..<numPatches {
                    sum += visionOutputPtr[p * 2560 + i]
                }
                pooledEmbedding[i] = sum / Float(numPatches)
            }
            
            let mag = sqrt(pooledEmbedding.reduce(0) { $0 + $1 * $1 })
            Swift.print("  ✓ Pooled embedding magnitude: \(mag)")
            
            // Normalize to match text embedding magnitude (~5)
            // Text embeddings have norm ~5, vision has norm ~2130
            // Scale down by factor of ~400 to match
            let targetMag: Float = 5.0
            let scale = targetMag / mag
            Swift.print("  Normalizing vision features: scaling by \(scale)")
            
            for i in 0..<2560 {
                pooledEmbedding[i] *= scale
            }
            
            let normalizedMag = sqrt(pooledEmbedding.reduce(0) { $0 + $1 * $1 })
            Swift.print("  ✓ Normalized magnitude: \(normalizedMag)")
            
            // Generate response
            let prompt = "Describe what you see in this image"
            let promptTokens = tokenizer.encode(text: prompt)
            
            Swift.print("\nPrompt: '\(prompt)'")
            Swift.print("Generating response...")
            
            let pooledBuffer = engine.device.makeBuffer(bytes: pooledEmbedding, length: pooledEmbedding.count * 4)!
            
            let generatedTokens = try inference.generate(
                textTokens: promptTokens,
                precomputedVisionEmbedding: pooledBuffer,
                maxTokens: 50
            )
            
            let decoded = tokenizer.decode(tokens: generatedTokens)
            Swift.print("\nFull output: '\(decoded)'")
            
            // Extract response
            let responseStart = promptTokens.count + 4  // Skip prompt + BOI + IMAGE + EOI
            if generatedTokens.count > responseStart {
                let responseTokens = Array(generatedTokens[responseStart...])
                let response = tokenizer.decode(tokens: responseTokens)
                Swift.print("\nResponse: '\(response)'")
                
                // Check if response makes sense
                Swift.print("\n=== Analysis ===")
                Swift.print("Vision features magnitude: \(mag)")
                Swift.print("Generated tokens count: \(responseTokens.count)")
                Swift.print("Note: Response quality depends on:")
                Swift.print("  - Proper vision preprocessing")
                Swift.print("  - Learned patch projection (we used raw pixels)")
                Swift.print("  - Model training quality")
            }
        } else {
            Swift.print("ERROR: Vision tower not loaded")
        }
    }
    
    func testRealImageInference() throws {
        Swift.print("\n=== Real Image Multimodal Inference Test ===")
        
        let imagePath = "/Users/accusys/MarkBase/tests/images/test_image.png"
        Swift.print("Loading image: \(imagePath)")
        
        // Check if image exists
        guard FileManager.default.fileExists(atPath: imagePath) else {
            Swift.print("ERROR: Test image not found at \(imagePath)")
            return
        }
        
        // Load image and create patch embeddings
        Swift.print("\nProcessing image into patch embeddings...")
        
        // Simple patch embedding: resize to 224x224, split into 16x16 patches
        // For Gemma-4 vision: patch_size=16, hidden_size=768
        let patchSize = 16
        let imageWidth = 224
        let imageHeight = 224
        let numPatches = (imageWidth / patchSize) * (imageHeight / patchSize)  // 14*14 = 196
        let hiddenSize = 768
        
        Swift.print("  Image size: \(imageWidth)x\(imageHeight)")
        Swift.print("  Patch size: \(patchSize)x\(patchSize)")
        Swift.print("  Num patches: \(numPatches)")
        Swift.print("  Hidden size: \(hiddenSize)")
        
        // Create simple patch embeddings (normalized pixel values)
        // In a real implementation, this would use a learned projection
        // For testing, we'll use raw pixel patches flattened and normalized
        var patchEmbeddings = [Float](repeating: 0, count: numPatches * hiddenSize)
        
        // Simulate patch extraction: each patch has 16*16*3 = 768 values
        for patchIdx in 0..<numPatches {
            let patchRow = patchIdx / 14
            let patchCol = patchIdx % 14
            
            // Each patch: 16x16 pixels x 3 channels = 768 floats
            for y in 0..<patchSize {
                for x in 0..<patchSize {
                    for c in 0..<3 {
                        let idx = patchIdx * hiddenSize + (y * patchSize + x) * 3 + c
                        // Simple pattern: use patch position as feature
                        patchEmbeddings[idx] = Float(patchRow * patchSize + y + patchCol * patchSize + x) / 224.0
                    }
                }
            }
        }
        
        Swift.print("  ✓ Created \(numPatches) patch embeddings")
        
        // Load multimodal model
        let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        let engine = try MarkBaseEngine(autoCompile: true)
        let multimodalModel = try MultimodalModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
        let inference = try MultimodalInference(model: multimodalModel)
        let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
        
        Swift.print("\nVision tower: \(multimodalModel.visionTowerFull != nil ? "✓ loaded" : "✗ not loaded")")
        
        // Generate with vision input
        let prompt = "What do you see in this image?"
        let promptTokens = tokenizer.encode(text: prompt)
        
        Swift.print("\nPrompt: '\(prompt)'")
        Swift.print("Generating response...")
        
        // Process through vision tower
        let visionBuffer = engine.device.makeBuffer(bytes: patchEmbeddings, length: patchEmbeddings.count * 4)!
        let visionOutputBuffer = engine.device.makeBuffer(length: numPatches * hiddenSize * 4)!
        
        // Run VisionTower
        Swift.print("Running vision tower forward pass...")
        if let tower = multimodalModel.visionTowerFull {
            try tower.forward(patchEmbeddings: visionBuffer, numPatches: numPatches, outputBuffer: visionOutputBuffer)
            Swift.print("  ✓ Vision tower forward complete")
            
            // Pool embeddings: mean of all 196 patches → single embedding
            Swift.print("\nPooling vision embeddings...")
            let visionOutputPtr = visionOutputBuffer.contents().assumingMemoryBound(to: Float.self)
            var pooledEmbedding = [Float](repeating: 0, count: hiddenSize)
            
            // Mean pooling
            for i in 0..<hiddenSize {
                var sum: Float = 0
                for p in 0..<numPatches {
                    sum += visionOutputPtr[p * hiddenSize + i]
                }
                pooledEmbedding[i] = sum / Float(numPatches)
            }
            
            Swift.print("  ✓ Pooled 196 embeddings into 1")
            Swift.print("  Pooled embedding magnitude: \(sqrt(pooledEmbedding.reduce(0) { $0 + $1 * $1 }))")
            
            // Create buffer for pooled embedding
            let pooledBuffer = engine.device.makeBuffer(bytes: pooledEmbedding, length: pooledEmbedding.count * 4)!
            
            // Use pooled embedding with MultimodalInference
            Swift.print("\nGenerating with pooled vision embedding...")
            let generatedTokens = try inference.generate(
                textTokens: promptTokens,
                precomputedVisionEmbedding: pooledBuffer,
                maxTokens: 30
            )
            
            // Decode result
            let decoded = tokenizer.decode(tokens: generatedTokens)
            Swift.print("\nGenerated: '\(decoded)'")
            
            // Extract response (skip prompt + BOI + IMAGE + EOI = 4 tokens)
            let newTextStart = promptTokens.count + 4
            if generatedTokens.count > newTextStart {
                let newText = tokenizer.decode(tokens: Array(generatedTokens[newTextStart...]))
                Swift.print("\nResponse: '\(newText)'")
            }
        } else {
            Swift.print("ERROR: Vision tower not loaded")
        }
    }
    
    func testMultimodalVisionInference() throws {
        Swift.print("\n=== Multimodal Vision Inference Test ===")
        
        // Check if MultimodalModel is available
        let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        let engine = try MarkBaseEngine(autoCompile: true)
        
        Swift.print("\nLoading multimodal model...")
        let multimodalModel = try MultimodalModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
        let inference = try MultimodalInference(model: multimodalModel)
        let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
        
        Swift.print("\nModel components:")
        Swift.print("  ✓ Text model: \(multimodalModel.textModel.numHiddenLayers) layers")
        Swift.print("  ✓ Vision tower full: \(multimodalModel.visionTowerFull != nil ? "loaded (\(multimodalModel.visionTowerFull!.config.numHiddenLayers) layers)" : "not loaded")")
        Swift.print("  ✓ Vision tower 12B: \(multimodalModel.visionTower != nil ? "loaded" : "not loaded")")
        Swift.print("  ✓ Audio tower full: \(multimodalModel.audioTowerFull != nil ? "loaded" : "not loaded")")
        Swift.print("  ✓ Audio tower 12B: \(multimodalModel.audioTower != nil ? "loaded" : "not loaded")")
        Swift.print("  ✓ BOI token: \(multimodalModel.boiTokenId)")
        Swift.print("  ✓ IMAGE token: \(multimodalModel.imageTokenId)")
        Swift.print("  ✓ EOI token: \(multimodalModel.eoiTokenId)")
        
        // Create a simple test image (dummy patches)
        // In reality, this would come from actual image processing
        // For testing, we'll create synthetic patch embeddings
        Swift.print("\nCreating synthetic vision input...")
        
        let hiddenSize = multimodalModel.textModel.hiddenSize
        let numPatches = 1  // Single patch for testing
        
        // Create a simple synthetic embedding
        // This simulates what the vision tower would output
        var patchEmbedding = [Float](repeating: 0.01, count: hiddenSize)
        
        // Add some variation to make it more realistic
        for i in 0..<hiddenSize {
            patchEmbedding[i] = Float.random(in: -0.1...0.1)
        }
        
        // Test prompt: "Describe this image"
        let promptTokens = tokenizer.encode(text: "Describe this image")
        Swift.print("Prompt: 'Describe this image'")
        Swift.print("Tokens: \(promptTokens)")
        
        // Generate with vision input
        Swift.print("\nGenerating with vision conditioning...")
        
        let generatedTokens = try inference.generate(
            textTokens: promptTokens,
            imagePatches: patchEmbedding,
            numImagePatches: numPatches,
            maxTokens: 20
        )
        
        // Decode result
        let decoded = tokenizer.decode(tokens: generatedTokens)
        Swift.print("\nGenerated tokens: \(generatedTokens)")
        Swift.print("Decoded: '\(decoded)'")
        
        // Check if generation makes more sense than text-only
        let newTextTokens = Array(generatedTokens.dropFirst(promptTokens.count + 3))  // Skip prompt + BOI/IMAGE/EOI
        let newText = tokenizer.decode(tokens: newTextTokens)
        Swift.print("\nNew text (after multimodal tokens): '\(newText)'")
        
        Swift.print("\n=== Summary ===")
        Swift.print("Multimodal inference pipeline executed successfully")
        Swift.print("Vision conditioning was injected into text generation")
        Swift.print("Note: Synthetic embeddings were used - real test needs actual image")
    }
    
    func testRealVisionPipeline() throws {
        Swift.print("\n=== Real Vision Pipeline Test ===")
        
        let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        let engine = try MarkBaseEngine(autoCompile: true)
        let multimodalModel = try MultimodalModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
        let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
        
        Swift.print("\n✓ Multimodal model loaded")
        
        // Create test image (red 224x224)
        let imageData = """
iVBORw0KGgoAAAANSUhEUgAAAOAAAADgCAIAAACVT/22AAACaUlEQVR4nO3SMQEAIAzAsIF/z+CAlx6Jgh5dZ6Br/w6AF4OSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGJc2gpBmUNIOSZlDSDEqaQUkzKGkGZcouY7ECvxTQrjoAAAAASUVORK5CYII=
"""
        
        let data = Data(base64Encoded: imageData)!
        Swift.print("\n✓ Test image loaded (red 224x224): \(data.count) bytes")
        
        // Preprocess image (create patch embeddings)
        
        guard let ciImage = CIImage(data: data) else {
            Swift.print("✗ Failed to create CIImage")
            return
        }
        
        // Resize to 224x224
        let resizeFilter = CIFilter(name: "CILanczosScaleTransform")!
        resizeFilter.setValue(ciImage, forKey: kCIInputImageKey)
        let scale = 224.0 / max(ciImage.extent.width, ciImage.extent.height)
        resizeFilter.setValue(scale, forKey: kCIInputScaleKey)
        resizeFilter.setValue(1.0, forKey: kCIInputAspectRatioKey)
        
        guard let resized = resizeFilter.outputImage else {
            Swift.print("✗ Failed to resize image")
            return
        }
        
        // Convert to CGImage
        let context = CIContext()
        guard let cgImage = context.createCGImage(resized, from: resized.extent) else {
            Swift.print("✗ Failed to create CGImage")
            return
        }
        
        // Extract RGB pixels
        let dataProvider = cgImage.dataProvider!
        let pixelData = dataProvider.data!
        let ptr = CFDataGetBytePtr(pixelData)!
        
        Swift.print("✓ First pixel RGB: (\(ptr[0]), \(ptr[1]), \(ptr[2]))")
        
        // Create patch embeddings (16x16 patches)
        let patchSize = 16
        let numPatches = 14 * 14  // 196
        let hiddenSize = 768
        
        var patchEmbeddings = [Float](repeating: 0, count: numPatches * hiddenSize)
        
        for patchIdx in 0..<numPatches {
            let patchRow = patchIdx / 14
            let patchCol = patchIdx % 14
            
            for y in 0..<patchSize {
                for x in 0..<patchSize {
                    let globalY = patchRow * patchSize + y
                    let globalX = patchCol * patchSize + x
                    
                    if globalY >= 224 || globalX >= 224 { continue }
                    
                    let pixelIdx = globalY * 224 + globalX
                    let offset = pixelIdx * 4  // RGBA
                    
                    let r = Float(ptr[offset]) / 255.0
                    let g = Float(ptr[offset + 1]) / 255.0
                    let b = Float(ptr[offset + 2]) / 255.0
                    
                    let embedIdx = patchIdx * hiddenSize + (y * patchSize + x) * 3
                    if embedIdx + 2 < patchEmbeddings.count {
                        patchEmbeddings[embedIdx] = r
                        patchEmbeddings[embedIdx + 1] = g
                        patchEmbeddings[embedIdx + 2] = b
                    }
                }
            }
        }
        
        Swift.print("✓ Patch embeddings created: \(patchEmbeddings.count) floats")
        
        // Vision tower forward pass
        let visionBuffer = engine.device.makeBuffer(bytes: patchEmbeddings, length: patchEmbeddings.count * 4)!
        let visionOutputBuffer = engine.device.makeBuffer(length: numPatches * 2560 * 4)!
        
        guard let tower = multimodalModel.visionTowerFull else {
            Swift.print("✗ Vision tower not available")
            return
        }
        
        try tower.forward(patchEmbeddings: visionBuffer, numPatches: numPatches, outputBuffer: visionOutputBuffer)
        Swift.print("✓ Vision tower forward pass complete")
        
        // Pool embeddings (196 patches → 1)
        let visionPtr = visionOutputBuffer.contents().assumingMemoryBound(to: Float.self)
        var pooled = [Float](repeating: 0, count: 2560)
        
        for i in 0..<2560 {
            var sum: Float = 0
            for p in 0..<numPatches {
                sum += visionPtr[p * 2560 + i]
            }
            pooled[i] = sum / Float(numPatches)
        }
        
        let mag = sqrt(pooled.reduce(0) { $0 + $1 * $1 })
        Swift.print("✓ Pooled vision embedding magnitude: \(mag)")
        
        // Normalize to match text embeddings (magnitude ~5)
        let scaleNorm: Float = 5.0 / mag
        for i in 0..<2560 {
            pooled[i] *= scaleNorm
        }
        
        let magAfterNorm = sqrt(pooled.reduce(0) { $0 + $1 * $1 })
        Swift.print("✓ Normalized magnitude: \(magAfterNorm)")
        
        // Multimodal inference
        let inference = try MultimodalInference(model: multimodalModel)
        let pooledBuffer = engine.device.makeBuffer(bytes: pooled, length: pooled.count * 4)!
        
        let promptTokens = tokenizer.encode(text: "What color is this image?")
        Swift.print("\nPrompt: 'What color is this image?'")
        Swift.print("Tokens: \(promptTokens)")
        
        let generatedTokens = try inference.generate(
            textTokens: promptTokens,
            precomputedVisionEmbedding: pooledBuffer,
            maxTokens: 20
        )
        
        let decoded = tokenizer.decode(tokens: generatedTokens)
        Swift.print("\nGenerated: '\(decoded)'")
        
        // Extract response (skip prompt + BOI + IMAGE + EOI)
        let responseStart = promptTokens.count + 4
        if generatedTokens.count > responseStart {
            let responseTokens = Array(generatedTokens[responseStart...])
            let response = tokenizer.decode(tokens: responseTokens)
            Swift.print("Response: '\(response)'")
        }
        
        Swift.print("\n=== Real Vision Pipeline Test Complete ===")
    }
    
    func test26BModelLoading() throws {
        Swift.print("\n=== 测试 Gemma-4 26B A4B 真正 4-bit 模型 ===")
        
        // 真正的 4-bit A4B model (不是 MXFP4)
        let modelDir = "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-a4b-it-4bit"
        
        // Check if model exists
        guard FileManager.default.fileExists(atPath: modelDir + "/config.json") else {
            Swift.print("✗ 转换后的26B模型未找到")
            Swift.print("  请运行: python3 convert_mlx_26b.py")
            return
        }
        
        Swift.print("✓ 转换后的26B模型找到")
        
        // Check files
        let configFile = modelDir + "/config.json"
        let tokenizerFile = modelDir + "/tokenizer.json"
        let weightsFile1 = modelDir + "/model-00001-of-00003.safetensors"
        
        for (file, desc) in [(configFile, "Config"), (tokenizerFile, "Tokenizer"), (weightsFile1, "Weights shard 1")] {
            if FileManager.default.fileExists(atPath: file) {
                let attrs = try FileManager.default.attributesOfItem(atPath: file)
                let size = (attrs[.size] as? Int64 ?? 0) / 1024 / 1024
                Swift.print("  ✓ \(desc): \(size) MB")
            } else {
                Swift.print("  ✗ \(desc): NOT FOUND")
                return
            }
        }
        
        // Check total weights size
        let totalWeights: Int64 = try FileManager.default.contentsOfDirectory(atPath: modelDir)
            .filter { $0.hasSuffix(".safetensors") }
            .reduce(0) { sum, file in
                let attrs = try FileManager.default.attributesOfItem(atPath: modelDir + "/" + file)
                return sum + (attrs[.size] as? Int64 ?? 0)
            }
        Swift.print("  ✓ Total weights: \(totalWeights / 1024 / 1024 / 1024) GB")
        
        Swift.print("\n步骤 2: 创建 Engine")
        let engine = try MarkBaseEngine(autoCompile: true)
        Swift.print("  ✓ Engine created")
        
        Swift.print("\n步骤 3: 加载 Model")
        Swift.print("  26B 转换后模型 (~15GB):")
        Swift.print("    - Hidden size: 2816")
        Swift.print("    - Layers: 42")
        Swift.print("    - Vocab: 262144")
        Swift.print("  警告: 可能需要1-2分钟...")
        
        do {
            // Use smaller context for testing
            let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 128)
            
            Swift.print("  ✓✓✓ Model loaded!")
            Swift.print("  Layers: \(model.numHiddenLayers)")
            Swift.print("  Hidden size: \(model.hiddenSize)")
            Swift.print("  Vocab size: \(model.vocabSize)")
            
            // Test tokenizer
            let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
            Swift.print("  ✓ Tokenizer loaded")
            
            // Test simple encoding
            let tokens = tokenizer.encode(text: "Hello")
            Swift.print("  Test tokens: \(tokens)")
            
            Swift.print("\n=== 26B 测试成功 ===")
            Swift.print("✓ 转换后的26B模型可以加载！")
            Swift.print("✓ Hidden size: \(model.hiddenSize) (比12B的2560大)")
            Swift.print("✓ 准备运行推理...")
            
            // Try simple forward pass (optional)
            Swift.print("\n尝试推理测试...")
            do {
                let logits = try model.forward(tokenId: tokens[0], position: 0)
                Swift.print("  ✓ Forward pass成功!")
                Swift.print("  Logits size: \(logits.count)")
                
                // Check top predictions
                let sorted = logits.enumerated().sorted { $0.element > $1.element }
                let top5 = sorted.prefix(5)
                Swift.print("  Top 5 predictions:")
                for (idx, val) in top5 {
                    let tokenStr = tokenizer.decode(tokens: [idx])
                    Swift.print("    Token \(idx) ('\(tokenStr)'): \(val)")
                }
                
                Swift.print("\n🎉 26B模型完全工作！")
                
                // ── Quick generation test (3 tokens) ──
                Swift.print("\n尝试简短的生成测试...")
                do {
                    let genConfig = GenerationConfig(maxTokens: 3, temperature: 0.7)
                    let generator = StreamingGenerator(model: model, tokenizer: tokenizer, engine: engine)
                    let response = try generator.generateComplete(prompt: "Hello", config: genConfig)
                    Swift.print("  Generated: '\(response)'")
                    XCTAssertGreaterThan(response.count, 0, "Should generate text")
                } catch {
                    Swift.print("  ⚠️ Generation test失败: \(error)")
                }
                
            } catch {
                Swift.print("  ⚠️ Forward pass失败: \(error)")
                Swift.print("  可能需要MoE支持或额外适配")
            }
            
        } catch {
            Swift.print("\n✗ 加载失败: \(error)")
            
            let nsError = error as NSError
            Swift.print("\n错误详情:")
            Swift.print("  Domain: \(nsError.domain)")
            Swift.print("  Code: \(nsError.code)")
            Swift.print("  Description: \(nsError.localizedDescription)")
            
            if let userInfo = nsError.userInfo as? [String: String] {
                for (key, value) in userInfo {
                    Swift.print("  \(key): \(value)")
                }
            }
            
            Swift.print("\n可能原因:")
            Swift.print("  1. 权重命名仍不完全匹配")
            Swift.print("  2. MoE结构需要实现")
            Swift.print("  3. scales转换有问题")
            Swift.print("  4. Memory不足 (需要~17GB)")
            
            Swift.print("\n建议:")
            Swift.print("  - 检查权重命名是否正确")
            Swift.print("  - 检查Memory是否充足")
            Swift.print("  - 查看详细错误日志")
        }
    }
    
    func testNaturalImageInference() throws {
        Swift.print("\n=== Natural Image Inference Test ===")
        Swift.print("Testing with sky gradient + sun (more realistic)")
        
        let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        let engine = try MarkBaseEngine(autoCompile: true)
        let multimodalModel = try MultimodalModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
        let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
        
        Swift.print("✓ Model loaded")
        
        guard let imageData = try? String(contentsOfFile: "/tmp/test_natural_image.txt") else {
            Swift.print("✗ Natural image not found")
            return
        }
        
        let data = Data(base64Encoded: imageData)!
        Swift.print("✓ Natural image: \(data.count) bytes (sky + sun)")
        
        let outputFile = "/tmp/natural_inference_output.txt"
        var outputs: [String] = ["=== Natural Image Inference Test ==="]
        
        // Preprocess
        guard let ciImage = CIImage(data: data) else {
            outputs.append("✗ Failed to create CIImage")
            try outputs.joined(separator: "\n").write(toFile: outputFile, atomically: true, encoding: .utf8)
            return
        }
        
        outputs.append("✓ Image size: \(Int(ciImage.extent.width))x\(Int(ciImage.extent.height))")
        
        let resizeFilter = CIFilter(name: "CILanczosScaleTransform")!
        resizeFilter.setValue(ciImage, forKey: kCIInputImageKey)
        let scale = 224.0 / max(ciImage.extent.width, ciImage.extent.height)
        resizeFilter.setValue(scale, forKey: kCIInputScaleKey)
        resizeFilter.setValue(1.0, forKey: kCIInputAspectRatioKey)
        
        guard let resized = resizeFilter.outputImage else { return }
        
        let context = CIContext()
        guard let cgImage = context.createCGImage(resized, from: resized.extent) else { return }
        
        let ptr = CFDataGetBytePtr(cgImage.dataProvider!.data!)!
        
        outputs.append("✓ Top-left pixel: (\(ptr[0]), \(ptr[1]), \(ptr[2]))")
        outputs.append("✓ Sun center (~40,40): (\(ptr[(40*224+40)*4]), \(ptr[(40*224+40)*4+1]), \(ptr[(40*224+40)*4+2]))")
        
        // Vision processing
        let numPatches = 196
        let hiddenSize = 768
        var patchEmbeddings = [Float](repeating: 0, count: numPatches * hiddenSize)
        
        for patchIdx in 0..<numPatches {
            let row = patchIdx / 14
            let col = patchIdx % 14
            
            for y in 0..<16 {
                for x in 0..<16 {
                    let gy = row * 16 + y
                    let gx = col * 16 + x
                    
                    if gy >= 224 || gx >= 224 { continue }
                    
                    let offset = (gy * 224 + gx) * 4
                    let r = Float(ptr[offset]) / 255.0
                    let g = Float(ptr[offset + 1]) / 255.0
                    let b = Float(ptr[offset + 2]) / 255.0
                    
                    let idx = patchIdx * hiddenSize + (y * 16 + x) * 3
                    if idx + 2 < patchEmbeddings.count {
                        patchEmbeddings[idx] = r
                        patchEmbeddings[idx + 1] = g
                        patchEmbeddings[idx + 2] = b
                    }
                }
            }
        }
        
        outputs.append("✓ Patch embeddings: \(patchEmbeddings.count) floats")
        
        // Vision tower
        let visionBuffer = engine.device.makeBuffer(bytes: patchEmbeddings, length: patchEmbeddings.count * 4)!
        let visionOutputBuffer = engine.device.makeBuffer(length: numPatches * 2560 * 4)!
        try multimodalModel.visionTowerFull!.forward(patchEmbeddings: visionBuffer, numPatches: numPatches, outputBuffer: visionOutputBuffer)
        
        outputs.append("✓ Vision tower forward complete")
        
        // Pool
        let visionPtr = visionOutputBuffer.contents().assumingMemoryBound(to: Float.self)
        var pooled = [Float](repeating: 0, count: 2560)
        
        for i in 0..<2560 {
            var sum: Float = 0
            for p in 0..<numPatches {
                sum += visionPtr[p * 2560 + i]
            }
            pooled[i] = sum / Float(numPatches)
        }
        
        let mag = sqrt(pooled.reduce(0) { $0 + $1 * $1 })
        outputs.append("✓ Pooled magnitude: \(mag)")
        
        let normScale: Float = 5.0 / mag
        for i in 0..<2560 {
            pooled[i] *= normScale
        }
        
        outputs.append("✓ Normalized magnitude: \(sqrt(pooled.reduce(0) { $0 + $1 * $1 }))")
        
        // Test multiple prompts
        let inference = try MultimodalInference(model: multimodalModel)
        let pooledBuffer = engine.device.makeBuffer(bytes: pooled, length: pooled.count * 4)!
        
        // Prompt 1: Describe scene
        let p1 = tokenizer.encode(text: "Describe what you see in this image")
        outputs.append("\n[Test 1] Prompt: 'Describe what you see in this image'")
        let g1 = try inference.generate(textTokens: p1, precomputedVisionEmbedding: pooledBuffer, maxTokens: 40)
        outputs.append("Generated: '\(tokenizer.decode(tokens: g1))'")
        
        // Prompt 2: Colors
        let p2 = tokenizer.encode(text: "What colors are present?")
        outputs.append("\n[Test 2] Prompt: 'What colors are present?'")
        let g2 = try inference.generate(textTokens: p2, precomputedVisionEmbedding: pooledBuffer, maxTokens: 40)
        outputs.append("Generated: '\(tokenizer.decode(tokens: g2))'")
        
        // Prompt 3: Scene type
        let p3 = tokenizer.encode(text: "Is this outdoor or indoor?")
        outputs.append("\n[Test 3] Prompt: 'Is this outdoor or indoor?'")
        let g3 = try inference.generate(textTokens: p3, precomputedVisionEmbedding: pooledBuffer, maxTokens: 40)
        outputs.append("Generated: '\(tokenizer.decode(tokens: g3))'")
        
        outputs.append("\n=== Natural Image Test Complete ===")
        outputs.append("Note: Compare outputs with model's expected behavior")
        
        try outputs.joined(separator: "\n").write(toFile: outputFile, atomically: true, encoding: .utf8)
        Swift.print("✓ Results: \(outputFile)")
        
        XCTAssertTrue(FileManager.default.fileExists(atPath: outputFile))
    }
    
    func testGradientImageInference() throws {
        Swift.print("\n=== Gradient Image Inference Test ===")
        Swift.print("Testing with more complex gradient pattern (not pure red)")
        
        let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        let engine = try MarkBaseEngine(autoCompile: true)
        let multimodalModel = try MultimodalModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
        let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
        
        Swift.print("✓ Multimodal model loaded")
        
        // Create gradient image (224x224)
        guard let imageData = try? String(contentsOfFile: "/tmp/test_gradient_image.txt") else {
            Swift.print("✗ Gradient image file not found, skipping test")
            return
        }
        
        let data = Data(base64Encoded: imageData)!
        Swift.print("✓ Gradient image loaded: \(data.count) bytes")
        
        // Save outputs to file for verification
        let outputFile = "/tmp/gradient_inference_output.txt"
        var outputs: [String] = []
        outputs.append("=== Gradient Image Inference Test ===")
        
        // Preprocess
        guard let ciImage = CIImage(data: data) else {
            outputs.append("✗ Failed to create CIImage")
            try outputs.joined(separator: "\n").write(toFile: outputFile, atomically: true, encoding: .utf8)
            return
        }
        
        let resizeFilter = CIFilter(name: "CILanczosScaleTransform")!
        resizeFilter.setValue(ciImage, forKey: kCIInputImageKey)
        let scale = 224.0 / max(ciImage.extent.width, ciImage.extent.height)
        resizeFilter.setValue(scale, forKey: kCIInputScaleKey)
        resizeFilter.setValue(1.0, forKey: kCIInputAspectRatioKey)
        
        guard let resized = resizeFilter.outputImage else {
            outputs.append("✗ Failed to resize")
            try outputs.joined(separator: "\n").write(toFile: outputFile, atomically: true, encoding: .utf8)
            return
        }
        
        let context = CIContext()
        guard let cgImage = context.createCGImage(resized, from: resized.extent) else {
            outputs.append("✗ Failed to create CGImage")
            try outputs.joined(separator: "\n").write(toFile: outputFile, atomically: true, encoding: .utf8)
            return
        }
        
        let dataProvider = cgImage.dataProvider!
        let pixelData = dataProvider.data!
        let ptr = CFDataGetBytePtr(pixelData)!
        
        outputs.append("✓ First pixel RGB: (\(ptr[0]), \(ptr[1]), \(ptr[2]))")
        
        // Create patch embeddings
        let patchSize = 16
        let numPatches = 14 * 14
        let hiddenSize = 768
        var patchEmbeddings = [Float](repeating: 0, count: numPatches * hiddenSize)
        
        for patchIdx in 0..<numPatches {
            let patchRow = patchIdx / 14
            let patchCol = patchIdx % 14
            
            for y in 0..<patchSize {
                for x in 0..<patchSize {
                    let globalY = patchRow * patchSize + y
                    let globalX = patchCol * patchSize + x
                    
                    if globalY >= 224 || globalX >= 224 { continue }
                    
                    let pixelIdx = globalY * 224 + globalX
                    let offset = pixelIdx * 4
                    
                    let r = Float(ptr[offset]) / 255.0
                    let g = Float(ptr[offset + 1]) / 255.0
                    let b = Float(ptr[offset + 2]) / 255.0
                    
                    let embedIdx = patchIdx * hiddenSize + (y * patchSize + x) * 3
                    if embedIdx + 2 < patchEmbeddings.count {
                        patchEmbeddings[embedIdx] = r
                        patchEmbeddings[embedIdx + 1] = g
                        patchEmbeddings[embedIdx + 2] = b
                    }
                }
            }
        }
        
        outputs.append("✓ Patch embeddings created: \(patchEmbeddings.count) floats")
        
        // Vision tower forward
        let visionBuffer = engine.device.makeBuffer(bytes: patchEmbeddings, length: patchEmbeddings.count * 4)!
        let visionOutputBuffer = engine.device.makeBuffer(length: numPatches * 2560 * 4)!
        try multimodalModel.visionTowerFull!.forward(patchEmbeddings: visionBuffer, numPatches: numPatches, outputBuffer: visionOutputBuffer)
        
        outputs.append("✓ Vision tower forward complete")
        
        // Pool and normalize
        let visionPtr = visionOutputBuffer.contents().assumingMemoryBound(to: Float.self)
        var pooled = [Float](repeating: 0, count: 2560)
        
        for i in 0..<2560 {
            var sum: Float = 0
            for p in 0..<numPatches {
                sum += visionPtr[p * 2560 + i]
            }
            pooled[i] = sum / Float(numPatches)
        }
        
        let mag = sqrt(pooled.reduce(0) { $0 + $1 * $1 })
        outputs.append("✓ Pooled magnitude: \(mag)")
        
        let scaleNorm: Float = 5.0 / mag
        for i in 0..<2560 {
            pooled[i] *= scaleNorm
        }
        
        let magAfter = sqrt(pooled.reduce(0) { $0 + $1 * $1 })
        outputs.append("✓ Normalized magnitude: \(magAfter)")
        
        // Multimodal inference with different prompts
        let inference = try MultimodalInference(model: multimodalModel)
        let pooledBuffer = engine.device.makeBuffer(bytes: pooled, length: pooled.count * 4)!
        
        // Test 1: Simple question
        let prompt1 = tokenizer.encode(text: "What do you see?")
        outputs.append("\n[Test 1] Prompt: 'What do you see?'")
        let gen1 = try inference.generate(textTokens: prompt1, precomputedVisionEmbedding: pooledBuffer, maxTokens: 30)
        let dec1 = tokenizer.decode(tokens: gen1)
        outputs.append("Generated: '\(dec1)'")
        
        // Extract response
        if gen1.count > prompt1.count + 4 {
            let respTokens = Array(gen1[(prompt1.count + 4)...])
            let resp = tokenizer.decode(tokens: respTokens)
            outputs.append("Response: '\(resp)'")
        }
        
        // Test 2: Describe request
        let prompt2 = tokenizer.encode(text: "Describe this image")
        outputs.append("\n[Test 2] Prompt: 'Describe this image'")
        let gen2 = try inference.generate(textTokens: prompt2, precomputedVisionEmbedding: pooledBuffer, maxTokens: 30)
        let dec2 = tokenizer.decode(tokens: gen2)
        outputs.append("Generated: '\(dec2)'")
        
        if gen2.count > prompt2.count + 4 {
            let respTokens = Array(gen2[(prompt2.count + 4)...])
            let resp = tokenizer.decode(tokens: respTokens)
            outputs.append("Response: '\(resp)'")
        }
        
        // Test 3: Color question
        let prompt3 = tokenizer.encode(text: "What colors are in this image?")
        outputs.append("\n[Test 3] Prompt: 'What colors are in this image?'")
        let gen3 = try inference.generate(textTokens: prompt3, precomputedVisionEmbedding: pooledBuffer, maxTokens: 30)
        let dec3 = tokenizer.decode(tokens: gen3)
        outputs.append("Generated: '\(dec3)'")
        
        if gen3.count > prompt3.count + 4 {
            let respTokens = Array(gen3[(prompt3.count + 4)...])
            let resp = tokenizer.decode(tokens: respTokens)
            outputs.append("Response: '\(resp)'")
        }
        
        outputs.append("\n=== Gradient Image Test Complete ===")
        outputs.append("Note: Output quality should be analyzed for coherence")
        
        // Write outputs to file
        try outputs.joined(separator: "\n").write(toFile: outputFile, atomically: true, encoding: .utf8)
        Swift.print("✓ Results saved to \(outputFile)")
        
        // Verify file was created
        XCTAssertTrue(FileManager.default.fileExists(atPath: outputFile), "Output file should be created")
    }
    
    func testMultimodalTextOnly() throws {
        Swift.print("\n=== Multimodal Model Text-Only Test ===")
        Swift.print("Note: E4B-MarkBase is Gemma4ForConditionalGeneration (multimodal)")
        Swift.print("Text-only generation may not work properly without vision/audio conditioning")
        
        let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        let engine = try MarkBaseEngine(autoCompile: true)
        
        // Check if multimodal inference is available
        Swift.print("\nChecking multimodal components:")
        let mmInferenceAvailable = NSClassFromString("G12B.MultimodalInference") != nil
        Swift.print("  MultimodalInference class: \(mmInferenceAvailable ? "available" : "NOT available")")
        
        // Check if vision/audio weights are loaded
        let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
        let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
        let sampler = Sampler()
        
        Swift.print("\nModel components loaded:")
        Swift.print("  ✓ 42 text layers")
        Swift.print("  Vision/Audio towers should be separate components")
        
        // The issue is: pure text generation without multimodal conditioning
        // The model expects vision/audio features to be interleaved with text
        Swift.print("\n=== Testing with multimodal placeholder tokens ===")
        
        // Try using image_token or audio_token as conditioning
        for cache in model.kvCaches {
            cache.reset()
        }
        
        // Test 1: Add image token before text (258880 = image_token)
        Swift.print("\nTest 1: Prompt with image token")
        let imageToken = 258880
        var tokens = [2, imageToken]  // BOS + IMAGE
        
        Swift.print("Processing IMAGE token at position 0...")
        // Note: This won't work properly because we need vision features
        // But let's see what happens
        
        // Generate a few tokens
        for i in 0..<5 {
            let pos = tokens.count - 1
            let logits = try model.forward(tokenId: tokens.last!, position: pos)
            let nextToken = sampler.sample(logits: logits, temperature: 0.1, topK: 50, topP: 0.95, filterUnusedTokens: true)
            tokens.append(nextToken)
            let tokenStr = tokenizer.decode(tokens: [nextToken])
            Swift.print("  Position \(i): \(nextToken) ('\(tokenStr)')")
        }
        
        // Test 2: Check special token IDs
        Swift.print("\n\nSpecial multimodal token IDs:")
        let specialTokens = [
            (258880, "image_token"),
            (258881, "audio_token"),
            (258882, "end_of_image"),
            (258883, "end_of_audio"),
            (255999, "boi"),
            (256000, "boa"),
        ]
        
        for (id, name) in specialTokens {
            let decoded = tokenizer.decode(tokens: [id])
            Swift.print("  Token \(id) (\(name)): '\(decoded)'")
        }
        
        // Test 3: What happens with end_of_turn token?
        Swift.print("\n\nTest 3: Check if model responds to format tokens")
        for cache in model.kvCaches {
            cache.reset()
        }
        
        // <|turn> token is 105
        let turnStart = 105
        let turnEnd = 106
        
        // BOS + <|turn> + text + <turn|>
        let prompt = "Describe this image"
        var promptTokens = tokenizer.encode(text: prompt)
        Swift.print("Prompt tokens: \(promptTokens)")
        
        // We need to insert <|turn> and <turn|> correctly
        // Actually, the tokenizer should handle these if they're in vocabulary
        Swift.print("Note: Proper multimodal inference requires vision/audio features")
        Swift.print("      Text-only generation is not the intended use case")
        
        Swift.print("\n=== Summary ===")
        Swift.print("E4B-MarkBase is a multimodal model (Gemma4ForConditionalGeneration)")
        Swift.print("It requires:")
        Swift.print("  1. Vision input (image) processed through vision_tower")
        Swift.print("  2. Audio input (audio) processed through audio_tower")
        Swift.print("  3. Multimodal embedding interleaved with text tokens")
        Swift.print("\nFor text-only generation, use a pure text model like Gemma-4-4B-IT")
    }
    
    func testContinuationMode() throws {
        Swift.print("\n=== Continuation Mode Test ===")
        let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
        let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
        let sampler = Sampler()
        
        // Test 1: Simple text continuation (no special format tokens)
        for cache in model.kvCaches {
            cache.reset()
        }
        
        // "The capital of France is" - expecting "Paris"
        let prompt1 = "The capital of France is"
        var tokens1 = tokenizer.encode(text: prompt1)
        Swift.print("\nPrompt: '\(prompt1)'")
        Swift.print("Tokens: \(tokens1)")
        Swift.print("Decoded: '\(tokenizer.decode(tokens: tokens1))'")
        
        // Process prompt
        for (pos, token) in tokens1.enumerated() {
            let logits = try model.forward(tokenId: token, position: pos)
            if pos == tokens1.count - 1 {
                // Show top 10 predictions at last prompt position
                let top10 = logits.enumerated().sorted { $0.element > $1.element }.prefix(10)
                Swift.print("\nTop 10 predictions at prompt end:")
                for (idx, val) in top10 {
                    let tokenStr = tokenizer.decode(tokens: [idx])
                    Swift.print("  Token \(idx) ('\(tokenStr)'): \(val)")
                }
            }
        }
        
        // Generate continuation
        Swift.print("\nGenerating continuation:")
        for i in 0..<10 {
            let pos = tokens1.count - 1
            let logits = try model.forward(tokenId: tokens1.last!, position: pos)
            let nextToken = sampler.sample(logits: logits, temperature: 0.1, topK: 50, topP: 0.95, filterUnusedTokens: true)
            tokens1.append(nextToken)
            let tokenStr = tokenizer.decode(tokens: [nextToken])
            Swift.print("  Token \(i + 1): \(nextToken) ('\(tokenStr)')")
        }
        
        Swift.print("\nContinuation: '\(tokenizer.decode(tokens: tokens1))'")
        
        // Test 2: Another continuation
        for cache in model.kvCaches {
            cache.reset()
        }
        
        let prompt2 = "Hello, my name is"
        var tokens2 = tokenizer.encode(text: prompt2)
        Swift.print("\n\nPrompt 2: '\(prompt2)'")
        
        // Process and generate
        for (pos, token) in tokens2.enumerated() {
            let logits = try model.forward(tokenId: token, position: pos)
            if pos == tokens2.count - 1 {
                let top5 = logits.enumerated().sorted { $0.element > $1.element }.prefix(5)
                Swift.print("Top 5 at prompt end:")
                for (idx, val) in top5 {
                    let tokenStr = tokenizer.decode(tokens: [idx])
                    Swift.print("  Token \(idx): \(val) ('\(tokenStr)')")
                }
            }
        }
        
        Swift.print("\nGenerating:")
        for i in 0..<10 {
            let pos = tokens2.count - 1
            let logits = try model.forward(tokenId: tokens2.last!, position: pos)
            let nextToken = sampler.sample(logits: logits, temperature: 0.1, topK: 50, topP: 0.95, filterUnusedTokens: true)
            tokens2.append(nextToken)
            let tokenStr = tokenizer.decode(tokens: [nextToken])
            Swift.print("  \(nextToken) ('\(tokenStr)')")
        }
        
        Swift.print("\nResult: '\(tokenizer.decode(tokens: tokens2))'")
        
        // Test 3: Check if "Paris" appears in predictions for "The capital of France is"
        for cache in model.kvCaches {
            cache.reset()
        }
        
        Swift.print("\n\n=== Checking 'Paris' prediction ===")
        let prompt3 = "The capital of France is"
        let tokens3 = tokenizer.encode(text: prompt3)
        
        // Process prompt
        for (pos, token) in tokens3.enumerated() {
            let logits = try model.forward(tokenId: token, position: pos)
            if pos == tokens3.count - 1 {
                // Find "Paris" token
                // Try different variations
                let parisVariants = [
                    tokenizer.encode(text: "Paris"),
                    tokenizer.encode(text: " paris"),
                    tokenizer.encode(text: "▁Paris"),
                ]
                Swift.print("Paris token variants: \(parisVariants)")
                
                // Check if any variant is in top 50
                let top50 = logits.enumerated().sorted { $0.element > $1.element }.prefix(50)
                for (idx, val) in top50 {
                    let tokenStr = tokenizer.decode(tokens: [idx])
                    if tokenStr.lowercased().contains("paris") {
                        Swift.print("Found 'Paris' variant at position \(idx) in top 50: rank \(top50.firstIndex(where: { $0.offset == idx }) ?? -1), logit=\(val)")
                    }
                }
            }
        }
    }
    
    func testKernelSelection() throws {
        Swift.print("\n=== Kernel Selection Test ===")
        let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        let engine = try MarkBaseEngine(autoCompile: true)
        
        // Check if SIMD kernel is available
        let simdAvailable = (try? engine.pipeline(named: "full_attention_simd")) != nil
        Swift.print("full_attention_simd available: \(simdAvailable)")
        
        let regularAvailable = (try? engine.pipeline(named: "full_attention")) != nil
        Swift.print("full_attention available: \(regularAvailable)")
        
        let simdSlidingAvailable = (try? engine.pipeline(named: "sliding_attention_simd")) != nil
        Swift.print("sliding_attention_simd available: \(simdSlidingAvailable)")
        
        let slidingAvailable = (try? engine.pipeline(named: "sliding_attention")) != nil
        Swift.print("sliding_attention available: \(slidingAvailable)")
        
        // Test with model
        let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
        let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
        let sampler = Sampler()
        
        Swift.print("\nModel layer types:")
        for i in [0, 5, 6, 11, 17, 23, 29, 35, 41] {
            let layer = model.layers[i]
            Swift.print("  Layer \(i): \(layer.config.isSliding ? "sliding" : "full") attention")
        }
        
        // Generate with temperature=0.1
        for cache in model.kvCaches {
            cache.reset()
        }
        
        Swift.print("\nGenerating with temperature=0.1:")
        var tokens = [2]
        for i in 0..<20 {
            let pos = tokens.count - 1
            let logits = try model.forward(tokenId: tokens.last!, position: pos)
            let nextToken = sampler.sample(logits: logits, temperature: 0.1, topK: 50, topP: 0.95, filterUnusedTokens: true)
            tokens.append(nextToken)
            let tokenStr = tokenizer.decode(tokens: [nextToken])
            Swift.print("Position \(i + 1): \(nextToken) ('\(tokenStr)')")
        }
        
        Swift.print("\nGenerated: '\(tokenizer.decode(tokens: tokens))'")
    }
    
    func testSimpleBOSGeneration() throws {
        Swift.print("\n=== Simple BOS Generation Test ===")
        let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
        let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
        let sampler = Sampler()
        
        for cache in model.kvCaches {
            cache.reset()
        }
        
        // Generate from BOS only
        Swift.print("\nGenerating from BOS token only (position 0):")
        let logits0 = try model.forward(tokenId: 2, position: 0)
        
        let top10 = logits0.enumerated().sorted { $0.element > $1.element }.prefix(10)
        Swift.print("Top 10 logits:")
        for (idx, val) in top10 {
            let tokenStr = tokenizer.decode(tokens: [idx])
            Swift.print("  Token \(idx) ('\(tokenStr)'): \(val)")
        }
        
        // Generate 20 tokens from BOS
        var tokens = [2]
        for i in 0..<20 {
            let pos = tokens.count - 1
            let logits = try model.forward(tokenId: tokens.last!, position: pos)
            let nextToken = sampler.sample(logits: logits, temperature: 0.7, topK: 50, topP: 0.95, filterUnusedTokens: true)
            tokens.append(nextToken)
            let tokenStr = tokenizer.decode(tokens: [nextToken])
            Swift.print("Position \(i + 1): \(nextToken) ('\(tokenStr)')")
        }
        
        Swift.print("\nGenerated: '\(tokenizer.decode(tokens: tokens))'")
        
        // Test with very low temperature (near greedy)
        for cache in model.kvCaches {
            cache.reset()
        }
        
        Swift.print("\n\nGenerating with temperature=0.1 (near-greedy):")
        var tokens2 = [2]
        for i in 0..<20 {
            let pos = tokens2.count - 1
            let logits = try model.forward(tokenId: tokens2.last!, position: pos)
            let nextToken = sampler.sample(logits: logits, temperature: 0.1, topK: 50, topP: 0.95, filterUnusedTokens: true)
            tokens2.append(nextToken)
            let tokenStr = tokenizer.decode(tokens: [nextToken])
            Swift.print("Position \(i + 1): \(nextToken) ('\(tokenStr)')")
        }
        
        Swift.print("\nGenerated: '\(tokenizer.decode(tokens: tokens2))'")
        
        // Test with top-K=1 (pure greedy)
        for cache in model.kvCaches {
            cache.reset()
        }
        
        Swift.print("\n\nGenerating with top-K=1 (greedy):")
        var tokens3 = [2]
        for i in 0..<20 {
            let pos = tokens3.count - 1
            let logits = try model.forward(tokenId: tokens3.last!, position: pos)
            let nextToken = sampler.sample(logits: logits, temperature: 1.0, topK: 1, filterUnusedTokens: true)
            tokens3.append(nextToken)
            let tokenStr = tokenizer.decode(tokens: [nextToken])
            Swift.print("Position \(i + 1): \(nextToken) ('\(tokenStr)')")
        }
        
        Swift.print("\nGenerated: '\(tokenizer.decode(tokens: tokens3))'")
    }
    
    func testTokenizerEncoding() throws {
        Swift.print("\n=== Tokenizer Encoding Test ===")
        let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
        
        // Test 1: Simple prompt without special tokens
        let prompt1 = "What is the capital of France?"
        let tokens1 = tokenizer.encode(text: prompt1)
        let decoded1 = tokenizer.decode(tokens: tokens1)
        Swift.print("\nTest 1: Simple prompt")
        Swift.print("  Prompt: '\(prompt1)'")
        Swift.print("  Tokens: \(tokens1)")
        Swift.print("  Decoded: '\(decoded1)'")
        
        // Check each token
        Swift.print("  Token breakdown:")
        for (i, token) in tokens1.enumerated() {
            let decodedToken = tokenizer.decode(tokens: [token])
            Swift.print("    \(i): Token \(token) -> '\(decodedToken)'")
        }
        
        // Test 2: Prompt with special format tokens
        let prompt2 = "<bos><|turn>What is the capital of France?<turn|><|turn>"
        let tokens2 = tokenizer.encode(text: prompt2)
        let decoded2 = tokenizer.decode(tokens: tokens2)
        Swift.print("\nTest 2: Prompt with format tokens")
        Swift.print("  Prompt: '\(prompt2)'")
        Swift.print("  Tokens: \(tokens2)")
        Swift.print("  Decoded: '\(decoded2)'")
        
        // Check first 20 tokens
        Swift.print("  Token breakdown (first 20):")
        for (i, token) in tokens2.enumerated().prefix(20) {
            let decodedToken = tokenizer.decode(tokens: [token])
            Swift.print("    \(i): Token \(token) -> '\(decodedToken)'")
        }
        
        // Test 3: Check if spaces are preserved
        let prompt3 = "Hello World"
        let tokens3 = tokenizer.encode(text: prompt3)
        let decoded3 = tokenizer.decode(tokens: tokens3)
        Swift.print("\nTest 3: Space preservation")
        Swift.print("  Prompt: '\(prompt3)'")
        Swift.print("  Tokens: \(tokens3)")
        Swift.print("  Decoded: '\(decoded3)'")
        Swift.print("  Are spaces preserved? \(decoded3.contains(" ") ? "YES" : "NO")")
        
        // Test 4: Check underscore prefix tokens
        // In Gemma, words usually start with ▁ (underscore) prefix
        Swift.print("\nTest 4: Check for ▁ prefix in vocabulary")
        // Manually check some vocabulary entries
        // Token IDs for common words should have ▁ prefix
        let commonTokens = [506, 532, 496, 107]  // the, and, a, newline
        for token in commonTokens {
            let decodedToken = tokenizer.decode(tokens: [token])
            Swift.print("  Token \(token): '\(decodedToken)' - has prefix? \(decodedToken.hasPrefix("▁") || decodedToken.hasPrefix(" ") ? "YES" : "NO")")
        }
    }
    
    func testSamplingStrategies() throws {
        Swift.print("\n=== Sampling Strategies Test ===")
        let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
        let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
        
        // Process token 126019 at position 0, then check next token predictions
        Swift.print("\nProcessing token 126019 at position 0...")
        let logits = try model.forward(tokenId: 126019, position: 0)
        
        // Test different sampling strategies
        let top5 = logits.enumerated().sorted { $0.element > $1.element }.prefix(5)
        Swift.print("Top 5 logits:")
        for (idx, val) in top5 {
            let tokenStr = tokenizer.decode(tokens: [idx])
            Swift.print("  Token \(idx) ('\(tokenStr)'): \(val)")
        }
        
        // Test 1: Greedy decoding (temperature=0, take argmax)
        let sampler = Sampler()
        let argmax = sampler.greedySample(logits: logits)
        Swift.print("\nGreedy decoding: Token \(argmax) ('\(tokenizer.decode(tokens: [argmax]))')")
        
        // Test 2: Temperature=0.1 (near-greedy)
        let samples01 = (0..<5).map { _ in sampler.sample(logits: logits, temperature: 0.1, topK: 50, topP: 0.95) }
        Swift.print("\nTemperature=0.1, 5 samples:")
        for (i, token) in samples01.enumerated() {
            let tokenStr = tokenizer.decode(tokens: [token])
            Swift.print("  Sample \(i): Token \(token) ('\(tokenStr)')")
        }
        
        // Test 3: Temperature=0.7 (default)
        let samples07 = (0..<5).map { _ in sampler.sample(logits: logits, temperature: 0.7, topK: 50, topP: 0.95) }
        Swift.print("\nTemperature=0.7, 5 samples:")
        for (i, token) in samples07.enumerated() {
            let tokenStr = tokenizer.decode(tokens: [token])
            Swift.print("  Sample \(i): Token \(token) ('\(tokenStr)')")
        }
        
        // Test 4: Temperature=1.0 (more random)
        let samples10 = (0..<5).map { _ in sampler.sample(logits: logits, temperature: 1.0, topK: 50, topP: 0.95) }
        Swift.print("\nTemperature=1.0, 5 samples:")
        for (i, token) in samples10.enumerated() {
            let tokenStr = tokenizer.decode(tokens: [token])
            Swift.print("  Sample \(i): Token \(token) ('\(tokenStr)')")
        }
        
        // Test 5: Top-K=1 (force argmax)
        let samplesK1 = (0..<5).map { _ in sampler.sample(logits: logits, temperature: 1.0, topK: 1) }
        Swift.print("\nTop-K=1 (force argmax), 5 samples:")
        for (i, token) in samplesK1.enumerated() {
            let tokenStr = tokenizer.decode(tokens: [token])
            Swift.print("  Sample \(i): Token \(token) ('\(tokenStr)')")
        }
        
        // Check if all top 5 tokens are unused tokens
        let top5AreUnused = top5.allSatisfy { (idx, _) in idx >= 258000 && idx < 259000 }
        Swift.print("\nAll top 5 are unused tokens (258000-258999): \(top5AreUnused)")
        
        // Test with filtering enabled
        Swift.print("\n=== Testing with unused token filtering enabled ===")
        let samplesFiltered = (0..<5).map { _ in sampler.sample(logits: logits, temperature: 0.7, topK: 50, topP: 0.95, filterUnusedTokens: true) }
        Swift.print("Temperature=0.7 with filtering, 5 samples:")
        for (i, token) in samplesFiltered.enumerated() {
            let tokenStr = tokenizer.decode(tokens: [token])
            Swift.print("  Sample \(i): Token \(token) ('\(tokenStr)')")
        }
        
        // Test generation with filtering
        Swift.print("\n=== Generation test with filtering ===")
        for cache in model.kvCaches {
            cache.reset()
        }
        var tokens2: [Int] = [2]  // BOS
        for pos in 0..<10 {
            let nextLogits = try model.forward(tokenId: tokens2.last!, position: pos)
            let nextToken = sampler.sample(logits: nextLogits, temperature: 0.7, topK: 50, topP: 0.95, filterUnusedTokens: true)
            tokens2.append(nextToken)
            let tokenStr = tokenizer.decode(tokens: [nextToken])
            Swift.print("Position \(pos + 1): sampled token \(nextToken) ('\(tokenStr)')")
        }
        Swift.print("\nGenerated text: '\(tokenizer.decode(tokens: tokens2))'")
        
        // Test with meaningful prompt
        Swift.print("\n=== Generation with meaningful prompt ===")
        for cache in model.kvCaches {
            cache.reset()
        }
        
        // Encode prompt
        let prompt = "What is the capital of France?"
        var tokens3 = tokenizer.encode(text: prompt)
        Swift.print("Prompt tokens: \(tokens3)")
        Swift.print("Prompt: '\(tokenizer.decode(tokens: tokens3))'")
        
        // Process prompt
        for (pos, token) in tokens3.enumerated() {
            let logits = try model.forward(tokenId: token, position: pos)
            if pos == tokens3.count - 1 {
                // Sample from last position
                let nextToken = sampler.sample(logits: logits, temperature: 0.7, topK: 50, topP: 0.95, filterUnusedTokens: true)
                tokens3.append(nextToken)
                let tokenStr = tokenizer.decode(tokens: [nextToken])
                Swift.print("First generated token: \(nextToken) ('\(tokenStr)')")
            }
        }
        
        // Generate more tokens
        for i in 0..<20 {
            let pos = tokens3.count - 1
            let logits = try model.forward(tokenId: tokens3.last!, position: pos)
            let nextToken = sampler.sample(logits: logits, temperature: 0.7, topK: 50, topP: 0.95, filterUnusedTokens: true)
            tokens3.append(nextToken)
            let tokenStr = tokenizer.decode(tokens: [nextToken])
            Swift.print("Generated token \(i + 1): \(nextToken) ('\(tokenStr)')")
        }
        
        Swift.print("\nFinal output: '\(tokenizer.decode(tokens: tokens3))'")
        
        // Test with proper format tokens
        Swift.print("\n=== Generation with proper format tokens ===")
        for cache in model.kvCaches {
            cache.reset()
        }
        
        // Use proper Gemma-4 format: <|turn> user question <turn|> <|turn> model response <turn|>
        // Format: <bos> <|turn> user prompt <turn|> <|turn> model response
        let promptFormatted = "<bos><|turn>What is the capital of France?<turn|><|turn>"
        
        Swift.print("Formatted prompt: '\(promptFormatted)'")
        
        // Encode
        var tokens4 = tokenizer.encode(text: promptFormatted)
        Swift.print("Tokens: \(tokens4)")
        Swift.print("Decoded: '\(tokenizer.decode(tokens: tokens4))'")
        
        // Process
        for (pos, token) in tokens4.enumerated() {
            let logits = try model.forward(tokenId: token, position: pos)
            if pos == tokens4.count - 1 {
                // Sample from last position
                let nextToken = sampler.sample(logits: logits, temperature: 0.7, topK: 50, topP: 0.95, filterUnusedTokens: true)
                tokens4.append(nextToken)
                let tokenStr = tokenizer.decode(tokens: [nextToken])
                Swift.print("First generated token: \(nextToken) ('\(tokenStr)')")
            }
        }
        
        // Generate 20 tokens
        var tokensGenerated = tokens4
        for i in 0..<20 {
            let pos = tokensGenerated.count - 1
            let logits = try model.forward(tokenId: tokensGenerated.last!, position: pos)
            let nextToken = sampler.sample(logits: logits, temperature: 0.7, topK: 50, topP: 0.95, filterUnusedTokens: true)
            tokensGenerated.append(nextToken)
            let tokenStr = tokenizer.decode(tokens: [nextToken])
            Swift.print("Generated token \(i + 1): \(nextToken) ('\(tokenStr)')")
            
            // Stop if we hit end of turn or EOS
            if nextToken == 106 || nextToken == 1 {
                break
            }
        }
        
        Swift.print("\nFinal output: '\(tokenizer.decode(tokens: tokensGenerated))'")
    }
    
    func testCompareEmbedding126019() throws {
        Swift.print("\n=== Compare Token 126019 Embedding: Swift vs Python ===")
        let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
        
        // Get embedding for token 126019
        let embedWeight = model.embedWeight
        
        // Token 126019
        let tokenId = 126019
        let numGroups = 40
        let groupSize = 64
        let hiddenDim = 2560
        
        // Dequantize in Swift
        var dequantized = [Float](repeating: 0, count: hiddenDim)
        
        let weightPtr = embedWeight.weight.contents().assumingMemoryBound(to: UInt32.self)
        let scalesPtr = embedWeight.scales.contents().assumingMemoryBound(to: Float.self)
        let biasesPtr = embedWeight.biases.contents().assumingMemoryBound(to: Float.self)
        
        for groupIdx in 0..<numGroups {
            let scale = scalesPtr[tokenId * numGroups + groupIdx]
            let bias = biasesPtr[tokenId * numGroups + groupIdx]
            
            for u32Idx in 0..<8 {
                let packed = weightPtr[tokenId * 320 + groupIdx * 8 + u32Idx]
                
                for bitIdx in 0..<8 {
                    let val4bit = Int32((packed >> (UInt32(bitIdx) * 4)) & 0xF)
                    let signed = val4bit > 7 ? val4bit - 16 : val4bit
                    let deqValue = Float(signed) * scale + bias
                    
                    let idx = groupIdx * groupSize + u32Idx * 8 + bitIdx
                    dequantized[idx] = deqValue
                }
            }
        }
        
        // Print first 20 values
        Swift.print("Swift embedding for token 126019 (first 20 values):")
        for i in 0..<20 {
            Swift.print("  [\(i)] = \(dequantized[i])")
        }
        
        // Print some middle values
        Swift.print("\nSwift embedding for token 126019 (values 1000-1020):")
        for i in 1000..<1020 {
            Swift.print("  [\(i)] = \(dequantized[i])")
        }
        
        // Print last 20 values
        Swift.print("\nSwift embedding for token 126019 (last 20 values):")
        for i in 2540..<2560 {
            Swift.print("  [\(i)] = \(dequantized[i])")
        }
        
        // Compute norm
        let norm = sqrt(dequantized.reduce(0) { $0 + $1 * $1 })
        Swift.print("\nSwift embedding norm: \(norm)")
        
        // Compare scales/biases
        Swift.print("\nSwift scales (first 10):")
        for i in 0..<10 {
            Swift.print("  [\(i)] = \(scalesPtr[tokenId * numGroups + i])")
        }
        
        Swift.print("\nSwift biases (first 10):")
        for i in 0..<10 {
            Swift.print("  [\(i)] = \(biasesPtr[tokenId * numGroups + i])")
        }
    }
    
    func testUnusedTokenEmbedding() throws {
        Swift.print("\n=== Unused Token Embedding Test ===")
        let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
        let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
        
        // Test unused token 258123 as input
        Swift.print("\nProcessing unused token 258123 ('<unused2211>') at position 0...")
        let logits = try model.forward(tokenId: 258123, position: 0)
        
        let top10 = logits.enumerated().sorted { $0.element > $1.element }.prefix(10)
        Swift.print("Top 10 logits:")
        for (idx, val) in top10 {
            let tokenStr = tokenizer.decode(tokens: [idx])
            Swift.print("  Token \(idx) ('\(tokenStr)'): \(val)")
        }
        
        // Test unused token 258090
        Swift.print("\nProcessing unused token 258090 ('<unused2178>') at position 0...")
        let logits2 = try model.forward(tokenId: 258090, position: 0)
        
        let top10_2 = logits2.enumerated().sorted { $0.element > $1.element }.prefix(10)
        Swift.print("Top 10 logits:")
        for (idx, val) in top10_2 {
            let tokenStr = tokenizer.decode(tokens: [idx])
            Swift.print("  Token \(idx) ('\(tokenStr)'): \(val)")
        }
        
        // Compare with BOS token
        Swift.print("\nCompare with BOS token (2) at position 0...")
        let logitsBOS = try model.forward(tokenId: 2, position: 0)
        
        let top10_BOS = logitsBOS.enumerated().sorted { $0.element > $1.element }.prefix(10)
        Swift.print("Top 10 logits:")
        for (idx, val) in top10_BOS {
            let tokenStr = tokenizer.decode(tokens: [idx])
            Swift.print("  Token \(idx) ('\(tokenStr)'): \(val)")
        }
        
        // Test token 126019 ('ccql') - the token that led to position 6 unused predictions
        Swift.print("\nProcessing token 126019 ('ccql') at position 0...")
        let logits_ccql = try model.forward(tokenId: 126019, position: 0)
        
        let top10_ccql = logits_ccql.enumerated().sorted { $0.element > $1.element }.prefix(10)
        Swift.print("Top 10 logits:")
        for (idx, val) in top10_ccql {
            let tokenStr = tokenizer.decode(tokens: [idx])
            Swift.print("  Token \(idx) ('\(tokenStr)'): \(val)")
        }
        
        // Test token 225448 ('globalMap') - another sampled token
        Swift.print("\nProcessing token 225448 ('globalMap') at position 0...")
        let logits_global = try model.forward(tokenId: 225448, position: 0)
        
        let top10_global = logits_global.enumerated().sorted { $0.element > $1.element }.prefix(10)
        Swift.print("Top 10 logits:")
        for (idx, val) in top10_global {
            let tokenStr = tokenizer.decode(tokens: [idx])
            Swift.print("  Token \(idx) ('\(tokenStr)'): \(val)")
        }
    }

    // ── 31B dense model test ─────────────────────────

    func test31BModelLoading() throws {
        Swift.print("\n=== Testing Gemma-4 31B (dense, 4-bit) ===")

        let modelDir = "/Users/accusys/MarkBaseEngine/models/gemma-4-31b-it-4bit"

        guard FileManager.default.fileExists(atPath: modelDir + "/config.json") else {
            Swift.print("✗ 31B model not found")
            return
        }
        Swift.print("✓ 31B model found")

        let totalWeights: Int64 = try FileManager.default.contentsOfDirectory(atPath: modelDir)
            .filter { $0.hasSuffix(".safetensors") }
            .reduce(0) { sum, file in
                let attrs = try FileManager.default.attributesOfItem(atPath: modelDir + "/" + file)
                return sum + (attrs[.size] as? Int64 ?? 0)
            }
        Swift.print("  Total weights: \(totalWeights / 1024 / 1024 / 1024) GB")

        let engine = try MarkBaseEngine(autoCompile: true)
        Swift.print("✓ Engine created")

        Swift.print("\nLoading 31B model (~18 GB, may take 2-3 minutes)...")
        Swift.print("  hidden_size=5376, 60 layers, 32 heads, 16 KV heads, intermediate=21504")

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

        Swift.print("✓✓✓ Model loaded!")
        Swift.print("  Layers: \(model.numHiddenLayers)")
        Swift.print("  Hidden size: \(model.hiddenSize)")
        Swift.print("  Vocab size: \(model.vocabSize)")
        let fullCount = model.layerTypesIsFull.filter { $0 }.count
        let slideCount = model.layerTypesIsFull.filter { !$0 }.count
        Swift.print("  Layer types: \(fullCount) full, \(slideCount) sliding")

        let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
        Swift.print("✓ Tokenizer loaded")

        // Forward pass
        Swift.print("\nForward pass test...")
        let tokens = tokenizer.encode(text: "Hello")
        Swift.print("Tokenized 'Hello': \(tokens)")

        let logits = try model.forward(tokenId: tokens[0], position: 0)
        Swift.print("✓ Forward pass: \(logits.count) logits, max=\(logits.max() ?? -999)")

        let sorted = logits.enumerated().sorted { $0.element > $1.element }
        let top5 = sorted.prefix(5)
        Swift.print("Top 5:")
        for (idx, val) in top5 {
            let tokenStr = tokenizer.decode(tokens: [idx])
            Swift.print("  Token \(idx) ('\(tokenStr)'): \(val)")
        }

        // Generation test (3 tokens)
        Swift.print("\nGeneration test (3 tokens)...")
        do {
            let genConfig = GenerationConfig(maxTokens: 3, temperature: 0.7)
            let generator = StreamingGenerator(model: model, tokenizer: tokenizer, engine: engine)
            let response = try generator.generateComplete(prompt: "Hello", config: genConfig)
            Swift.print("Generated: '\(response)'")
            XCTAssertGreaterThan(response.count, 0, "Should generate text")
        } catch {
            Swift.print("⚠️ Generation failed: \(error)")
        }

        Swift.print("\n✅ 31B test complete!")
    }
}