import Foundation

public struct AudioConfig: Codable {
    public let hiddenSize: Int
    public let numAttentionHeads: Int
    public let numHiddenLayers: Int
    public let convKernelSize: Int
    public let attentionChunkSize: Int
    public let attentionContextLeft: Int
    public let attentionContextRight: Int
    public let attentionLogitCap: Float
    public let hiddenAct: String
    public let rmsNormEps: Float
    public let outputProjDims: Int
    public let subsamplingConvChannels: [Int]
    public let residualWeight: Float
    
    public init(
        hiddenSize: Int = 1024,
        numAttentionHeads: Int = 8,
        numHiddenLayers: Int = 12,
        convKernelSize: Int = 5,
        attentionChunkSize: Int = 12,
        attentionContextLeft: Int = 13,
        attentionContextRight: Int = 0,
        attentionLogitCap: Float = 50.0,
        hiddenAct: String = "silu",
        rmsNormEps: Float = 1e-6,
        outputProjDims: Int = 1536,
        subsamplingConvChannels: [Int] = [128, 32],
        residualWeight: Float = 0.5
    ) {
        self.hiddenSize = hiddenSize
        self.numAttentionHeads = numAttentionHeads
        self.numHiddenLayers = numHiddenLayers
        self.convKernelSize = convKernelSize
        self.attentionChunkSize = attentionChunkSize
        self.attentionContextLeft = attentionContextLeft
        self.attentionContextRight = attentionContextRight
        self.attentionLogitCap = attentionLogitCap
        self.hiddenAct = hiddenAct
        self.rmsNormEps = rmsNormEps
        self.outputProjDims = outputProjDims
        self.subsamplingConvChannels = subsamplingConvChannels
        self.residualWeight = residualWeight
    }
    
    public var headDim: Int {
        hiddenSize / numAttentionHeads
    }
}