import XCTest
@testable import MarkBase

class E4BMarkBaseTest: XCTestCase {
    
    func testE4BTextPerformance() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  E4B-MarkBase TEXT Performance Test")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        let modelPath = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        
        guard FileManager.default.fileExists(atPath: modelPath) else {
            print("⚠ Model not found")
            return
        }
        
        // Check scales quality
        print("1. Scales Quality Check:")
        let reader = try SafeTensorsReader(path: "\(modelPath)/model.safetensors")
        let tensors = reader.allTensors
        
        let embedScales = tensors.first { $0.name.contains("embed_tokens.scales") }
        if let s = embedScales {
            let data = try reader.read(tensor: s)
            let scales = data.withUnsafeBytes { ptr in
                Array(ptr.assumingMemoryBound(to: Float.self).prefix(20))
            }
            print("  Scales shape: \(s.shape), dtype: \(s.dtype)")
            print("  Sample: \(scales)")
            let negCount = scales.filter { $0 < 0 }.count
            print("  Negative: \(negCount)")
        }
        
        // Test TEXT inference
        print("\n2. TEXT Inference Test:")
        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelPath, engine: engine, maxContextLength: 128)
        print("  ✓ Model loaded (Layers: \(model.numHiddenLayers), Hidden: \(model.hiddenSize))")
        
        // Warmup
        _ = try model.forwardOptimized(tokenId: 2, position: 0)
        
        // Test NaN
        print("\n3. NaN Test (tokenIds 0-5):")
        var nanCount = 0
        for tokenId in 0..<5 {
            let result = try model.forwardOptimized(tokenId: tokenId, position: 0)
            let nans = result.filter { $0.isNaN }.count
            if nans > 0 { nanCount += nans }
        }
        print("  NaN count: \(nanCount)")
        
        // Test speed
        print("\n4. Speed Test (10 tokens):")
        let testStart = Date()
        var currentToken = 2
        for i in 0..<10 {
            let result = try model.forwardOptimized(tokenId: currentToken, position: i)
            var maxIdx = 0
            var maxVal = result[0]
            for j in 1..<result.count {
                if result[j] > maxVal {
                    maxVal = result[j]
                    maxIdx = j
                }
            }
            currentToken = maxIdx
        }
        let testTime = Date().timeIntervalSince(testStart) * 1000
        let avgTime = testTime / 10.0
        
        print("  Average: \(String(format: "%.1f", avgTime))ms per token")
        print("  Speed: \(String(format: "%.1f", 1000.0 / avgTime)) tok/s")
        
        if avgTime < 100 && nanCount == 0 {
            print("\n✓✓✓ E4B-MarkBase PRODUCTION READY")
        } else {
            print("\n⚠ Issues detected")
        }
        
        print("\n═══════════════════════════════════════════════════════════════════")
    }
}