import XCTest
@testable import MarkBase

class A4BDiagnosticTest: XCTestCase {
    
    func testEmbeddingScalesNaNCheck() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  26B-A4B Embedding Scales NaN Diagnostic")
        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")
            return
        }
        
        // Load SafeTensors readers
        let index = try SafeTensorsIndex(modelDir: modelPath)
        
        var readers: [String: SafeTensorsReader] = [:]
        for shardFile in index.weightMap.values {
            if readers[shardFile] == nil {
                readers[shardFile] = try SafeTensorsReader(path: "\(modelPath)/\(shardFile)")
            }
        }
        
        print("Loaded \(readers.count) shard readers")
        
        // Find embed_tokens tensors
        let allTensors = readers.values.flatMap { $0.allTensors }
        
        let embedWeightTensor = allTensors.first { $0.name.contains("embed_tokens.weight") }
        let embedScalesTensor = allTensors.first { $0.name.contains("embed_tokens.scales") }
        let embedBiasesTensor = allTensors.first { $0.name.contains("embed_tokens.biases") }
        
        print("\nEmbedding tensors:")
        if let wt = embedWeightTensor {
            print("  weight: shape=\(wt.shape), dtype=\(wt.dtype), size=\(wt.dataSize)")
        }
        if let st = embedScalesTensor {
            print("  scales: shape=\(st.shape), dtype=\(st.dtype), size=\(st.dataSize)")
        }
        if let bt = embedBiasesTensor {
            print("  biases: shape=\(bt.shape), dtype=\(bt.dtype), size=\(bt.dataSize)")
        }
        
        // Read scales and check for NaN
        if let st = embedScalesTensor {
            print("\nChecking scales for NaN...")
            let reader = readers[index.weightMap[st.name] ?? "model-00001-of-00003.safetensors"]!
            let scalesData = try reader.read(tensor: st)
            
            // Parse scales as Float array
            let scales = scalesData.withUnsafeBytes { ptr in
                Array(ptr.assumingMemoryBound(to: Float.self))
            }
            
            let nanCount = scales.filter { $0.isNaN }.count
            let infCount = scales.filter { $0.isInfinite }.count
            
            print("  Total scales: \(scales.count)")
            print("  NaN count: \(nanCount)")
            print("  Inf count: \(infCount)")
            
            if nanCount > 0 {
                // Find NaN positions
                let nanIndices = scales.enumerated().filter { $0.element.isNaN }.map { $0.offset }
                print("  NaN positions (first 20): \(nanIndices.prefix(20))")
                
                // Calculate which tokens these belong to
                // Assuming scales shape = [vocabSize, groupSize] or [vocabSize]
                let vocabSize = 262144
                let groupSize = st.shape.count > 1 ? st.shape[1] : 1
                
                print("\n  Token mapping:")
                for idx in nanIndices.prefix(10) {
                    let tokenId = idx / groupSize
                    let groupId = idx % groupSize
                    print("    scales[\(idx)] = NaN → token \(tokenId), group \(groupId)")
                }
            }
            
            // Check scales distribution
            let validScales = scales.filter { !$0.isNaN && !$0.isInfinite }
            if !validScales.isEmpty {
                let avgScale = validScales.reduce(0, +) / Float(validScales.count)
                let minScale = validScales.min() ?? 0
                let maxScale = validScales.max() ?? 0
                print("\n  Valid scales distribution:")
                print("    min=\(minScale), max=\(maxScale), avg=\(avgScale)")
            }
        }
        
        // Read biases and check for NaN
        if let bt = embedBiasesTensor {
            print("\nChecking biases for NaN...")
            let reader = readers[index.weightMap[bt.name] ?? "model-00001-of-00003.safetensors"]!
            let biasesData = try reader.read(tensor: bt)
            
            let biases = biasesData.withUnsafeBytes { ptr in
                Array(ptr.assumingMemoryBound(to: Float.self))
            }
            
            let nanCount = biases.filter { $0.isNaN }.count
            let infCount = biases.filter { $0.isInfinite }.count
            
            print("  Total biases: \(biases.count)")
            print("  NaN count: \(nanCount)")
            print("  Inf count: \(infCount)")
            
            if nanCount > 0 {
                let nanIndices = biases.enumerated().filter { $0.element.isNaN }.map { $0.offset }
                print("  NaN positions (first 20): \(nanIndices.prefix(20))")
            }
        }
        
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  Diagnosis Complete")
        print("═══════════════════════════════════════════════════════════════════\n")
    }
}