import XCTest
@testable import MarkBase

class TwelveBSpecialTokenTest: XCTestCase {
    func testSpecialTokenDebug() throws {
        print("\n=== 12B特殊Token深度Debug測試 ===\n")
        
        let engine = try MarkBaseEngine(autoCompile: true)
        let modelPath = "/Users/accusys/MarkBaseEngine/models/gemma-4-12b-it-4bit"
        
        let model = try E4BModel(modelDir: modelPath, engine: engine, maxContextLength: 128)
        
        print("Model info:")
        print("  Layers: \(model.numHiddenLayers)")
        print("  Hidden: \(model.hiddenSize)")
        print("  Vocab: \(model.vocabSize)")
        print()
        
        let specialTokens = [2, 255999, 256000]
        
        for tokenId in specialTokens {
            print("Testing Token \(tokenId):")
            
            do {
                let result = try model.forwardOptimized(tokenId: tokenId, position: 0)
                let nanCount = result.filter { $0.isNaN }.count
                
                print("  Total logits: \(result.count)")
                print("  NaN count: \(nanCount)")
                
                if nanCount > 0 {
                    print("  NaN indices: ")
                    for (idx, val) in result.enumerated() {
                        if val.isNaN {
                            print("    Index \(idx): NaN")
                        }
                    }
                }
                
                let validLogits = result.filter { !$0.isNaN }
                if validLogits.count > 0 {
                    print("  Valid logits: min=\(validLogits.min() ?? 0), max=\(validLogits.max() ?? 0)")
                }
                
            } catch {
                print("  ✗ Error: \(error)")
            }
            print()
        }
        
        print("Testing Token 100 (normal):")
        let normalResult = try model.forwardOptimized(tokenId: 100, position: 0)
        let normalNan = normalResult.filter { $0.isNaN }.count
        print("  NaN count: \(normalNan)")
        print("  Min/Max: \(normalResult.min() ?? 0) / \(normalResult.max() ?? 0)")
    }
}