"""

automl_lib.architecture

~~~~~~~~~~~~~~~~~~~~~~~

Dynamic neural-network builder from a user-specified layer config.
"""

from __future__ import annotations


import torch.nn as nn

from .config import trainiqConfig
from .utils import get_logger

logger = get_logger(__name__)


_ACTIVATION_MAP = {

    "relu": nn.ReLU,

    "gelu": nn.GELU,

    "silu": nn.SiLU,

    "tanh": nn.Tanh,

    "leaky_relu": nn.LeakyReLU,

    "elu": nn.ELU,

    "sigmoid": nn.Sigmoid,

}



def build_custom_mlp(

    input_dim: int,

    output_dim: int,

    config: trainiqConfig,

) -> nn.Module:

    """Build a fully-connected network from ``config.layers`` and

    ``config.activations``.


    Example

    -------

    >>> cfg = trainiqConfig(layers=[512, 256, 128], activations="gelu", dropout=0.4)

    >>> model = build_custom_mlp(100, 10, cfg)
    """

    act_cls = _ACTIVATION_MAP.get(config.activations.lower(), nn.ReLU)

    modules: list[nn.Module] = []

    in_features = input_dim

    for hidden in config.layers:

        modules.extend([

            nn.Linear(in_features, hidden),

            nn.BatchNorm1d(hidden),
            act_cls(),

            nn.Dropout(config.dropout),

        ])

        in_features = hidden


    modules.append(nn.Linear(in_features, output_dim))

    model = nn.Sequential(*modules)

    total_params = sum(p.numel() for p in model.parameters())
    logger.info(

        f"Custom MLP: {config.layers} → {output_dim} | "

        f"activation={config.activations} | params={total_params:,}"
    )
    return model



def build_custom_cnn(

    num_classes: int,

    in_channels: int = 3,

    config: trainiqConfig | None = None,

) -> nn.Module:

    """Build a lightweight CNN for image classification.


    Uses config.layers to define channel sizes of conv blocks.
    """

    channels = config.layers if config and config.layers else [32, 64, 128]

    act_cls = _ACTIVATION_MAP.get(

        config.activations.lower() if config else "relu", nn.ReLU
    )

    dropout = config.dropout if config else 0.3


    conv_blocks: list[nn.Module] = []

    in_ch = in_channels

    for out_ch in channels:

        conv_blocks.extend([

            nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1),

            nn.BatchNorm2d(out_ch),
            act_cls(),

            nn.MaxPool2d(2),

        ])

        in_ch = out_ch


    model = nn.Sequential(

        *conv_blocks,

        nn.AdaptiveAvgPool2d(1),

        nn.Flatten(),

        nn.Dropout(dropout),

        nn.Linear(channels[-1], num_classes),
    )

    total_params = sum(p.numel() for p in model.parameters())

    logger.info(f"Custom CNN: channels={channels} | params={total_params:,}")
    return model

