import argparse


def get_args(command=None):
    parser = argparse.ArgumentParser()
    # World generator args
    parser.add_argument(
        "--num_rel",
        default=4,
        type=int,
        help="number of relations for the current rule",
    )
    parser.add_argument(
        "--per_inverse", default=0.5, type=float, help="percentage of inverse relations"
    )
    parser.add_argument(
        "--corrupt_eps",
        default=0.1,
        type=float,
        help="corruption epsilon greedy percentage for each rule",
    )
    parser.add_argument(
        "--rule_name", default="", type=str, help="placeholder for rule name"
    )
    parser.add_argument(
        "--uniform_prob",
        default=False,
        action="store_true",
        help="make all rules with uniform probability",
    )
    parser.add_argument(
        "--fix_num_relations",
        default=False,
        action="store_true",
        help="if true, then create the first rule world, and subsequently only allow the heads of the first rule world in the rest",
    )
    parser.add_argument(
        "--policy", default="flip", type=str, help="random/flip/sanity/fsrl_1"
    )
    parser.add_argument(
        "--num_splits",
        default=4,
        type=int,
        help="number of splits for the current config",
    )
    # Graph Generator args
    parser.add_argument("--num_nodes", default=50, type=int, help="number of nodes")
    parser.add_argument(
        "--expand_steps",
        default=5,
        type=int,
        help="max number of steps expansion takes place",
    )
    parser.add_argument(
        "--add_noise", default=False, action="store_true", help="add noise"
    )
    parser.add_argument(
        "--save_path",
        default="mlp/lgw/data/",
        type=str,
        help="save path (to be prepended by $HOME/mlp)",
    )
    parser.add_argument(
        "--sanity",
        default=False,
        action="store_true",
        help="if true, generate graphs from the same rule world",
    )
    parser.add_argument(
        "--bidirectional",
        default=False,
        action="store_true",
        help="if true, then the main resolution path contain bidirectional edges for the no-noise case",
    )
    parser.add_argument(
        "--gen_graph_cyles",
        default=10,
        type=int,
        help="cycles in the to generate in the big graph",
    )
    parser.add_argument("--randomize_steps", action="store_true")
    parser.add_argument(
        "--world_graph_expand_steps",
        default=10,
        type=int,
        help="max world graph expansions",
    )
    parser.add_argument(
        "--world_graph_per_edges",
        default=1.0,
        type=float,
        help="world graph percentage edges",
    )
    parser.add_argument("--sample_worlds", default=False, action="store_true", help="")
    parser.add_argument("--sample_graphs", default=False, action="store_true", help="")
    parser.add_argument(
        "--num_train_rows", default=5000, type=int, help="number of train"
    )
    parser.add_argument(
        "--num_valid_rows", default=1000, type=int, help="number of valid"
    )
    parser.add_argument(
        "--num_test_rows", default=1000, type=int, help="number of test"
    )
    parser.add_argument(
        "--easy_mode",
        default=False,
        action="store_true",
        help="draw test graphs from train distribution",
    )
    parser.add_argument(
        "--path_cutoff", default=5, type=int, help="max length of the resolution path"
    )

    # Meta args
    parser.add_argument("--num_worlds", default=3, type=int, help="number of worlds")
    parser.add_argument(
        "--graphs_per_world", default=300, type=int, help="number of graphs per world"
    )
    parser.add_argument(
        "--train_test_split", default=0.9, type=float, help="train test split"
    )
    parser.add_argument(
        "--train_val_split", default=0.9, type=float, help="train val split"
    )
    parser.add_argument(
        "--world_train_val_test_split",
        default=0.8,
        type=float,
        help="train val test split for worlds",
    )
    parser.add_argument(
        "--num_rel_choices", default="8", type=str, help="number of relations per world"
    )
    parser.add_argument(
        "--per_inverse_choices",
        default="0.3,0.5,0.7",
        type=str,
        help="comma separated `per_world`",
    )
    parser.add_argument(
        "--corrupt_eps_choices",
        default="0,0.1,0.2,0.5",
        type=str,
        help="comma separated `corrupt_eps`",
    )
    parser.add_argument(
        "--expand_steps_choices",
        default="3,4,5,6,7",
        type=str,
        help="Expand step choices",
    )
    parser.add_argument("--folder_name", default="data", type=str)
    # parser.add_argument("--exp_name", default="simple", type=str, help="experiment name of the data")
    parser.add_argument(
        "--load_rule", default="", type=str, help="load rule world for sanity testing"
    )

    # Eval args (for scripts)
    parser.add_argument(
        "--config_path", type=str, default="supervised/supervised_exp_icml"
    )
    parser.add_argument("--eval_k_shot", type=int, default=0)
    parser.add_argument("--eval_k_epoch", type=int, default=-1)
    parser.add_argument("--output", type=str, default="eval_results.csv")
    parser.add_argument(
        "--config_toggle_true",
        type=str,
        default="",
        help="comma separated flag to toggle true in model",
    )
    parser.add_argument("--eval_data_mode", type=str, default="hard")
    parser.add_argument("--eval_load_epoch", type=int, default=-1)
    parser.add_argument("--eval_rules", type=str, help="comma separated", default="")
    parser.add_argument("--eval_data_folder", type=str, default="train", help="")
    parser.add_argument("--eval_store_rep", action="store_true", default=False)

    if command:
        return parser.parse_args(command.split(" "))
    else:
        return parser.parse_args()
