import argparse
import os
import shutil

import yaml

while "Algorithms" in os.getcwd():
    os.chdir("..")
import sys
sys.path.append(os.getcwd())


def seed_everything(seed: int):
    import random, os
    import numpy as np
    import torch

    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = True

def run_main_general(args):
    seed_pool = [6, 36, 216, 1296, 7776, 46656]
    cfg_yaml=f"Algorithms/{args.algo}/cfg/{args.cfg}.yaml"
    with open(cfg_yaml,"r") as f:
        cfg_dict=yaml.load(f,Loader=yaml.FullLoader)

    if args.algo =="Slav3":
        from Algorithms.Slav3.config_slav3 import Slav3Config as Config_
        from Algorithms.Slav3.main_slav3 import main_slav3 as main
    elif args.algo =="RND":
        from Algorithms.RND.config_rnd import RNDConfig as Config_
        from Algorithms.RND.main_rnd import main_rnd as main
    elif args.algo =="ETD":
        from Algorithms.ETD.config_etd import ETDConfig as Config_
        from Algorithms.ETD.main_etd import main_etd as main
    elif args.algo =="RIDE":
        from Algorithms.RIDE.config_ride import RIDEConfig as Config_
        from Algorithms.RIDE.main_ride import main_ride as main
    elif args.algo =="NovelD":
        from Algorithms.NovelD.config_noveld import NovelDConfig as Config_
        from Algorithms.NovelD.main_noveld import main_noveld as main
    else:
        raise NotImplementedError(f"Algorithm {args.algo} not implemented.")
    idx = 0 if args.idx is None else args.idx
    if idx!=-1:
        cfg_dict.update({
            "seed": seed_pool[idx],
        })

    if args.clean:
        cfg_save_name=args.cfg.replace("/","_")
        cfg_dict.update({
            "log_dir":f"{args.algo}_{cfg_save_name}_{idx}",
        })

    config= Config_(cfg_dict)
    if args.chdir is not None:
        os.makedirs(f"{args.chdir}/{args.algo}",exist_ok=True)
        os.chdir(f"{args.chdir}/{args.algo}")

    if args.clean:

        if os.path.exists(os.path.join(os.getcwd(), "Logs", config.log_dir)):
            print(f"{config.log_dir} already exists, removing.")
            shutil.rmtree(os.path.join(os.getcwd(), "Logs", config.log_dir))
        if os.path.exists(os.path.join(os.getcwd(), "Models", config.log_dir)):
            shutil.rmtree(os.path.join(os.getcwd(), "Models", config.log_dir))
    if hasattr(config,"seed") and config.seed is not None:
        seed_everything(config.seed)
    main(config)


def parse_args():
    args = argparse.ArgumentParser()
    args.add_argument('--algo', type=str, default="Slav3")
    args.add_argument(
        "--cfg",
        type=str,
        default="AT_PF")
    args.add_argument("--idx",type=int,default=None)
    args.add_argument("--chdir", type=str, default="Runs")
    args.add_argument("--clean", action="store_false")
    args.add_argument("--record_video",  action="store_true")
    args = args.parse_args()
    return args

if __name__ == "__main__":
    args=parse_args()
    run_main_general(args)
