ACTIVATIONS = {
    "relu": "nn.ReLU()",
    "tanh": "nn.Tanh()",
    "sigmoid": "nn.Sigmoid()",
    "leaky_relu": "nn.LeakyReLU()",
}


def build_mlp(input_dim, hidden, output_dim, activation="relu"):
    act = ACTIVATIONS.get(activation, "nn.ReLU()")
    layers = []
    prev = input_dim
    for h in hidden:
        layers.append(f"            nn.Linear({prev}, {h}),")
        layers.append(f"            {act},")
        prev = h
    layers.append(f"            nn.Linear({prev}, {output_dim}),")
    return "\n".join(layers)
