"""Train an algorithm."""
import argparse
import json
import os
from harl.utils.configs_tools import get_defaults_yaml_args, get_task_name, update_args
import wandb
import glob

def main():
    """Main function."""
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument(
        "--algo",
        type=str,
        default="happo",
        choices=[
            "happo",
            "hatrpo",
            "haa2c",
            "haddpg",
            "hatd3",
            "hasac",
            "had3qn",
            "maddpg",
            "matd3",
            "mappo",
        ],
        help="Algorithm name. Choose from: happo, hatrpo, haa2c, haddpg, hatd3, hasac, had3qn, maddpg, matd3, mappo.",
    )
    parser.add_argument(
        "--env",
        type=str,
        default="pettingzoo_mpe",
        choices=[
            "smac",
            "mamujoco",
            "pettingzoo_mpe",
            "gym",
            "football",
            "dexhands",
            "smacv2",
            "lag",
        ],
        help="Environment name. Choose from: smac, mamujoco, pettingzoo_mpe, gym, football, dexhands, smacv2, lag.",
    )
    parser.add_argument(
        "--exp_name", type=str, default="installtest", help="Experiment name."
    )
    parser.add_argument(
        "--load_config",
        type=str,
        default="",
        help="If set, load existing experiment config file instead of reading from yaml config file.",
    )
    args, unparsed_args = parser.parse_known_args()
    
    # Generate a descriptive name for the experiment
    seed = getattr(args, "seed", None) or (int(unparsed_args[unparsed_args.index("--seed") + 1]) if "--seed" in unparsed_args else None)
    wandb_exp_name = f"{args.exp_name}-{args.algo}-{args.env}-{seed}"
    tags= [args.algo, args.env, "icml-baselines",args.exp_name, "neurips","ablations"]
    group = f"{args.exp_name}-{args.algo}-{args.env}"

    # Start a wandb run with `sync_tensorboard=True` and the descriptive name
    wandb.init(project="happo-comparison", sync_tensorboard=True, name=wandb_exp_name, tags=tags, group=group)

    def process(arg):
        try:
            return eval(arg)
        except:
            print("error in processing argument: ", arg)
            return arg

    keys = [k[2:] for k in unparsed_args[0::2]]  # remove -- from argument
    values = [process(v) for v in unparsed_args[1::2]]
    unparsed_dict = {k: v for k, v in zip(keys, values)}
    args = vars(args)  # convert to dict
    if args["load_config"] != "":  # load config from existing config file
        with open(args["load_config"], encoding="utf-8") as file:
            all_config = json.load(file)
        args["algo"] = all_config["main_args"]["algo"]
        args["env"] = all_config["main_args"]["env"]
        algo_args = all_config["algo_args"]
        env_args = all_config["env_args"]
        
    else:  # load config from corresponding yaml file
        algo_args, env_args = get_defaults_yaml_args(args["algo"], args["env"])
    update_args(unparsed_dict, algo_args, env_args)  # update args from command line

    if args["env"] == "dexhands":
        import isaacgym  # isaacgym has to be imported before PyTorch

    # note: isaac gym does not support multiple instances, thus cannot eval separately
    if args["env"] == "dexhands":
        algo_args["eval"]["use_eval"] = False
        algo_args["train"]["episode_length"] = env_args["hands_episode_length"]

     # Log algo_args and env_args to WandB
    wandb.config.update({"algo_args": algo_args, "env_args": env_args, "EXP_TAGS": tags})
    print(f"Logged algo_args and env_args to WandB.")
    
    # start training
    from harl.runners import RUNNER_REGISTRY

    runner = RUNNER_REGISTRY[args["algo"]](args, algo_args, env_args)
    runner.run()
    runner.close()
    
    # save models 
    from pathlib import Path
    models_folder = Path(runner.run_dir) / "models"
    absolute_models_folder = models_folder.resolve()
    print(f"Absolute path to models folder: {absolute_models_folder}")

    # Check if the folder exists
    if not absolute_models_folder.exists():
        print(f"Models folder does not exist: {absolute_models_folder}")
    else:
        # Log models using W&B artifacts (preferred for managing multiple files)
        artifact = wandb.Artifact('model-weights', type='model')
        artifact.add_dir(str(absolute_models_folder))  # Add the entire directory to the artifact
        wandb.log_artifact(artifact)
        print(f"Uploaded models from {absolute_models_folder} to W&B.")
    wandb.finish()


if __name__ == "__main__":
    main()
