import XCTest
@testable import MarkBase

class Model12BVisionCheckTest: XCTestCase {
    
    func test12BVisionWeights() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  12B Standard Vision Weights Check")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        let modelPath = "/Users/accusys/.cache/huggingface/hub/models--mlx-community--gemma-4-12B-it-4bit/snapshots/73bcf09092aa277861d5a191b989b666f7f32e8f"
        
        guard FileManager.default.fileExists(atPath: modelPath) else {
            print("⚠ Model not found")
            return
        }
        
        // Check shard 2 (VisionTower12B.load uses model-00002-of-00002.safetensors)
        let shard2Path = "\(modelPath)/model-00002-of-00002.safetensors"
        guard FileManager.default.fileExists(atPath: shard2Path) else {
            print("⚠ Shard 2 not found")
            return
        }
        
        let reader = try SafeTensorsReader(path: shard2Path)
        let tensors = reader.allTensors
        
        print("1. Shard 2 Total Tensors: \(tensors.count)")
        
        // Check vision_embedder weights
        print("\n2. Vision Embedder Weights:")
        let visionEmbedderTensors = tensors.filter { $0.name.contains("vision_embedder") }
        print("  vision_embedder tensors: \(visionEmbedderTensors.count)")
        for tensor in visionEmbedderTensors {
            print("    \(tensor.name): shape=\(tensor.shape), dtype=\(tensor.dtype)")
        }
        
        // Check embed_vision weights
        print("\n3. Embed Vision Weights:")
        let embedVisionTensors = tensors.filter { $0.name.contains("embed_vision") }
        print("  embed_vision tensors: \(embedVisionTensors.count)")
        for tensor in embedVisionTensors {
            print("    \(tensor.name): shape=\(tensor.shape), dtype=\(tensor.dtype)")
        }
        
        // Check audio weights
        print("\n4. Audio Weights:")
        let audioTensors = tensors.filter { $0.name.contains("audio") }
        print("  audio tensors: \(audioTensors.count)")
        for tensor in audioTensors.prefix(10) {
            print("    \(tensor.name): shape=\(tensor.shape)")
        }
        
        // Summary
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  Vision Weights Summary")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        if visionEmbedderTensors.count > 0 {
            print("✓ 12B HAS VISION EMBEDDER: \(visionEmbedderTensors.count) tensors")
        } else {
            print("✗ 12B NO VISION EMBEDDER")
        }
        
        if embedVisionTensors.count > 0 {
            print("✓ 12B HAS EMBED VISION: \(embedVisionTensors.count) tensors")
        } else {
            print("✗ 12B NO EMBED VISION")
        }
        
        if audioTensors.count > 0 {
            print("✓ 12B HAS AUDIO: \(audioTensors.count) tensors")
        } else {
            print("✗ 12B NO AUDIO")
        }
        
        print("\nVision Capability: \(visionEmbedderTensors.count > 0 && embedVisionTensors.count > 0 ? "YES ✓" : "NO ✗")")
        print("Audio Capability: \(audioTensors.count > 0 ? "YES ✓" : "NO ✗")")
        
        print("\n═══════════════════════════════════════════════════════════════════")
    }
}