import re
import optuna
import subprocess
import tempfile


def objective(trial):
    # Suggest parameters.
    trial.suggest_int("train.batch_size", 4, 100)
    trial.suggest_loguniform("train.learning_rate", 0.0001, 1.0)

    # Replace config entries with the suggested parameters.
    config = open("config.gin").read()
    for name, value in trial.params.items():
        config = re.sub("(?<=" + name + ") *=.*", "=" + str(value), config)

    # Create a temporary config file.
    temp = tempfile.NamedTemporaryFile(mode="w", encoding="utf-8")
    temp.write(config)
    temp.flush()

    # Run train script with the temporary config.
    result = subprocess.run(
        ["python3", "train.py", "--config-path", temp.name],
        stdout=subprocess.PIPE,
        encoding="utf-8",
    )

    # Parse the script output to get the objective value.
    return float(result.stdout)


study = optuna.create_study()
study.optimize(objective, n_trials=10)