import copy
import json
import math
import os

import deepobs
from deepobs.tuner import RandomSearch
import tensorflow as tf

import utils


def generate_commands_script(
    tuner,
    testproblem,
    optimizer_name,
    run_script,
    output_dir,
    generation_dir,
    num_epochs,
    batch_size,
    random_seed=42,
):

    os.makedirs(os.path.join(generation_dir, testproblem), exist_ok=True)
    file_path = os.path.join(
        generation_dir, testproblem, f"jobs_{optimizer_name}_{tuner._search_name}.json"
    )

    if os.path.exists(file_path):
        print(f"Error: Command script already exists at {file_path}.")
        print("Please delete it if you would like to overwrite it.")
        return

    tuner._set_seed(random_seed)
    params = tuner._sample()

    contents = [
        {
            "optimizer_name": optimizer_name,
            "runner_params": {
                "testproblem": testproblem,
                "hyperparams": sample,
                "random_seed": random_seed,
                "outut_dir": output_dir,
                "num_epochs": num_epochs,
                "batch_size": batch_size,
            },
        }
        for sample in params
    ]
    with open(file_path, "w") as f:
        json.dump(contents, f, indent=4)


def main():
    opt = "rgd"
    problem = "fmnist_vae"

    tuner = RandomSearch(
        utils.optimizers[opt],
        utils.hyperparams[opt],
        utils.hp_distributions[opt],
        runner=deepobs.tensorflow.runners.StandardRunner,
        ressources=25,
    )
    generate_commands_script(
        tuner,
        problem,
        opt,
        run_script="./runner.py",
        num_epochs=utils.problems[problem]["num_epochs"],
        batch_size=utils.problems[problem]["batch_size"],
        output_dir="./results",
        generation_dir="./command_scripts",
    )


if __name__ == "__main__":
    main()
