import XCTest
@testable import MarkBase

class TwentySixBA4BNaNLocationTest: XCTestCase {
    
    func test26BA4BNaLLocation() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  26B-A4B NaN位置精确定位测试")
        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 at \(modelPath)")
            return
        }
        
        let engine = try MarkBaseEngine(autoCompile: true)
        
        print("Step 1: 載入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("Step 2: 測試不同tokens的NaN位置...")
        print("預期：如果NaN位置固定在[255999, 256000]，則是設計特性（像12B）")
        print()
        
        let testTokens = [2, 100, 200, 255999, 256000, 258880, 258881]
        var allNanPositions: [[Int]] = []
        
        for tokenId in testTokens {
            print("測試 Token \(tokenId):")
            do {
                let result = try model.forwardOptimized(tokenId: tokenId, position: 0)
                let nanIndices = result.enumerated()
                    .filter { $0.element.isNaN }
                    .map { $0.offset }
                
                print("  Total logits: \(result.count)")
                print("  NaN count: \(nanIndices.count)")
                
                if nanIndices.count > 0 {
                    print("  NaN positions: \(nanIndices)")
                    allNanPositions.append(nanIndices)
                    
                    // 檢查logit分布
                    let validLogits = result.filter { !$0.isNaN }
                    if validLogits.count > 0 {
                        print("  Valid logits: min=\(validLogits.min() ?? 0), max=\(validLogits.max() ?? 0)")
                    }
                } else {
                    print("  ✓✓ Zero NaN")
                    allNanPositions.append([])
                    let minVal = result.min() ?? 0
                    let maxVal = result.max() ?? 0
                    print("  Min/Max: \(minVal) / \(maxVal)")
                }
            } catch {
                print("  ✗ Error: \(error)")
                allNanPositions.append([])
            }
            print()
        }
        
        print("═══════════════════════════════════════════════════════════════════")
        print("  分析結果")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        // 檢查是否所有NaN位置都相同
        let nonEmptyPositions = allNanPositions.filter { !$0.isEmpty }
        if nonEmptyPositions.count > 0 {
            let firstPositions = nonEmptyPositions[0]
            let allSame = nonEmptyPositions.allSatisfy { $0 == firstPositions }
            
            if allSame {
                print("✅ 關鍵發現：")
                print("  NaN位置完全固定！")
                print("  固定位置: \(firstPositions)")
                print()
                
                // 檢查是否為多模態token位置
                let multimodalPositions = [255999, 256000, 258880, 258881, 258882, 258883]
                let isMultimodalRelated = firstPositions.allSatisfy { multimodalPositions.contains($0) }
                
                if isMultimodalRelated {
                    print("🎯 確認結論：")
                    print("  這是多模態設計特性，不是bug！")
                    print("  NaN位置對應多模態特殊tokens（像12B一樣）")
                    print()
                    print("多模態Token映射：")
                    if firstPositions.contains(255999) {
                        print("  255999 = BOI (Begin of Image)")
                    }
                    if firstPositions.contains(256000) {
                        print("  256000 = BOA (Begin of Audio)")
                    }
                    if firstPositions.contains(258880) {
                        print("  258880 = Image token")
                    }
                    if firstPositions.contains(258881) {
                        print("  258881 = Audio token")
                    }
                    print()
                    print("建議：")
                    print("  ✅ 可以繼續使用26B-A4B")
                    print("  ✅ 僅影響\(firstPositions.count)個位置（\(String(format: "%.6f", Double(firstPositions.count) / Double(model.vocabSize) * 100))%）")
                    print("  ✅ 使用26B-Standard作為替代（無NaN）")
                } else {
                    print("⚠️ 注意：")
                    print("  NaN位置不在預期的多模態token範圍內")
                    print("  需要進一步分析")
                }
            } else {
                print("⚠️ 警告：")
                print("  NaN位置不固定，可能存在其他問題")
                print("  不同tokens的NaN位置：")
                for (i, positions) in allNanPositions.enumerated() {
                    if !positions.isEmpty {
                        print("    Token \(testTokens[i]): \(positions)")
                    }
                }
            }
        } else {
            print("✅ 所有tokens的forward pass都無NaN")
            print("  這可能表示問題已被解決或間歇性問題")
        }
        
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  測試完成")
        print("═══════════════════════════════════════════════════════════════════\n")
    }
    
    func testCompareWith26BStandard() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  26B-A4B vs 26B-Standard 對比測試")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        let a4bPath = "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-a4b-it-4bit"
        let stdPath = "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-standard"
        
        guard FileManager.default.fileExists(atPath: a4bPath) else {
            print("⚠️ 26B-A4B not found")
            return
        }
        
        guard FileManager.default.fileExists(atPath: stdPath) else {
            print("⚠️ 26B-Standard not found")
            return
        }
        
        let engine = try MarkBaseEngine(autoCompile: true)
        
        // 測試26B-A4B
        print("=== 測試 26B-A4B ===")
        do {
            let a4bModel = try E4BModel(modelDir: a4bPath, engine: engine, maxContextLength: 128)
            print("✓ 26B-A4B載入成功")
            print("  Layers: \(a4bModel.numHiddenLayers)")
            print("  Hidden: \(a4bModel.hiddenSize)")
            
            let a4bResult = try a4bModel.forwardOptimized(tokenId: 2, position: 0)
            let a4bNan = a4bResult.filter { $0.isNaN }.count
            let a4bNanIndices = a4bResult.enumerated()
                .filter { $0.element.isNaN }
                .map { $0.offset }
            
            print("  Forward pass: \(a4bNan) NaN")
            if a4bNan > 0 {
                print("  NaN positions: \(a4bNanIndices)")
            }
        } catch {
            print("  ✗ Failed: \(error)")
        }
        
        print()
        
        // 測試26B-Standard
        print("=== 測試 26B-Standard ===")
        do {
            let stdModel = try E4BModel(modelDir: stdPath, engine: engine, maxContextLength: 128)
            print("✓ 26B-Standard載入成功")
            print("  Layers: \(stdModel.numHiddenLayers)")
            print("  Hidden: \(stdModel.hiddenSize)")
            
            let stdResult = try stdModel.forwardOptimized(tokenId: 2, position: 0)
            let stdNan = stdResult.filter { $0.isNaN }.count
            
            print("  Forward pass: \(stdNan) NaN")
            if stdNan == 0 {
                print("  ✓✓✓ 完美！無NaN")
            }
        } catch {
            print("  ✗ Failed: \(error)")
        }
        
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  對比完成")
        print("═══════════════════════════════════════════════════════════════════\n")
    }
}