import XCTest
@testable import MarkBase

class Model12BArchitectureTest: XCTestCase {
    
    func testModel12BFullArchitecture() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  12B Standard Complete Architecture 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
        }
        
        // Load tensors
        let index = try SafeTensorsIndex(modelDir: modelPath)
        var readers: [String: SafeTensorsReader] = [:]
        for shardFile in Set(index.weightMap.values) {
            readers[shardFile] = try SafeTensorsReader(path: "\(modelPath)/\(shardFile)")
        }
        
        let allTensors = readers.values.flatMap { $0.allTensors }
        
        print("1. Total Tensors: \(allTensors.count)")
        
        // Audio tower analysis
        print("\n2. Audio Tower:")
        let audioTensors = allTensors.filter { $0.name.contains("audio_tower") || $0.name.contains("audio_model") }
        print("  Audio tensors: \(audioTensors.count)")
        if !audioTensors.isEmpty {
            print("  Sample tensors:")
            for tensor in audioTensors.prefix(10) {
                print("    \(tensor.name): shape=\(tensor.shape)")
            }
        }
        
        // Vision tower analysis
        print("\n3. Vision Tower:")
        let visionTensors = allTensors.filter { $0.name.contains("vision_tower") || $0.name.contains("vision_model") }
        print("  Vision tensors: \(visionTensors.count)")
        if !visionTensors.isEmpty {
            print("  Sample tensors:")
            for tensor in visionTensors.prefix(10) {
                print("    \(tensor.name): shape=\(tensor.shape)")
            }
        }
        
        // TEXT model analysis
        print("\n4. TEXT Model:")
        let textTensors = allTensors.filter { $0.name.contains("language_model") || $0.name.contains("layers.") }
        print("  TEXT-related tensors: \(textTensors.count)")
        
        // Embed tokens
        let embedTokens = allTensors.filter { $0.name.contains("embed_tokens") }
        print("  Embed tokens: \(embedTokens.count)")
        for tensor in embedTokens.prefix(5) {
            print("    \(tensor.name): shape=\(tensor.shape)")
        }
        
        // Summary
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  Architecture Summary")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        if audioTensors.count > 0 {
            print("✓ 12B HAS AUDIO TOWER: \(audioTensors.count) tensors")
        } else {
            print("✗ 12B NO AUDIO TOWER")
        }
        
        if visionTensors.count > 0 {
            print("✓ 12B HAS VISION TOWER: \(visionTensors.count) tensors")
        } else {
            print("✗ 12B NO VISION TOWER")
        }
        
        print("✓ TEXT Model: \(textTensors.count) tensors")
        
        print("\nModel Type: \(audioTensors.count > 0 || visionTensors.count > 0 ? "Multimodal" : "Pure TEXT")")
        
        print("\n═══════════════════════════════════════════════════════════════════")
    }
}