import numpy as np
import pybullet_envs.deep_mimic.learning.agent_builder as AgentBuilder
import pybullet_envs.deep_mimic.learning.tf_util as TFUtil
from pybullet_envs.deep_mimic.learning.rl_agent import RLAgent
from pybullet_utils.logger import Logger
import pybullet_data

class RLWorld(object):
    def __init__(self, env, arg_parser):
        TFUtil.disable_gpu()

        self.env = env
        self.arg_parser = arg_parser
        self._enable_training = True
        self.train_agents = []
        self.parse_args(arg_parser)

        self.build_agents()
        
        return

    def get_enable_training(self):
        return self._enable_training
    
    def set_enable_training(self, enable):
        self._enable_training = enable
        for i in range(len(self.agents)):
            curr_agent = self.agents[i]
            if curr_agent is not None:
                enable_curr_train = self.train_agents[i] if (len(self.train_agents) > 0) else True
                curr_agent.enable_training = self.enable_training and enable_curr_train

        if (self._enable_training):
            self.env.set_mode(RLAgent.Mode.TRAIN)
        else:
            self.env.set_mode(RLAgent.Mode.TEST)

        return

    enable_training = property(get_enable_training, set_enable_training)
    
    def parse_args(self, arg_parser):
        self.train_agents = self.arg_parser.parse_bools('train_agents')
        num_agents = self.env.get_num_agents()
        assert(len(self.train_agents) == num_agents or len(self.train_agents) == 0)

        return

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

    def build_agents(self):
        num_agents = self.env.get_num_agents()
        print("num_agents=",num_agents)
        self.agents = []

        Logger.print2('')
        Logger.print2('Num Agents: {:d}'.format(num_agents))

        agent_files = self.arg_parser.parse_strings('agent_files')
        print("len(agent_files)=",len(agent_files))
        assert(len(agent_files) == num_agents or len(agent_files) == 0)

        model_files = self.arg_parser.parse_strings('model_files')
        assert(len(model_files) == num_agents or len(model_files) == 0)

        output_path = self.arg_parser.parse_string('output_path')
        int_output_path = self.arg_parser.parse_string('int_output_path')

        for i in range(num_agents):
            curr_file = agent_files[i]
            curr_agent = self._build_agent(i, curr_file)

            if curr_agent is not None:
                curr_agent.output_dir = output_path
                curr_agent.int_output_dir = int_output_path
                Logger.print2(str(curr_agent))

                if (len(model_files) > 0):
                    curr_model_file = model_files[i]
                    if curr_model_file != 'none':
                        curr_agent.load_model(pybullet_data.getDataPath()+"/"+curr_model_file)

            self.agents.append(curr_agent)
            Logger.print2('')

        self.set_enable_training(self.enable_training)

        return

    def update(self, timestep):
        #print("world update!\n")
        self._update_agents(timestep)
        self._update_env(timestep)
        return

    def reset(self):
        self._reset_agents()
        self._reset_env()
        return

    def end_episode(self):
        self._end_episode_agents();
        return

    def _update_env(self, timestep):
        self.env.update(timestep)
        return

    def _update_agents(self, timestep):
        #print("len(agents)=",len(self.agents))
        for agent in self.agents:
            if (agent is not None):
                agent.update(timestep)
        return

    def _reset_env(self):
        self.env.reset()
        return

    def _reset_agents(self):
        for agent in self.agents:
            if (agent != None):
                agent.reset()
        return

    def _end_episode_agents(self):
        for agent in self.agents:
            if (agent != None):
                agent.end_episode()
        return

    def _build_agent(self, id, agent_file):
        Logger.print2('Agent {:d}: {}'.format(id, agent_file))
        if (agent_file == 'none'):
            agent = None
        else:
            agent = AgentBuilder.build_agent(self, id, agent_file)
            assert (agent != None), 'Failed to build agent {:d} from: {}'.format(id, agent_file)
        
        return agent
        
