markbaseengine/Sources/MarkBase/Multimodal.swift

import Metal

public final class MultimodalModel {
    public let textModel: E4BModel
    public let audioTower: AudioTower12B?
    public let audioTowerFull: AudioTower?
    public let audioTowerE2B: AudioTowerE2B?
    public let visionTower: VisionTower12B?
    public let visionTowerFull: VisionTower?
    public let visionTowerE2B: VisionTowerE2B?

    public let audioTokenId: Int
    public let boaTokenId: Int
    public let eoaTokenId: Int
    public let imageTokenId: Int
    public let boiTokenId: Int
    public let eoiTokenId: Int

    private let audioEmbedBuffer: MTLBuffer
    private let visionEmbedBuffer: MTLBuffer

    public init(modelDir: String, engine: MarkBaseEngine, maxContextLength: Int) throws {
        audioTokenId = 258881
        boaTokenId = 256000
        eoaTokenId = 258883
        imageTokenId = 258882
        boiTokenId = 256001
        eoiTokenId = 258884

        textModel = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: maxContextLength)

        let device = engine.device
        let hs = textModel.hiddenSize
        audioEmbedBuffer = device.makeBuffer(length: 1024 * hs * 4)!
        visionEmbedBuffer = device.makeBuffer(length: 1024 * hs * 4)!

        // Try full VisionTower first (E4B-MarkBase format), fall back to E2B, then 12B
        print("Loading vision tower...")
        var vt: VisionTower? = nil
        var vtE2B: VisionTowerE2B? = nil
        let vcfg = loadVisionConfig(modelDir: modelDir)
        
        // Detect format: E4B (uint32 quantized) vs E2B (bfloat16)
        var isE2BVisionFormat = false
        if let reader = try? SafeTensorsReader(path: modelDir + "/model.safetensors") {
            let descriptors = reader.allDescriptors()
            let hasLinearWeight = descriptors.contains { $0.name.contains(".linear.weight") && $0.name.hasPrefix("vision_tower.") }
            let hasQuantized = descriptors.contains { $0.name.contains(".scales") && $0.name.hasPrefix("vision_tower.") }
            isE2BVisionFormat = hasLinearWeight && !hasQuantized
            print("  Detected format: \(isE2BVisionFormat ? "E2B (bfloat16)" : "E4B (uint32 quantized)")")
        }
        
        if let reader = try? SafeTensorsReader(path: modelDir + "/model.safetensors") {
            if isE2BVisionFormat {
                do {
                    vtE2B = try loadVisionTowerE2B(reader: reader, config: vcfg, engine: engine)
                    print("  ✓ VisionTowerE2B loaded successfully!")
                } catch {
                    print("  ✗ VisionTowerE2B loading failed: \(error)")
                }
            } else {
                do {
                    vt = try loadVisionTower(reader: reader, config: vcfg, engine: engine)
                    print("  ✓ Vision tower loaded successfully!")
                } catch {
                    print("  ✗ Vision tower loading failed: \(error)")
                }
            }
        } else {
            print("  ✗ Failed to create safetensors reader")
        }
        
        visionTowerFull = vt
        visionTowerE2B = vtE2B
        if vt != nil {
            print("  ✓ Full VisionTower (\(vt!.config.numHiddenLayers) layers)")
        } else if vtE2B != nil {
            print("  ✓ VisionTowerE2B (\(vtE2B!.config.numHiddenLayers) layers)")
        } else {
            print("  Full VisionTower not available, trying 12B variant...")
        }
        visionTower = try? VisionTower12B.load(modelDir: modelDir, engine: engine)
        if visionTower != nil {
            print("  ✓ VisionTower12B")
        }

        // Try full AudioTower - detect format (E2B bfloat16 vs E4B uint32 quantized)
        print("Loading audio tower...")
        let acfg = loadAudioConfig(modelDir: modelDir)
        
        // Detect format by checking first layer weight structure
        var isE2BFormat = false
        if let reader = try? SafeTensorsReader(path: modelDir + "/model.safetensors") {
            let descriptors = reader.allDescriptors()
            let hasLinearWeight = descriptors.contains { $0.name.contains(".linear.weight") && $0.name.hasPrefix("audio_tower.") }
            let hasScales = descriptors.contains { $0.name.contains(".scales") && $0.name.hasPrefix("audio_tower.") }
            isE2BFormat = hasLinearWeight && !hasScales
            print("  Detected format: \(isE2BFormat ? "E2B (bfloat16)" : "E4B (uint32 quantized)")")
        }
        
        // Load appropriate tower based on format
        if let reader = try? SafeTensorsReader(path: modelDir + "/model.safetensors") {
            if isE2BFormat {
                audioTowerE2B = try? loadAudioTowerE2B(reader: reader, config: acfg, engine: engine)
                audioTowerFull = nil
                if audioTowerE2B != nil {
                    print("  ✓ AudioTowerE2B (\(audioTowerE2B!.config.numHiddenLayers) layers)")
                }
            } else {
                audioTowerFull = try? loadAudioTower(reader: reader, config: acfg, engine: engine)
                audioTowerE2B = nil
                if audioTowerFull != nil {
                    print("  ✓ Full AudioTower (\(audioTowerFull!.config.numHiddenLayers) layers)")
                } else {
                    print("  Full AudioTower not available, trying 12B variant...")
                }
            }
        } else {
            audioTowerFull = nil
            audioTowerE2B = nil
        }
        
        audioTower = try? AudioTower12B.load(modelDir: modelDir, engine: engine)
        if audioTower != nil {
            print("  ✓ AudioTower12B")
        }
    }

    public var engine: MarkBaseEngine { textModel.engine }

    public func generateText(tokens: [Int], maxTokens: Int) throws -> [Int] {
        var generated: [Int] = tokens
        for _ in 0..<maxTokens {
            let logits = try textModel.forward(tokenId: generated.last ?? 0, position: generated.count - 1)
            var maxLogit = logits[0]
            var maxIdx = 0
            for j in 1..<logits.count {
                if logits[j] > maxLogit { maxLogit = logits[j]; maxIdx = j }
            }
            generated.append(maxIdx)
        }
        return generated
    }

    public func processAudio(audioFeatures: [[Float]]) throws -> [Float] {
        if let tower = audioTowerFull {
            let numFrames = audioFeatures.count
            let flatFeatures = audioFeatures.flatMap { $0 }
            let inputBuffer = engine.device.makeBuffer(bytes: flatFeatures, length: flatFeatures.count * 4)!
            let hs = tower.config.outputProjDims
            let outputBuffer = engine.device.makeBuffer(length: numFrames / 4 * hs * 4)!
            try tower.forward(inputBuffer: inputBuffer, seqLen: numFrames, outputBuffer: outputBuffer)
            let ptr = outputBuffer.contents().assumingMemoryBound(to: Float.self)
            return Array(UnsafeBufferPointer(start: ptr, count: numFrames / 4 * hs))
        } else if let tower = audioTowerE2B {
            let numFrames = audioFeatures.count
            let flatFeatures = audioFeatures.flatMap { $0 }
            let inputBuffer = engine.device.makeBuffer(bytes: flatFeatures, length: flatFeatures.count * 4)!
            let hs = tower.config.outputProjDims
            let outputBuffer = engine.device.makeBuffer(length: numFrames / 4 * hs * 4)!
            try tower.forward(inputBuffer: inputBuffer, seqLen: numFrames, outputBuffer: outputBuffer)
            let ptr = outputBuffer.contents().assumingMemoryBound(to: Float.self)
            return Array(UnsafeBufferPointer(start: ptr, count: numFrames / 4 * hs))
        } else if let tower = audioTower {
            let numFrames = audioFeatures.count
            let flatFeatures = audioFeatures.flatMap { $0 }
            let inputBuffer = engine.device.makeBuffer(bytes: flatFeatures, length: flatFeatures.count * 4)!
            try tower.forward(inputBuffer: inputBuffer, seqLen: numFrames, outputBuffer: audioEmbedBuffer)
            return Array(repeating: 0.0, count: 100)
        }
        throw WeightError.tensorNotFound("Audio tower not loaded")
    }

    public func processVision(patchEmbeddings: [Float], numPatches: Int) throws -> [Float] {
        if let tower = visionTowerFull {
            let inputBuffer = engine.device.makeBuffer(bytes: patchEmbeddings, length: patchEmbeddings.count * 4)!
            let hs = tower.config.hiddenSize
            let outputBuffer = engine.device.makeBuffer(length: numPatches * hs * 4)!
            try tower.forward(patchEmbeddings: inputBuffer, numPatches: numPatches, outputBuffer: outputBuffer)
            let ptr = outputBuffer.contents().assumingMemoryBound(to: Float.self)
            return Array(UnsafeBufferPointer(start: ptr, count: numPatches * hs))
        } else if let tower = visionTower {
            let inputBuffer = engine.device.makeBuffer(bytes: patchEmbeddings, length: patchEmbeddings.count * 4)!
            try tower.forward(patchEmbeddings: inputBuffer, numPatches: numPatches, outputBuffer: visionEmbedBuffer)
            let ptr = visionEmbedBuffer.contents().assumingMemoryBound(to: Float.self)
            return Array(UnsafeBufferPointer(start: ptr, count: numPatches * 3840))
        }
        throw WeightError.tensorNotFound("Vision tower not loaded")
    }
}

// ── Full VisionTower loading ────────────────────────────

func loadVisionConfig(modelDir: String) -> VisionConfig {
    let path = modelDir + "/config.json"
    guard let data = FileManager.default.contents(atPath: path),
          let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
          let vc = json["vision_config"] as? [String: Any] else {
        return VisionConfig()
    }
    return VisionConfig(
        hiddenSize: vc["hidden_size"] as? Int ?? 768,
        numAttentionHeads: vc["num_attention_heads"] as? Int ?? 12,
        numHiddenLayers: vc["num_hidden_layers"] as? Int ?? 16,
        headDim: vc["head_dim"] as? Int ?? 64,
        globalHeadDim: 64,
        intermediateSize: vc["intermediate_size"] as? Int ?? 3072,
        hiddenAct: "gelu_pytorch_tanh",
        rmsNormEps: (vc["rms_norm_eps"] as? NSNumber)?.floatValue ?? 1e-6,
        outputProjDims: 2560,
        patchSize: vc["patch_size"] as? Int ?? 16,
        imageSize: 224
    )
}

func loadVisionTower(reader: SafeTensorsReader, config: VisionConfig,
                       engine: MarkBaseEngine) throws -> VisionTower {
    print("Loading E4B Vision Tower with preload optimization...")
    let startTime = Date()
    
    // Collect all vision tensor descriptors
    let visionPrefix = "vision_tower."
    let embedPrefix = "embed_vision."
    let visionDescriptors = reader.allDescriptors().filter {
        $0.name.hasPrefix(visionPrefix) || $0.name.hasPrefix(embedPrefix)
    }
    
    print("  Found \(visionDescriptors.count) vision tensors")
    
    // Parallel preload all vision tensors
    let dispatchGroup = DispatchGroup()
    let loadQueue = DispatchQueue(label: "vision-preload-e4b", attributes: .concurrent)
    var loadedData: [Data?] = Array(repeating: nil, count: visionDescriptors.count)
    var loadErrors: [Error?] = Array(repeating: nil, count: visionDescriptors.count)
    
    for (idx, desc) in visionDescriptors.enumerated() {
        dispatchGroup.enter()
        loadQueue.async {
            do {
                let data = try reader.read(tensor: desc)
                loadedData[idx] = data
            } catch {
                loadErrors[idx] = error
            }
            dispatchGroup.leave()
        }
    }
    
    dispatchGroup.wait()
    
    // Check for errors
    for (idx, error) in loadErrors.enumerated() {
        if let err = error {
            throw WeightError.readFailed("Failed to preload vision tensor \(visionDescriptors[idx].name): \(err)")
        }
    }
    
    let preloadTime = Date().timeIntervalSince(startTime) * 1000
    print("  ✓ Parallel preloaded \(visionDescriptors.count) vision tensors in \(String(format: "%.1f", preloadTime))ms")
    
    // Convert to tensors/floats dictionaries (sequential, but from preloaded data)
    var tensors: [String: Data] = [:]
    var floats: [String: [Float]] = [:]
    
    for (idx, desc) in visionDescriptors.enumerated() {
        guard let data = loadedData[idx] else { continue }
        let name = desc.name
        switch desc.dtype {
        case .u32:
            tensors[name] = data
        case .f32:
            floats[name] = data.withUnsafeBytes {
                Array($0.assumingMemoryBound(to: Float.self))
            }
        case .bf16:
            floats[name] = SafeTensorsReader.bf16ToFloat32(data)
        default:
            break
        }
    }

    guard !tensors.isEmpty, !floats.isEmpty else {
        throw WeightError.tensorNotFound("Vision tower tensors")
    }

    let weights = try VisionWeights(device: engine.device, config: config,
                                     tensors: tensors, floats: floats)
    
    let totalTime = Date().timeIntervalSince(startTime) * 1000
    print("  ✓ E4B Vision Tower loaded in \(String(format: "%.1f", totalTime))ms")
    
    return try VisionTower(config: config, engine: engine, weights: weights)
}

// ── Full AudioTower loading ────────────────────────────

func loadAudioConfig(modelDir: String) -> AudioConfig {
    let path = modelDir + "/config.json"
    guard let data = FileManager.default.contents(atPath: path),
          let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
          let ac = json["audio_config"] as? [String: Any] else {
        return AudioConfig()
    }
    return AudioConfig(
        hiddenSize: ac["hidden_size"] as? Int ?? 1024,
        numAttentionHeads: ac["num_attention_heads"] as? Int ?? 8,
        numHiddenLayers: ac["num_hidden_layers"] as? Int ?? 12,
        convKernelSize: ac["conv_kernel_size"] as? Int ?? 5,
        attentionChunkSize: ac["attention_chunk_size"] as? Int ?? 12,
        attentionContextLeft: ac["attention_context_left"] as? Int ?? 13,
        attentionContextRight: ac["attention_context_right"] as? Int ?? 0,
        attentionLogitCap: (ac["attention_logit_cap"] as? NSNumber)?.floatValue ?? 50.0,
        hiddenAct: ac["hidden_act"] as? String ?? "silu",
        rmsNormEps: (ac["rms_norm_eps"] as? NSNumber)?.floatValue ?? 1e-6,
        outputProjDims: ac["output_proj_dims"] as? Int ?? 1536,
        subsamplingConvChannels: [128, 32],
        residualWeight: 0.5
    )
}

func loadAudioTower(reader: SafeTensorsReader, config: AudioConfig,
                    engine: MarkBaseEngine) throws -> AudioTower {
    print("Loading E4B Audio Tower with preload optimization...")
    let startTime = Date()
    
    // Collect all audio tensor descriptors
    let audioPrefix = "audio_tower."
    let audioDescriptors = reader.allDescriptors().filter {
        $0.name.hasPrefix(audioPrefix)
    }
    
    print("  Found \(audioDescriptors.count) audio tensors")
    
    // Parallel preload all audio tensors
    let dispatchGroup = DispatchGroup()
    let loadQueue = DispatchQueue(label: "audio-preload-e4b", attributes: .concurrent)
    var loadedData: [Data?] = Array(repeating: nil, count: audioDescriptors.count)
    var loadErrors: [Error?] = Array(repeating: nil, count: audioDescriptors.count)
    
    for (idx, desc) in audioDescriptors.enumerated() {
        dispatchGroup.enter()
        loadQueue.async {
            do {
                let data = try reader.read(tensor: desc)
                loadedData[idx] = data
            } catch {
                loadErrors[idx] = error
            }
            dispatchGroup.leave()
        }
    }
    
    dispatchGroup.wait()
    
    // Check for errors
    for (idx, error) in loadErrors.enumerated() {
        if let err = error {
            throw WeightError.readFailed("Failed to preload audio tensor \(audioDescriptors[idx].name): \(err)")
        }
    }
    
    let preloadTime = Date().timeIntervalSince(startTime) * 1000
    print("  ✓ Parallel preloaded \(audioDescriptors.count) audio tensors in \(String(format: "%.1f", preloadTime))ms")
    
    // Convert to tensors/floats/descriptors dictionaries
    var tensors: [String: Data] = [:]
    var floats: [String: [Float]] = [:]
    var descriptors: [String: TensorDescriptor] = [:]

    for (idx, desc) in audioDescriptors.enumerated() {
        guard let data = loadedData[idx] else { continue }
        let name = desc.name
        descriptors[name] = desc
        switch desc.dtype {
        case .u32:
            tensors[name] = data
        case .f32:
            floats[name] = data.withUnsafeBytes {
                Array($0.assumingMemoryBound(to: Float.self))
            }
        case .bf16:
            floats[name] = SafeTensorsReader.bf16ToFloat32(data)
        default:
            break
        }
    }

    guard !tensors.isEmpty, !floats.isEmpty else {
        throw WeightError.tensorNotFound("Audio tower tensors")
    }

    let weights = try AudioWeights(device: engine.device, config: config,
                                    tensors: tensors, floats: floats,
                                    descriptors: descriptors)
    
    let totalTime = Date().timeIntervalSince(startTime) * 1000
    print("  ✓ E4B Audio Tower loaded in \(String(format: "%.1f", totalTime))ms")
    
    return try AudioTower(config: config, engine: engine, weights: weights)
}