from __future__ import annotations

import torch
from torch import nn

try:
    from .config import ExperimentConfig, ModelConfig
except ImportError:
    from config import ExperimentConfig, ModelConfig


class FeatureMLP(nn.Module):
    def __init__(self, config: ModelConfig) -> None:
        super().__init__()
        layers: list[nn.Module] = []
        input_dim = config.input_dim
        for hidden_dim in config.hidden_dims:
            layers.append(nn.Linear(input_dim, hidden_dim))
            layers.append(nn.GELU())
            layers.append(nn.Dropout(config.dropout))
            input_dim = hidden_dim
        layers.append(nn.Linear(input_dim, 1))
        self.network = nn.Sequential(*layers)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.network(x)


def build_model(config: ExperimentConfig | ModelConfig) -> FeatureMLP:
    model_config = config.model if isinstance(config, ExperimentConfig) else config
    return FeatureMLP(model_config)