import jax
import random
import argparse

from rlhf_agents.train import train_rlhf_agent


def jax_debug_wrapper(f):
    def wrapped_fn(*x):
        jax.config.update("jax_debug_nans", True)
        if False:
            with jax.disable_jit():
                return f(*x)
        else:
            return f(*x)

    return wrapped_fn


arg_parser = argparse.ArgumentParser()
arg_parser.add_argument("--env_name", type=str, default="ant")
arg_parser.add_argument("--data_type", type=str, default="basic")
arg_parser.add_argument("--reward_agent1", type=str, default="5100")
arg_parser.add_argument("--reward_agent2", type=str, default="5100")
arg_parser.add_argument("--num_data_points", type=int, default=1280)
arg_parser.add_argument(
    "--dataset_path",
    type=str,
)
arg_parser.add_argument("--loss_type", type=str, default="orpo")
arg_parser.add_argument("--tracking", type=int, default=0)

arg_parser.add_argument("--seed", type=int, default=random.randint(0, 100000))
arg_parser.add_argument("--num_eval_agents", type=int, default=100)
arg_parser.add_argument("--nolog", action="store_true")
arg_parser.add_argument("--wandb_project", type=str)
arg_parser.add_argument("--wandb_entity", type=str)
arg_parser.add_argument("--wandb_group", type=str)
arg_parser.add_argument("--not_wandb_tuning", action="store_true")
arg_parser.add_argument("--reference_agent", type=str, default=None)
arg_parser.add_argument("--temporally_aware", type=int, default=0)
arg_parser.add_argument("--map_location", type=str, default=None)
arg_parser.add_argument("--random_init", type=int, default=0)
arg_parser.add_argument("--indeces", type=int, nargs="+", default=None)
arg_parser.add_argument(
    "--main_folder_path", type=str
)
arg_parser.add_argument("--update_epochs_multiplier", type=float, default=1.0)
arg_parser.add_argument("--parametrised_reward_model", type=int, default=0)
arg_parser.add_argument("--add_logsimoid_bias", type=int, default=0)


# Prevent a wandb bug
arg_parser.add_argument("--LR", type=float)
arg_parser.add_argument("--LR_END", type=float)
arg_parser.add_argument("--UPDATE_EPOCHS", type=int)
arg_parser.add_argument("--ANNEAL_LR", type=bool)
arg_parser.add_argument("--MAX_GRAD_NORM", type=float)
arg_parser.add_argument("--BETA", type=float)
arg_parser.add_argument("--ALPHA", type=float)
arg_parser.add_argument("--ALPHA1", type=float)
arg_parser.add_argument("--ALPHA2", type=float)
arg_parser.add_argument("--GAMMA_BETA_RATIO", type=float)
arg_parser.add_argument("--MINIBATCH_SIZE", type=int)


args = arg_parser.parse_args()
experiment_fn = jax_debug_wrapper(train_rlhf_agent)


if __name__ == "__main__":
    experiment_fn(args)
