from rl_games.common.ivecenv import IVecEnv
from rl_games.common.env_configurations import configurations
from rl_games.common.tr_helpers import dicts_to_dict_with_arrays
import numpy as np
import random
from time import sleep
import torch


class RayWorker:
    """Wrapper around a third-party (gym for example) environment class that enables parallel training.

    The RayWorker class wraps around another environment class to enable the use of this 
    environment within an asynchronous parallel training setup

    """
    def __init__(self, config_name, config):
        """Initialise the class. Sets up the environment creator using the `rl_games.common.env_configurations.configuraitons` dict

        Args:
            config_name (:obj:`str`): Key of the environment to create.
            config: Misc. kwargs passed on to the environment creator function

        """
        self.env = configurations[config_name]['env_creator'](**config)

    def _obs_to_fp32(self, obs):
        if isinstance(obs, dict):
            for k, v in obs.items():
                if isinstance(v, dict):
                    for dk, dv in v.items():
                        if dv.dtype == np.float64:
                            v[dk] = dv.astype(np.float32)
                else:
                    if v.dtype == np.float64:
                        obs[k] = v.astype(np.float32)
        else:
            if obs.dtype == np.float64:
                obs = obs.astype(np.float32)
        return obs

    def step(self, action):
        """Step the environment and reset if done

        Args:
            action (type depends on env): Action to take.

        """
        next_state, reward, is_done, info = self.env.step(action)

        if np.isscalar(is_done):
            episode_done = is_done
        else:
            episode_done = is_done.all()
        if episode_done:
            next_state = self.reset()
        next_state = self._obs_to_fp32(next_state)
        return next_state, reward, is_done, info

    def seed(self, seed):
        if hasattr(self.env, 'seed'):
            torch.manual_seed(seed)
            torch.cuda.manual_seed_all(seed)
            np.random.seed(seed)
            random.seed(seed)
            self.env.seed(seed)

    def render(self):
        self.env.render()

    def reset(self):
        obs = self.env.reset()
        obs = self._obs_to_fp32(obs)
        return obs

    def get_action_mask(self):
        return self.env.get_action_mask()

    def get_number_of_agents(self):
        if hasattr(self.env, 'get_number_of_agents'):
            return self.env.get_number_of_agents()
        else:
            return 1

    def set_weights(self, weights):
        self.env.update_weights(weights)

    def close(self):
        if hasattr(self.env, 'close'):
            self.env.close()

    def can_concat_infos(self):
        if hasattr(self.env, 'concat_infos'):
            return self.env.concat_infos
        else:
            return False

    def get_env_info(self):
        info = {}
        observation_space = self.env.observation_space

        #if isinstance(observation_space, gym.spaces.dict.Dict):
        #    observation_space = observation_space['observations']

        info['action_space'] = self.env.action_space
        info['observation_space'] = observation_space
        info['state_space'] = None
        info['use_global_observations'] = False
        info['agents'] = self.get_number_of_agents()
        info['value_size'] = 1
        if hasattr(self.env, 'use_central_value'):
            info['use_global_observations'] = self.env.use_central_value
        if hasattr(self.env, 'value_size'):
            info['value_size'] = self.env.value_size
        if hasattr(self.env, 'state_space'):
            info['state_space'] = self.env.state_space
        return info


class RayVecEnv(IVecEnv):
    """Main env class that manages several `rl_games.common.vecenv.Rayworker` objects for parallel training

    The RayVecEnv class manages a set of individual environments and wraps around the methods from RayWorker.
    Each worker is executed asynchronously.

    """
    def __init__(self, config_name, num_actors, **kwargs):
        """Initialise the class. Sets up the config for the environment and creates individual workers to manage.

        Args:
            config_name (:obj:`str`): Key of the environment to create.
            num_actors (:obj:`int`): Number of environments (actors) to create
            **kwargs: Misc. kwargs passed on to the environment creator function within the RayWorker __init__

        """
        # Import and initialize Ray only when RayVecEnv is actually used
        try:
            import ray
            self.ray = ray
        except ImportError:
            raise ImportError("Ray is required for RayVecEnv. Please install it with: pip install ray")

        # Initialize Ray if not already initialized
        if not ray.is_initialized():
            # Get Ray initialization parameters from config
            # Example usage in YAML under params.config.env_config:
            #   ray_config:
            #     num_cpus: 8                      # Number of CPUs Ray can use
            #     num_gpus: 1                      # Number of GPUs Ray can use  
            #     object_store_memory: 2000000000  # Object store memory in bytes (2GB)
            #     local_mode: True                 # Run Ray in local mode for debugging
            #     ignore_reinit_error: True        # Suppress errors if Ray is already initialized
            #     include_dashboard: False         # Disable Ray dashboard
            #     dashboard_host: '0.0.0.0'        # Ray dashboard host
            ray_config = kwargs.pop('ray_config', {})

            # Set default object store memory if not specified
            if 'object_store_memory' not in ray_config:
                ray_config['object_store_memory'] = 1024*1024*1000

            ray.init(**ray_config)

        self.config_name = config_name
        self.num_actors = num_actors
        self.use_torch = False
        self.seed = kwargs.pop('seed', None)

        self.remote_worker = self.ray.remote(RayWorker)
        self.workers = [self.remote_worker.remote(self.config_name, kwargs) for i in range(self.num_actors)]

        if self.seed is not None:
            seeds = range(self.seed, self.seed + self.num_actors)
            seed_set = []
            for (seed, worker) in zip(seeds, self.workers):	        
                seed_set.append(worker.seed.remote(seed))
            self.ray.get(seed_set)

        res = self.workers[0].get_number_of_agents.remote()
        self.num_agents = self.ray.get(res)

        res = self.workers[0].get_env_info.remote()
        env_info = self.ray.get(res)
        res = self.workers[0].can_concat_infos.remote()
        can_concat_infos = self.ray.get(res)
        self.use_global_obs = env_info['use_global_observations']
        self.concat_infos = can_concat_infos
        obs_space = env_info.get('observation_space')
        state_space = env_info.get('state_space')
        self.obs_type_dict = obs_space is not None and type(obs_space).__name__ == 'Dict'
        self.state_type_dict = state_space is not None and type(state_space).__name__ == 'Dict'
        if self.num_agents == 1:
            self.concat_func = np.stack
        else:
            self.concat_func = np.concatenate

    def step(self, actions):
        """Step all individual environments (using the created workers). 
        Returns a concatenated array of observations, rewards, done states, and infos if the env allows concatenation.
        Else returns a nested dict.

        Args:
            action (type depends on env): Action to take.

        """
        newobs, newstates, newrewards, newdones, newinfos = [], [], [], [], []
        res_obs = []
        if self.num_agents == 1:
            for (action, worker) in zip(actions, self.workers):	        
                res_obs.append(worker.step.remote(action))
        else:
            for num, worker in enumerate(self.workers):
                res_obs.append(worker.step.remote(actions[self.num_agents * num: self.num_agents * num + self.num_agents]))

        all_res = self.ray.get(res_obs)
        for res in all_res:
            cobs, crewards, cdones, cinfos = res
            if self.use_global_obs:
                newobs.append(cobs["obs"])
                newstates.append(cobs["state"])
            else:
                newobs.append(cobs)
            newrewards.append(crewards)
            newdones.append(cdones)
            newinfos.append(cinfos)

        if self.obs_type_dict:
            ret_obs = dicts_to_dict_with_arrays(newobs, self.num_agents == 1)
        else:
            ret_obs = self.concat_func(newobs)

        if self.use_global_obs:
            newobsdict = {}
            newobsdict["obs"] = ret_obs

            if self.state_type_dict:
                newobsdict["states"] = dicts_to_dict_with_arrays(newstates, True)
            else:
                newobsdict["states"] = np.stack(newstates)            
            ret_obs = newobsdict
        if self.concat_infos:
            newinfos = dicts_to_dict_with_arrays(newinfos, False)
        return ret_obs, self.concat_func(newrewards), self.concat_func(newdones), newinfos

    def get_env_info(self):
        res = self.workers[0].get_env_info.remote()
        return self.ray.get(res)

    def set_weights(self, indices, weights):
        res = []
        for ind in indices:
            res.append(self.workers[ind].set_weights.remote(weights))
        self.ray.get(res)

    def has_action_masks(self):
        return True

    def get_action_masks(self):
        mask = [worker.get_action_mask.remote() for worker in self.workers]
        masks = self.ray.get(mask)
        return np.concatenate(masks, axis=0)

    def reset(self):
        res_obs = [worker.reset.remote() for worker in self.workers]
        newobs, newstates = [],[]
        for res in res_obs:
            cobs = self.ray.get(res)
            if self.use_global_obs:
                newobs.append(cobs["obs"])
                newstates.append(cobs["state"])
            else:
                newobs.append(cobs)

        if self.obs_type_dict:
            ret_obs = dicts_to_dict_with_arrays(newobs, self.num_agents == 1)
        else:
            ret_obs = self.concat_func(newobs)

        if self.use_global_obs:
            newobsdict = {}
            newobsdict["obs"] = ret_obs

            if self.state_type_dict:
                newobsdict["states"] = dicts_to_dict_with_arrays(newstates, True)
            else:
                newobsdict["states"] = np.stack(newstates)            
            ret_obs = newobsdict
        return ret_obs
    
    def close(self):
        """Close all workers and shutdown Ray if we initialized it."""
        # Close all worker environments
        for worker in self.workers:
            self.ray.get(worker.close.remote())
        
        # Shutdown Ray if we initialized it
        if hasattr(self, 'ray') and self.ray.is_initialized():
            # Only shutdown if no other processes are using Ray
            try:
                self.ray.shutdown()
            except:
                pass  # Ray might be used by other processes

vecenv_config = {}

def register(config_name, func):
    """Add an environment type (for example RayVecEnv) to the list of available types `rl_games.common.vecenv.vecenv_config`
    Args:
        config_name (:obj:`str`): Key of the environment to create.
        func (:obj:`func`): Function that creates the environment 

    """
    vecenv_config[config_name] = func

def create_vec_env(config_name, num_actors, **kwargs):
    vec_env_name = configurations[config_name]['vecenv_type']
    return vecenv_config[vec_env_name](config_name, num_actors, **kwargs)

register('RAY', lambda config_name, num_actors, **kwargs: RayVecEnv(config_name, num_actors, **kwargs))

from rl_games.envs.brax import BraxEnv
register('BRAX', lambda config_name, num_actors, **kwargs: BraxEnv(config_name, num_actors, **kwargs))

from rl_games.envs.envpool import Envpool
register('ENVPOOL', lambda config_name, num_actors, **kwargs: Envpool(config_name, num_actors, **kwargs))

from rl_games.envs.cule import CuleEnv
register('CULE', lambda config_name, num_actors, **kwargs: CuleEnv(config_name, num_actors, **kwargs))

from rl_games.envs.maniskill import ManiskillEnv
register('MANISKILL', lambda config_name, num_actors, **kwargs: ManiskillEnv(config_name, num_actors, **kwargs))