import XCTest
@testable import MarkBase

class A4BComparisonTest: XCTestCase {
    
    func testCompareEmbeddingTensors() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  26B-A4B vs 26B-Standard Embedding Tensor Comparison")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        // Load 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 }
        
        // Load 26B-Standard (single file)
        let standardPath = "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-standard"
        let standardReader = try SafeTensorsReader(path: "\(standardPath)/model.safetensors")
        let standardTensors = standardReader.allTensors
        
        print("26B-A4B tensors: \(a4bTensors.count)")
        print("26B-Standard tensors: \(standardTensors.count)")
        
        // Compare embed_tokens
        print("\nEmbedding tensor comparison:")
        
        // A4B
        let a4bWeight = a4bTensors.first { $0.name.contains("embed_tokens.weight") }
        let a4bScales = a4bTensors.first { $0.name.contains("embed_tokens.scales") }
        let a4bBiases = a4bTensors.first { $0.name.contains("embed_tokens.biases") }
        
        print("\n26B-A4B:")
        if let w = a4bWeight { print("  weight: \(w.shape), \(w.dtype)") }
        if let s = a4bScales { print("  scales: \(s.shape), \(s.dtype)") }
        if let b = a4bBiases { print("  biases: \(b.shape), \(b.dtype)") }
        
        // Standard
        let standardWeight = standardTensors.first { $0.name.contains("embed_tokens.weight") }
        let standardScales = standardTensors.first { $0.name.contains("embed_tokens.scales") }
        let standardBiases = standardTensors.first { $0.name.contains("embed_tokens.biases") }
        
        print("\n26B-Standard:")
        if let w = standardWeight { print("  weight: \(w.shape), \(w.dtype)") }
        if let s = standardScales { print("  scales: \(s.shape), \(s.dtype)") }
        if let b = standardBiases { print("  biases: \(b.shape), \(b.dtype)") }
        
        // Compare first few weight values (if uint32, sample them)
        if let w = a4bWeight, w.dtype == .u32 {
            let reader = a4bReaders[a4bIndex.weightMap[w.name] ?? "model-00001-of-00003.safetensors"]!
            let data = try reader.read(tensor: w)
            let uint32s = data.withUnsafeBytes { ptr in
                Array(ptr.bindMemory(to: UInt32.self).prefix(20))
            }
            print("\n26B-A4B weight sample (first 20 uint32): \(uint32s)")
        }
        
        if let w = standardWeight, w.dtype == .u32 {
            let data = try standardReader.read(tensor: w)
            let uint32s = data.withUnsafeBytes { ptr in
                Array(ptr.bindMemory(to: UInt32.self).prefix(20))
            }
            print("\n26B-Standard weight sample (first 20 uint32): \(uint32s)")
        }
        
        // Compare 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("\n26B-A4B scales sample (first 20): \(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("\n26B-Standard scales sample (first 20): \(scales)")
        }
        
        print("\n═══════════════════════════════════════════════════════════════════")
    }
}