"""A minimal MLP with tanh activations, packaged as a HuggingFace custom model.

Self-contained on purpose: only imports from torch / transformers so that it can
be downloaded and executed by `AutoModel.from_pretrained(..., trust_remote_code=True)`
without any project-local dependency.
"""

import torch
import torch.nn as nn
from transformers import PreTrainedModel
from transformers.modeling_outputs import BaseModelOutput

from .configuration_mlp import MLPConfig


class MLPModel(PreTrainedModel):
    # `config_class` wires the model into the AutoModel registry.
    config_class = MLPConfig

    def __init__(self, config: MLPConfig):
        super().__init__(config)

        dims = (
            [config.input_dim]
            + [config.hidden_dim] * config.num_hidden_layers
            + [config.output_dim]
        )
        layers = []
        for i in range(len(dims) - 1):
            layers.append(nn.Linear(dims[i], dims[i + 1]))
            # tanh after every layer except the final projection.
            if i < len(dims) - 2:
                layers.append(nn.Tanh())
        self.mlp = nn.Sequential(*layers)

        # Standard HF weight init hook.
        self.post_init()

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=0.02)
            if module.bias is not None:
                module.bias.data.zero_()

    def forward(self, x: torch.Tensor = None, **kwargs) -> BaseModelOutput:
        # Accept `input_ids` as an alias so generic HF tooling doesn't choke.
        if x is None:
            x = kwargs.get("input_ids")
        output = self.mlp(x)
        return BaseModelOutput(last_hidden_state=output)