"""
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 AutoMLConfig
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: AutoMLConfig,
) -> nn.Module:
    """Build a fully-connected network from ``config.layers`` and
    ``config.activations``.

    Example
    -------
    >>> cfg = AutoMLConfig(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: AutoMLConfig | 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
