import XCTest
@testable import MarkBase

class TwentySixBA4BLayerByLayerDebugTest: XCTestCase {
    
    func testSimpleLayerByLayerCheck() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  26B-A4B 简单逐层检查")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        let modelPath = "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-a4b-it-4bit"
        
        guard FileManager.default.fileExists(atPath: modelPath) else {
            print("⚠️ Model not found")
            return
        }
        
        let engine = try MarkBaseEngine(autoCompile: true)
        
        print("载入26B-A4B...")
        let model = try E4BModel(modelDir: modelPath, engine: engine, maxContextLength: 128)
        print("✓ 载入成功")
        print("  Layers: \(model.numHiddenLayers)")
        print("  Hidden: \(model.hiddenSize)")
        print("  Vocab: \(model.vocabSize)")
        print()
        
        print("=== 测试不同position的NaN ===")
        for pos in [0, 1, 2] {
            print("\nPosition \(pos):")
            let tokenId = 2
            let logits = try model.forwardOptimized(tokenId: tokenId, position: pos)
            let nanCount = logits.filter { $0.isNaN }.count
            let nanIndices = logits.enumerated().filter { $0.element.isNaN }.map { $0.offset }
            
            print("  NaN count: \(nanCount)")
            if nanCount > 0 {
                print("  NaN positions: \(nanIndices)")
                
                let validLogits = logits.filter { !$0.isNaN }
                if validLogits.count > 0 {
                    print("  Valid logits: min=\(validLogits.min() ?? 0), max=\(validLogits.max() ?? 0)")
                    print("  Mean: \(validLogits.reduce(0, +) / Float(validLogits.count))")
                }
            } else {
                let minVal = logits.min() ?? 0
                let maxVal = logits.max() ?? 0
                print("  ✓ No NaN! Min=\(minVal), Max=\(maxVal)")
            }
        }
        
        print("\n=== 关键观察 ===")
        print("1. Embedding始终正常（之前测试已确认）")
        print("2. Forward pass产生NaN")
        print("3. NaN位置固定在[2, 98]")
        print()
        
        print("假设：")
        print("- Token 2和98可能有特殊含义")
        print("- 可能是token ID被用作logits索引")
        print("- 可能是某个特定的vocab index有问题")
        
        print("\n检查vocab index 2和98的含义...")
        print("如果它们是特殊tokens（如padding, unk, bos），可能是设计特性")
        
        print("\n═══════════════════════════════════════════════════════════════════\n")
    }
    
    func testTokenIdsAsLogitsIndices() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  检查Token ID是否用作Logits索引")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        let modelPath = "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-a4b-it-4bit"
        guard FileManager.default.fileExists(atPath: modelPath) else { return }
        
        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelPath, engine: engine, maxContextLength: 128)
        
        print("假设：")
        print("如果Token ID被用作logits索引，那么：")
        print("- Token 2 → logit[2]产生NaN")
        print("- Token 98 → logit[98]产生NaN")
        print()
        
        print("测试不同Token IDs:")
        for tokenId in [2, 50, 98, 100, 500] {
            print("\nToken \(tokenId):")
            let logits = try model.forwardOptimized(tokenId: tokenId, position: 0)
            let nanIndices = logits.enumerated().filter { $0.element.isNaN }.map { $0.offset }
            
            print("  NaN positions: \(nanIndices)")
            
            if nanIndices.contains(tokenId) {
                print("  ⭐ 发现！Token ID \(tokenId) 在 NaN位置中")
                print("  这可能证实假设：Token ID用作logits索引")
            }
        }
        
        print("\n如果所有Token的NaN位置都包含Token ID本身：")
        print("→ 这可能是设计特性（像12B的多模态token屏蔽）")
        print("→ 不需要修复，是正常行为")
        
        print("\n═══════════════════════════════════════════════════════════════════\n")
    }
}