import XCTest
@testable import MarkBase

class SimpleLogitsDebugTest: XCTestCase {
    func testLogitsDebug() throws {
        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("\n=== Testing Token 2 ===")
        let logits = try model.forward(tokenId: 2, position: 0, debug: true)
        
        print("\nLogits sample (first 20):")
        print(logits.prefix(20))
        
        print("\nNaN count: \(logits.filter { $0.isNaN }.count)")
        print("Inf count: \(logits.filter { $0.isInfinite }.count)")
        
        let validLogits = logits.filter { !$0.isNaN && !$0.isInfinite }
        if validLogits.count > 0 {
            print("Max valid logit: \(validLogits.max() ?? 0)")
        }
        
        let nanIndices = logits.enumerated().filter { $0.element.isNaN }.map { $0.offset }
        print("NaN positions: \(nanIndices)")
    }
}