import Foundation
import Metal

// ─────────────────────────────────────────────────────────────
// Float16 Support for MarkBase
// ─────────────────────────────────────────────────────────────

/// Float16 data type for memory-efficient computation
public typealias Float16 = Swift.Float16

/// Float16 buffer utilities
public enum Float16Utils {
    /// Convert Float32 array to Float16
    public static func toFloat16(_ values: [Float]) -> [Float16] {
        return values.map { Float16($0) }
    }
    
    /// Convert Float16 array to Float32
    public static func toFloat32(_ values: [Float16]) -> [Float] {
        return values.map { Float($0) }
    }
    
    /// Create MTLBuffer from Float16 array
    public static func makeBuffer(device: MTLDevice, values: [Float16]) -> MTLBuffer? {
        return device.makeBuffer(
            bytes: values,
            length: values.count * MemoryLayout<Float16>.stride,
            options: .storageModeShared
        )
    }
    
    /// Read Float16 from MTLBuffer
    public static func readFloat16(from buffer: MTLBuffer, count: Int) -> [Float16] {
        let ptr = buffer.contents().assumingMemoryBound(to: Float16.self)
        return Array(UnsafeBufferPointer(start: ptr, count: count))
    }
    
    /// Convert Float32 MTLBuffer to Float16
    public static func convertBuffer(
        from buffer: MTLBuffer,
        device: MTLDevice,
        count: Int
    ) -> MTLBuffer? {
        let float32Ptr = buffer.contents().assumingMemoryBound(to: Float.self)
        let float32Values = Array<Float>(UnsafeBufferPointer(start: float32Ptr, count: count))
        let float16Values = toFloat16(float32Values)
        return makeBuffer(device: device, values: float16Values)
    }
}

/// Float16 quantization for model weights
public struct Float16Quantizer {
    /// Quantize Float32 weights to Float16
    public static func quantize(weights: [Float]) -> [Float16] {
        return Float16Utils.toFloat16(weights)
    }
    
    /// Dequantize Float16 weights to Float32
    public static func dequantize(weights: [Float16]) -> [Float] {
        return Float16Utils.toFloat32(weights)
    }
    
    /// Calculate memory savings
    public static func memorySavings(float32Count: Int, float16Count: Int) -> Double {
        let float32Size = float32Count * MemoryLayout<Float>.stride
        let float16Size = float16Count * MemoryLayout<Float16>.stride
        return 1.0 - (Double(float16Size) / Double(float32Size))
    }
}