import argparse
from dataset_gen.dataset_gen import gen_dataset
import os
import jax.numpy as jnp
import jax.random as jr

if __name__ == "__main__":
    arg_parser = argparse.ArgumentParser()
    arg_parser.add_argument("--env_name", type=str, default="hopper")
    arg_parser.add_argument("--num_data_points", type=int, default=1000)
    arg_parser.add_argument("--seed", type=int, default=0)
    arg_parser.add_argument(
        "--save_path", type=str
    )
    arg_parser.add_argument("--agent1", type=str)
    arg_parser.add_argument("--agent2", type=str)
    arg_parser.add_argument("--noise", type=float, default=0.0)
    arg_parser.add_argument("--shuffle_agents", type=int, default=0)
    arg_parser.add_argument("--judge_temp", type=float)
    arg_parser.add_argument("--penalty_indeces", type=int, nargs="+", default=None)
    arg_parser.add_argument("--tot_reward", type=str, default="sum", choices=["sum", "max"])
    args = arg_parser.parse_args()
    preferences = gen_dataset(jr.PRNGKey(args.seed), args, indeces=args.penalty_indeces)

    os.makedirs(os.path.dirname(args.save_path), exist_ok=True)
    with open(args.save_path, "wb") as f:
        jnp.save(f, preferences)
