# Autogenerated code. DO NOT EDIT.
import py_gen_ml as pgml
import typing

import pydantic
import typer

from . import config_base as base


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

    learning_rate: typing.Annotated[typing.Optional[float], typer.Option(help="Learning rate. Maps to 'learning_rate'"), pydantic.Field(None), pgml.ArgRef("learning_rate")]
    """Learning rate"""

    beta1: typing.Annotated[typing.Optional[float], typer.Option(help="Decay rate. Maps to 'beta1'"), pydantic.Field(None), pgml.ArgRef("beta1")]
    """Decay rate"""



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

    batch_size: typing.Annotated[typing.Optional[int], typer.Option(help="Batch size for a single GPU. Maps to 'batch_size'"), pydantic.Field(None), pgml.ArgRef("batch_size")]
    """Batch size for a single GPU"""

    num_epochs: typing.Annotated[typing.Optional[int], typer.Option(help="Number of epochs to train. Maps to 'num_epochs'"), pydantic.Field(None), pgml.ArgRef("num_epochs")]
    """Number of epochs to train"""



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

    out_channels: typing.Annotated[typing.Optional[int], typer.Option(help="Number of output channels. Maps to 'out_channels'"), pydantic.Field(None), pgml.ArgRef("out_channels")]
    """Number of output channels"""

    kernel_size: typing.Annotated[typing.Optional[int], typer.Option(help="Square kernel size. Maps to 'kernel_size'"), pydantic.Field(None), pgml.ArgRef("kernel_size")]
    """Square kernel size"""

    pool_size: typing.Annotated[typing.Optional[int], typer.Option(help="Square pool size. Maps to 'pool_size'"), pydantic.Field(None), pgml.ArgRef("pool_size")]
    """Square pool size"""

    activation: typing.Annotated[typing.Optional[base.Activation], typer.Option(help="Activation function. Maps to 'activation'"), pydantic.Field(None), pgml.ArgRef("activation")]
    """Activation function"""



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

    out_features: typing.Annotated[typing.Optional[int], typer.Option(help="Number of output features. Maps to 'out_features'"), pydantic.Field(None), pgml.ArgRef("out_features")]
    """Number of output features"""

    activation: typing.Annotated[typing.Optional[base.Activation], typer.Option(help="Activation function. Maps to 'activation'"), pydantic.Field(None), pgml.ArgRef("activation")]
    """Activation function"""



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

    num_layers: typing.Annotated[typing.Optional[int], typer.Option(help="Number of layers. Maps to 'num_layers'"), pydantic.Field(None), pgml.ArgRef("num_layers")]
    """Number of layers"""

    out_channels: typing.Annotated[typing.Optional[int], typer.Option(help="Number of output channels. Maps to 'block.out_channels'"), pydantic.Field(None), pgml.ArgRef("block.out_channels")]
    """Number of output channels"""

    kernel_size: typing.Annotated[typing.Optional[int], typer.Option(help="Square kernel size. Maps to 'block.kernel_size'"), pydantic.Field(None), pgml.ArgRef("block.kernel_size")]
    """Square kernel size"""

    pool_size: typing.Annotated[typing.Optional[int], typer.Option(help="Square pool size. Maps to 'block.pool_size'"), pydantic.Field(None), pgml.ArgRef("block.pool_size")]
    """Square pool size"""

    activation: typing.Annotated[typing.Optional[base.Activation], typer.Option(help="Activation function. Maps to 'block.activation'"), pydantic.Field(None), pgml.ArgRef("block.activation")]
    """Activation function"""



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

    num_layers: typing.Annotated[typing.Optional[int], typer.Option(help="Number of layers. Maps to 'num_layers'"), pydantic.Field(None), pgml.ArgRef("num_layers")]
    """Number of layers"""

    out_features: typing.Annotated[typing.Optional[int], typer.Option(help="Number of output features. Maps to 'block.out_features'"), pydantic.Field(None), pgml.ArgRef("block.out_features")]
    """Number of output features"""

    activation: typing.Annotated[typing.Optional[base.Activation], typer.Option(help="Activation function. Maps to 'block.activation'"), pydantic.Field(None), pgml.ArgRef("block.activation")]
    """Activation function"""



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

    out_channels: typing.Annotated[typing.Optional[int], typer.Option(help="Number of output channels. Maps to 'conv_net.block.out_channels'"), pydantic.Field(None), pgml.ArgRef("conv_net.block.out_channels")]
    """Number of output channels"""

    kernel_size: typing.Annotated[typing.Optional[int], typer.Option(help="Square kernel size. Maps to 'conv_net.block.kernel_size'"), pydantic.Field(None), pgml.ArgRef("conv_net.block.kernel_size")]
    """Square kernel size"""

    pool_size: typing.Annotated[typing.Optional[int], typer.Option(help="Square pool size. Maps to 'conv_net.block.pool_size'"), pydantic.Field(None), pgml.ArgRef("conv_net.block.pool_size")]
    """Square pool size"""

    out_features: typing.Annotated[typing.Optional[int], typer.Option(help="Number of output features. Maps to 'head.block.out_features'"), pydantic.Field(None), pgml.ArgRef("head.block.out_features")]
    """Number of output features"""



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

    out_channels: typing.Annotated[typing.Optional[int], typer.Option(help="Number of output channels. Maps to 'net.conv_net.block.out_channels'"), pydantic.Field(None), pgml.ArgRef("net.conv_net.block.out_channels")]
    """Number of output channels"""

    kernel_size: typing.Annotated[typing.Optional[int], typer.Option(help="Square kernel size. Maps to 'net.conv_net.block.kernel_size'"), pydantic.Field(None), pgml.ArgRef("net.conv_net.block.kernel_size")]
    """Square kernel size"""

    pool_size: typing.Annotated[typing.Optional[int], typer.Option(help="Square pool size. Maps to 'net.conv_net.block.pool_size'"), pydantic.Field(None), pgml.ArgRef("net.conv_net.block.pool_size")]
    """Square pool size"""

    out_features: typing.Annotated[typing.Optional[int], typer.Option(help="Number of output features. Maps to 'net.head.block.out_features'"), pydantic.Field(None), pgml.ArgRef("net.head.block.out_features")]
    """Number of output features"""

    batch_size: typing.Annotated[typing.Optional[int], typer.Option(help="Batch size for a single GPU. Maps to 'data.batch_size'"), pydantic.Field(None), pgml.ArgRef("data.batch_size")]
    """Batch size for a single GPU"""

    num_epochs: typing.Annotated[typing.Optional[int], typer.Option(help="Number of epochs to train. Maps to 'data.num_epochs'"), pydantic.Field(None), pgml.ArgRef("data.num_epochs")]
    """Number of epochs to train"""

    learning_rate: typing.Annotated[typing.Optional[float], typer.Option(help="Learning rate. Maps to 'optimizer.learning_rate'"), pydantic.Field(None), pgml.ArgRef("optimizer.learning_rate")]
    """Learning rate"""

    beta1: typing.Annotated[typing.Optional[float], typer.Option(help="Decay rate. Maps to 'optimizer.beta1'"), pydantic.Field(None), pgml.ArgRef("optimizer.beta1")]
    """Decay rate"""

    conv_activation: typing.Annotated[typing.Optional[base.Activation], typer.Option(help="Activation function. Maps to 'net.conv_net.block.activation'"), pydantic.Field(None), pgml.ArgRef("net.conv_net.block.activation")]
    """Activation function"""

    head_activation: typing.Annotated[typing.Optional[base.Activation], typer.Option(help="Activation function. Maps to 'net.head.block.activation'"), pydantic.Field(None), pgml.ArgRef("net.head.block.activation")]
    """Activation function"""

    num_conv_layers: typing.Annotated[typing.Optional[int], typer.Option(help="Number of layers. Maps to 'net.conv_net.num_layers'"), pydantic.Field(None), pgml.ArgRef("net.conv_net.num_layers")]
    """Number of layers"""

    num_mlp_layers: typing.Annotated[typing.Optional[int], typer.Option(help="Number of layers. Maps to 'net.head.num_layers'"), pydantic.Field(None), pgml.ArgRef("net.head.num_layers")]
    """Number of layers"""


