import os
# from datetime import datetime
# from typing import Tuple
import torch
import numpy as np
import random
# from rsl_rl.env import VecEnv
# from roboimi.rsl_rl.runners import OnPolicyRunner
from roboimi.rsl_rl.rsl_rl.runners.on_policy_runner import OnPolicyRunner
from roboimi.rsl_rl.rsl_rl.env import VecEnv

# from legged_gym import LEGGED_GYM_ROOT_DIR, LEGGED_GYM_ENVS_DIR
from roboimi.envs.vs_rl_task.config.vs_robot_config import VSRobotCfg, VSRobotCfgPPO
from .vs_robot import VisualServoDualArmRobot




class VS_TaskRegistry():
    def __init__(self):
        self.task_classes = {}
        self.task_ = None
        self.env_cfgs = {}
        self.env_cfg = None
        self.train_cfgs = {}
        self.train_cfg = None

    # name fixed
    def register(self, name="vs_robot", task_class=VisualServoDualArmRobot, env_cfg=VSRobotCfg, train_cfg=VSRobotCfgPPO):
        self.task_classes[name] = task_class
        self.task_ = task_class
        self.env_cfg = env_cfg
        self.train_cfg = train_cfg
        self.env_cfgs[name] = env_cfg
        self.train_cfgs[name] = train_cfg

    # def get_task_class(self, name: str) -> VecEnv:
    #     return self.task_classes[name]

    def get_task(self):
        return self.task_
    
    def get_cfgs(self, name='vs_robot'):
        # train_cfg = self.train_cfgs[name]
        # env_cfg = self.env_cfgs[name]
        # copy seed
        train_cfg = self.train_cfg
        env_cfg = self.env_cfg
        env_cfg.seed = self.train_cfg.seed
        return env_cfg, train_cfg
    
    def make_vs_env(self,name='vs_robot', device='cuda:0',headless=True):
        task_class = self.get_task()
        env_cfg, _ = self.get_cfgs(name)
        self.set_seed(env_cfg.seed)
        env = task_class(cfg=env_cfg, device=device, headless=headless)
        return env, env_cfg
    
    def make_vs_alg_runner(self, env, name='vs_robot'):
        _, train_cfg = self.get_cfgs(name)
        train_cfg_dict = env.class_to_dict(train_cfg)
        log_dir = None
        runner = OnPolicyRunner(env, train_cfg_dict, log_dir, device='cuda:0')
        resume = train_cfg.runner.resume
        if resume:
            # load previously trained model
            resume_path = self.get_load_path(log_dir, load_run=train_cfg.runner.load_run, checkpoint=train_cfg.runner.checkpoint)
            print(f"Loading model from: {resume_path}")
            runner.load(resume_path)
        return runner, train_cfg
    
    def set_seed(seed):
        if seed == -1:
            seed = np.random.randint(0, 10000)
        print("Setting seed: {}".format(seed))

        random.seed(seed)
        np.random.seed(seed)
        torch.manual_seed(seed)
        os.environ['PYTHONHASHSEED'] = str(seed)
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)

    def get_load_path(root, load_run=-1, checkpoint=-1):
        try:
            runs = os.listdir(root)
            #TODO sort by date to handle change of month
            runs.sort()
            if 'exported' in runs: runs.remove('exported')
            last_run = os.path.join(root, runs[-1])
        except:
            raise ValueError("No runs in this directory: " + root)
        if load_run==-1:
            load_run = last_run
        else:
            load_run = os.path.join(root, load_run)

        if checkpoint==-1:
            models = [file for file in os.listdir(load_run) if 'model' in file]
            models.sort(key=lambda m: '{0:0>15}'.format(m))
            model = models[-1]
        else:
            model = "model_{}.pt".format(checkpoint) 

        load_path = os.path.join(load_run, model)
        return load_path

    # def make_env(self, name='vs_robot', env_cfg=None):
    #     """ Creates an environment either from a registered namme or from the provided config file.

    #     Args:
    #         name (string): Name of a registered env.
    #         args (Args, optional): Isaac Gym comand line arguments. If None get_args() will be called. Defaults to None.
    #         env_cfg (Dict, optional): Environment config file used to override the registered config. Defaults to None.

    #     Raises:
    #         ValueError: Error if no registered env corresponds to 'name' 

    #     Returns:
    #         isaacgym.VecTaskPython: The created environment
    #         Dict: the corresponding config file
    #     """
    #     # if no args passed get command line arguments
    #     # if args is None:
    #     #     pass
    #     #     # args = self.get_args()
    #     # # check if there is a registered env with that name
    #     # if name in self.task_classes:
    #     #     # task_class = self.get_task_class(name)
    #     #     task_class = self.get_task()
    #     # else:
    #     #     raise ValueError(f"Task with name: {name} was not registered")
    #     task_class = self.get_task()
    #     # if env_cfg is None:
    #         # load config files
    #     env_cfg, _ = self.get_cfgs(name)
    #     # override cfg from args (if specified)
    #     # env_cfg, _ = update_cfg_from_args(env_cfg, None, args)
    #     self.set_seed(env_cfg.seed)
    #     #task_class = VisualServoDualArmRobot
    #     env = task_class(cfg=env_cfg,device='cuda:0',headless=True)
    #     return env, env_cfg
    

    # def make_alg_runner(self, env, name=None, args=None, train_cfg=None, log_root="default") -> Tuple[OnPolicyRunner, LeggedRobotCfgPPO]:
    #     """ Creates the training algorithm  either from a registered namme or from the provided config file.

    #     Args:
    #         env (isaacgym.VecTaskPython): The environment to train (TODO: remove from within the algorithm)
    #         name (string, optional): Name of a registered env. If None, the config file will be used instead. Defaults to None.
    #         args (Args, optional): Isaac Gym comand line arguments. If None get_args() will be called. Defaults to None.
    #         train_cfg (Dict, optional): Training config file. If None 'name' will be used to get the config file. Defaults to None.
    #         log_root (str, optional): Logging directory for Tensorboard. Set to 'None' to avoid logging (at test time for example). 
    #                                 Logs will be saved in <log_root>/<date_time>_<run_name>. Defaults to "default"=<path_to_LEGGED_GYM>/logs/<experiment_name>.

    #     Raises:
    #         ValueError: Error if neither 'name' or 'train_cfg' are provided
    #         Warning: If both 'name' or 'train_cfg' are provided 'name' is ignored

    #     Returns:
    #         PPO: The created algorithm
    #         Dict: the corresponding config file
    #     """
    #     # if no args passed get command line arguments
    #     if args is None:
    #         args = get_args()
    #     # if config files are passed use them, otherwise load from the name
    #     if train_cfg is None:
    #         if name is None:
    #             raise ValueError("Either 'name' or 'train_cfg' must be not None")
    #         # load config files
    #         _, train_cfg = self.get_cfgs(name)
    #     else:
    #         if name is not None:
    #             print(f"'train_cfg' provided -> Ignoring 'name={name}'")
    #     # override cfg from args (if specified)
    #     _, train_cfg = update_cfg_from_args(None, train_cfg, args)

    #     if log_root=="default":
    #         log_root = os.path.join(LEGGED_GYM_ROOT_DIR, 'logs', train_cfg.runner.experiment_name)
    #         log_dir = os.path.join(log_root, datetime.now().strftime('%b%d_%H-%M-%S') + '_' + train_cfg.runner.run_name)
    #     elif log_root is None:
    #         log_dir = None
    #     else:
    #         log_dir = os.path.join(log_root, datetime.now().strftime('%b%d_%H-%M-%S') + '_' + train_cfg.runner.run_name)
        
    #     train_cfg_dict = class_to_dict(train_cfg)
    #     runner = OnPolicyRunner(env, train_cfg_dict, log_dir, device=args.rl_device)
    #     #save resume path before creating a new log_dir
    #     resume = train_cfg.runner.resume
    #     if resume:
    #         # load previously trained model
    #         resume_path = get_load_path(log_root, load_run=train_cfg.runner.load_run, checkpoint=train_cfg.runner.checkpoint)
    #         print(f"Loading model from: {resume_path}")
    #         runner.load(resume_path)
    #     return runner, train_cfg
    



    # def get_args():
    #     custom_parameters = [
    #         {"name": "--task", "type": str, "default": "anymal_c_flat", "help": "Resume training or start testing from a checkpoint. Overrides config file if provided."},
    #         {"name": "--resume", "action": "store_true", "default": False,  "help": "Resume training from a checkpoint"},
    #         {"name": "--experiment_name", "type": str,  "help": "Name of the experiment to run or load. Overrides config file if provided."},
    #         {"name": "--run_name", "type": str,  "help": "Name of the run. Overrides config file if provided."},
    #         {"name": "--load_run", "type": str,  "help": "Name of the run to load when resume=True. If -1: will load the last run. Overrides config file if provided."},
    #         {"name": "--checkpoint", "type": int,  "help": "Saved model checkpoint number. If -1: will load the last checkpoint. Overrides config file if provided."},

    #         {"name": "--headless", "action": "store_true", "default": False, "help": "Force display off at all times"},
    #         {"name": "--horovod", "action": "store_true", "default": False, "help": "Use horovod for multi-gpu training"},
    #         {"name": "--rl_device", "type": str, "default": "cuda:0", "help": 'Device used by the RL algorithm, (cpu, gpu, cuda:0, cuda:1 etc..)'},
    #         {"name": "--num_envs", "type": int, "help": "Number of environments to create. Overrides config file if provided."},
    #         {"name": "--seed", "type": int, "help": "Random seed. Overrides config file if provided."},
    #         {"name": "--max_iterations", "type": int, "help": "Maximum number of training iterations. Overrides config file if provided."},
    #     ]
    #     # parse arguments
    #     args = gymutil.parse_arguments(
    #         description="RL Policy",
    #         custom_parameters=custom_parameters)

    #     # name allignment
    #     args.sim_device_id = args.compute_device_id
    #     args.sim_device = args.sim_device_type
    #     if args.sim_device=='cuda':
    #         args.sim_device += f":{args.sim_device_id}"
    #     return args

    # def update_cfg_from_args(env_cfg, cfg_train, args):
    #     # seed
    #     if env_cfg is not None:
    #         # num envs
    #         if args.num_envs is not None:
    #             env_cfg.env.num_envs = args.num_envs
    #     if cfg_train is not None:
    #         if args.seed is not None:
    #             cfg_train.seed = args.seed
    #         # alg runner parameters
    #         if args.max_iterations is not None:
    #             cfg_train.runner.max_iterations = args.max_iterations
    #         if args.resume:
    #             cfg_train.runner.resume = args.resume
    #         if args.experiment_name is not None:
    #             cfg_train.runner.experiment_name = args.experiment_name
    #         if args.run_name is not None:
    #             cfg_train.runner.run_name = args.run_name
    #         if args.load_run is not None:
    #             cfg_train.runner.load_run = args.load_run
    #         if args.checkpoint is not None:
    #             cfg_train.runner.checkpoint = args.checkpoint

    #     return env_cfg, cfg_train