import XCTest
@testable import MarkBase

class ModelScalesComparisonTest: XCTestCase {
    
    func testCompareAllModelScales() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  26B-Standard vs 26B-A4B vs 31B Scales Comparison")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        // 26B-Standard (single file, custom quant)
        print("1. 26B-Standard:")
        let standardPath = "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-standard"
        let standardReader = try SafeTensorsReader(path: "\(standardPath)/model.safetensors")
        let standardTensors = standardReader.allTensors
        
        let standardScales = standardTensors.first { $0.name.contains("embed_tokens.scales") }
        if let s = standardScales {
            let data = try standardReader.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 nanCount = scales.filter { $0.isNaN }.count
            let negCount = scales.filter { $0 < 0 }.count
            print("  NaN: \(nanCount), Negative: \(negCount)")
        }
        
        // 26B-A4B (sharded, MLX-vlm 0.4.3)
        print("\n2. 26B-A4B:")
        let a4bPath = "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-a4b-it-4bit"
        let a4bIndex = try SafeTensorsIndex(modelDir: a4bPath)
        var a4bReaders: [String: SafeTensorsReader] = [:]
        for shardFile in a4bIndex.weightMap.values {
            if a4bReaders[shardFile] == nil {
                a4bReaders[shardFile] = try SafeTensorsReader(path: "\(a4bPath)/\(shardFile)")
            }
        }
        let a4bTensors = a4bReaders.values.flatMap { $0.allTensors }
        
        let a4bScales = a4bTensors.first { $0.name.contains("embed_tokens.scales") }
        if let s = a4bScales {
            let reader = a4bReaders[a4bIndex.weightMap[s.name] ?? "model-00001-of-00003.safetensors"]!
            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 nanCount = scales.filter { $0.isNaN }.count
            let negCount = scales.filter { $0 < 0 }.count
            print("  NaN: \(nanCount), Negative: \(negCount)")
        }
        
        // 31B (sharded, MLX-vlm 0.4.3)
        print("\n3. 31B:")
        let model31BPath = "/Users/accusys/MarkBaseEngine/models/gemma-4-31b-it-4bit"
        let model31BIndex = try SafeTensorsIndex(modelDir: model31BPath)
        var model31BReaders: [String: SafeTensorsReader] = [:]
        for shardFile in model31BIndex.weightMap.values {
            if model31BReaders[shardFile] == nil {
                model31BReaders[shardFile] = try SafeTensorsReader(path: "\(model31BPath)/\(shardFile)")
            }
        }
        let model31BTensors = model31BReaders.values.flatMap { $0.allTensors }
        
        let model31BScales = model31BTensors.first { $0.name.contains("embed_tokens.scales") }
        if let s = model31BScales {
            let reader = model31BReaders[model31BIndex.weightMap[s.name] ?? "model-00001-of-00004.safetensors"]!
            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 nanCount = scales.filter { $0.isNaN }.count
            let negCount = scales.filter { $0 < 0 }.count
            print("  NaN: \(nanCount), Negative: \(negCount)")
        }
        
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  Comparison Summary")
        print("═══════════════════════════════════════════════════════════════════\n")
        print("26B-Standard: custom quant, group_size=32")
        print("26B-A4B: MLX-vlm 0.4.3, group_size=64, affine")
        print("31B: MLX-vlm 0.4.3, group_size=64, affine")
        print("\nIf 31B scales similar to 26B-A4B (±0.01), it may have same problem!")
        print("═══════════════════════════════════════════════════════════════════\n")
    }
}