import Foundation

public struct VisionConfig: Codable {
    public let hiddenSize: Int
    public let numAttentionHeads: Int
    public let numHiddenLayers: Int
    public let headDim: Int
    public let globalHeadDim: Int
    public let intermediateSize: Int
    public let hiddenAct: String
    public let rmsNormEps: Float
    public let outputProjDims: Int
    public let patchSize: Int
    public let imageSize: Int
    
    public init(
        hiddenSize: Int = 768,
        numAttentionHeads: Int = 12,
        numHiddenLayers: Int = 12,
        headDim: Int = 64,
        globalHeadDim: Int = 64,
        intermediateSize: Int = 3072,
        hiddenAct: String = "gelu_pytorch_tanh",
        rmsNormEps: Float = 1e-6,
        outputProjDims: Int = 1536,
        patchSize: Int = 14,
        imageSize: Int = 224
    ) {
        self.hiddenSize = hiddenSize
        self.numAttentionHeads = numAttentionHeads
        self.numHiddenLayers = numHiddenLayers
        self.headDim = headDim
        self.globalHeadDim = globalHeadDim
        self.intermediateSize = intermediateSize
        self.hiddenAct = hiddenAct
        self.rmsNormEps = rmsNormEps
        self.outputProjDims = outputProjDims
        self.patchSize = patchSize
        self.imageSize = imageSize
    }
    
    public var numPatches: Int {
        (imageSize / patchSize) * (imageSize / patchSize)
    }
}