# Autogenerated code. DO NOT EDIT.
import enum

import py_gen_ml as pgml

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    import cifar10.modules


class Activation(str, enum.Enum):
    """Activation function"""

    GELU = "GELU"
    """GELU activation"""

    RELU = "RELU"
    """ReLU activation"""


class Optimizer(pgml.YamlBaseModel):
    """Optimizer configuration"""

    learning_rate: float = 0.0001
    """Learning rate"""

    beta1: float = 0.99
    """Decay rate"""


class Data(pgml.YamlBaseModel):
    """Data configuration"""

    batch_size: int = 32
    """Batch size for a single GPU"""

    num_epochs: int = 10
    """Number of epochs to train"""


class ConvBlock(pgml.YamlBaseModel):
    """Convolutional layer configuration"""

    out_channels: int = 128
    """Number of output channels"""

    kernel_size: int = 3
    """Square kernel size"""

    pool_size: int = 2
    """Square pool size"""

    activation: Activation = Activation.GELU
    """Activation function"""

    def build(self) -> "cifar10.modules.ConvBlock":
        import cifar10.modules

        return cifar10.modules.ConvBlock(
            out_channels=self.out_channels,
            kernel_size=self.kernel_size,
            pool_size=self.pool_size,
            activation=self.activation,
        )


class LinearBlock(pgml.YamlBaseModel):
    """Linear layer configuration"""

    out_features: int = 128
    """Number of output features"""

    activation: Activation = Activation.GELU
    """Activation function"""

    def build(self) -> "cifar10.modules.LinearBlock":
        import cifar10.modules

        return cifar10.modules.LinearBlock(
            out_features=self.out_features,
            activation=self.activation,
        )


class ConvNet(pgml.YamlBaseModel):
    """Convolutional neural network configuration"""

    block: ConvBlock
    """Conv layer configuration"""

    num_layers: int = 2
    """Number of layers"""


class MLP(pgml.YamlBaseModel):
    """Multi-layer perceptron configuration"""

    block: LinearBlock
    """Linear layer configuration"""

    num_layers: int = 2
    """Number of layers"""


class Model(pgml.YamlBaseModel):
    """Model configuration"""

    conv_net: ConvNet
    """Conv blocks"""

    head: MLP
    """MLP head"""


class Project(pgml.YamlBaseModel):
    """Global configuration"""

    net: Model
    """Model configuration"""

    optimizer: Optimizer
    """Optimizer configuration"""

    data: Data
    """Data configuration"""
