//
//  QuantConfig.swift
//  Zyquo MLX
//
//  Author: Simon-Pierre Boucher
//  Mail: contact@spboucher.ai
//

import Foundation

/// Quantization recipe (docs/MLX-RESEARCH.md §5: affine supports bits
/// 2/3/4/5/6/8 with group sizes 32/64/128; mxfp4/nvfp4/mxfp8 have fixed
/// bits/groups).
struct QuantConfig: Codable, Hashable, Sendable {
    enum Mode: String, Codable, CaseIterable, Sendable {
        case affine
        case mxfp4
        case nvfp4
        case mxfp8
    }

    var bits: Int = 4
    var groupSize: Int = 64
    var mode: Mode = .affine

    static let affineBits = [2, 3, 4, 5, 6, 8]
    static let affineGroupSizes = [32, 64, 128]

    /// Effective bytes/param for size previews (group-scale overhead included).
    var bytesPerParameter: Double {
        (Double(bits) + 32.0 / Double(groupSize)) / 8.0
    }

    /// Predicted weight size after quantizing a model with the given
    /// parameter count (before/after preview in the Convert UI).
    func predictedWeightBytes(parameterCount: Int64) -> Int64 {
        Int64(Double(parameterCount) * bytesPerParameter)
    }

    var label: String { "\(bits)-bit \(mode.rawValue) (g\(groupSize))" }
}