import XCTest
@testable import MarkBase

class TwentySixBA4BDeepDebugTest: XCTestCase {
    
    func test26BA4BTokenByTokenDebug() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  26B-A4B Token-by-Token NaN Debug")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        let engine = try MarkBaseEngine(autoCompile: true)
        let modelPath = "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-a4b-it-4bit"
        
        print("Step 1: 載入26B-A4B...")
        let model = try E4BModel(modelDir: modelPath, engine: engine, maxContextLength: 128)
        print("✓ 載入成功")
        print("  Layers: \(model.numHiddenLayers)")
        print("  Vocab: \(model.vocabSize)")
        print()
        
        print("Step 2: Token-by-Token測試（找出NaN模式）...")
        print()
        
        // 測試關鍵token範圍
        let testRanges = [
            (0, 10),     // 特殊tokens
            (90, 110),   // 包含98的範圍
            (250, 260),  // 中間範圍
            (255990, 256010), // BOI/BOA附近
        ]
        
        var nanPattern: [Int: [Int]] = [:] // tokenId -> NaN positions
        
        for (start, end) in testRanges {
            print("測試範圍 \(start)-\(end):")
            for tokenId in stride(from: start, through: end, by: 1) {
                do {
                    let result = try model.forwardOptimized(tokenId: tokenId, position: 0)
                    let nanIndices = result.enumerated()
                        .filter { $0.element.isNaN }
                        .map { $0.offset }
                    
                    if nanIndices.count > 0 {
                        print("  Token \(tokenId): \(nanIndices.count) NaN at \(nanIndices)")
                        nanPattern[tokenId] = nanIndices
                    }
                } catch {
                    print("  Token \(tokenId): Error")
                }
            }
            print()
        }
        
        print("═══════════════════════════════════════════════════════════════════")
        print("  NaN模式分析")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        if nanPattern.isEmpty {
            print("✅ 所有測試tokens無NaN")
        } else {
            print("⚠️ 發現NaN模式：")
            for (tokenId, positions) in nanPattern.sorted(by: { $0.key < $1.key }) {
                print("  Token \(tokenId): NaN at \(positions)")
            }
            
            // 分析NaN位置模式
            let allPositions = Set(nanPattern.values.flatMap { $0 })
            print("\n所有NaN位置集合: \(allPositions.sorted())")
            
            // 檢查是否固定
            let uniquePositionSets = Set(nanPattern.values.map { Set($0) })
            if uniquePositionSets.count == 1 {
                print("✅ NaN位置固定（所有tokens相同）")
                print("   這可能是設計特性或forward pass bug")
            } else {
                print("⚠️ NaN位置不固定（不同tokens有不同NaN）")
                print("   這暗示可能是embedding或計算問題")
            }
        }
        
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  測試完成")
        print("═══════════════════════════════════════════════════════════════════\n")
    }
    
    func test26BA4BNaNPatternSimple() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  26B-A4B NaN Pattern Simple Analysis")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        let engine = try MarkBaseEngine(autoCompile: true)
        let modelPath = "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-a4b-it-4bit"
        
        print("載入模型...")
        let model = try E4BModel(modelDir: modelPath, engine: engine, maxContextLength: 128)
        print("✓ 載入成功\n")
        
        print("測試已知NaN位置：Token 2和98")
        print("對比：Token 100（應無NaN）\n")
        
        // Token 2
        print("Token 2 (BOS):")
        let logits2 = try model.forwardOptimized(tokenId: 2, position: 0)
        let nanLogits2 = logits2.filter { $0.isNaN }.count
        let nanIndices2 = logits2.enumerated().filter { $0.element.isNaN }.map { $0.offset }
        print("  Forward NaN: \(nanLogits2)/\(logits2.count)")
        print("  NaN positions: \(nanIndices2)")
        
        // Token 98  
        print("\nToken 98:")
        let logits98 = try model.forwardOptimized(tokenId: 98, position: 0)
        let nanLogits98 = logits98.filter { $0.isNaN }.count
        let nanIndices98 = logits98.enumerated().filter { $0.element.isNaN }.map { $0.offset }
        print("  Forward NaN: \(nanLogits98)/\(logits98.count)")
        print("  NaN positions: \(nanIndices98)")
        
        // Token 100（对照组）
        print("\nToken 100 (对照组):")
        let logits100 = try model.forwardOptimized(tokenId: 100, position: 0)
        let nanLogits100 = logits100.filter { $0.isNaN }.count
        let nanIndices100 = logits100.enumerated().filter { $0.element.isNaN }.map { $0.offset }
        print("  Forward NaN: \(nanLogits100)/\(logits100.count)")
        print("  NaN positions: \(nanIndices100)")
        
        // Token 200（对照组）
        print("\nToken 200 (对照组):")
        let logits200 = try model.forwardOptimized(tokenId: 200, position: 0)
        let nanLogits200 = logits200.filter { $0.isNaN }.count
        let nanIndices200 = logits200.enumerated().filter { $0.element.isNaN }.map { $0.offset }
        print("  Forward NaN: \(nanLogits200)/\(logits200.count)")
        print("  NaN positions: \(nanIndices200)")
        
        print("\n分析:")
        if nanIndices2 == nanIndices98 {
            print("  ✅ Token 2和98的NaN位置完全相同")
            print("  → NaN位置固定（可能是forward pass bug）")
        } else {
            print("  ⚠️ Token 2和98的NaN位置不同")
            print("  → NaN位置不固定（需要进一步分析）")
        }
        
        print("\n═══════════════════════════════════════════════════════════════════\n")
    }
}