import ray
import torch

from copy import deepcopy
from itertools import count

import random

from ...env.env import BitcoinTradeEnv
from ..utils import inv_rescale, rescale

class LocalMemory(object):
    def __init__(self):
        self.actions = []
        self.rewards = []
        self.obs = []
        self.next_obs = []
        self.dones = []

    def __len__(self):
        return len(self.dones)

    def reset(self, num_instance_to_keep=0):
        if num_instance_to_keep == 0:
            self.actions = []
            self.rewards = []
            self.obs = []
            self.next_obs = []
            self.dones = []
        else:
            del self.actions[:-num_instance_to_keep]
            del self.rewards[:-num_instance_to_keep]
            del self.obs[:-num_instance_to_keep]
            del self.next_obs[:-num_instance_to_keep]
            del self.dones[:-num_instance_to_keep]

@ray.remote(num_cpus=1)
class Actor(object):
    def __init__(self,
                 agent_core_net,
                 actor_id,
                 actor_total_num,
                 gamma,
                 nstep,
                 actor_update_frequency,
                 memory_size_bound,
                 device,
                 memory_server,
                 actor_epsilon,
                 actor_alpha,
                 parameter_server,
                 update_lambda,
                 env_config,
                 trade_data_path):
        self.actor_net = deepcopy(agent_core_net).cpu()
        self.actor_net.load_state_dict({k: v.cpu() for k, v in agent_core_net.state_dict().items()})
        self.actor_net.eval()
        self.actor_id = actor_id
        self.actor_total_num = actor_total_num
        self.gamma = gamma
        self.nstep = nstep
        self.actor_update_frequency = actor_update_frequency
        self.memory_size_bound = memory_size_bound
        self.device = device
        self.memory_server = memory_server
        self.parameter_server = parameter_server
        self.update_lambda = update_lambda
        self.epsilon = actor_epsilon ** (1 + (actor_id * actor_alpha) / (self.actor_total_num - 1))
        self.env = BitcoinTradeEnv(trade_data_path, env_config)
        self.local_memory = LocalMemory()
        self.episode_count = 0
        self.act_count = 0
        self.memory_size_bound = memory_size_bound

    def epsilon_greedy_policy(self, state):
        if random.random() < self.epsilon:
            action = torch.from_numpy(self.env.get_action_mask()).float().multinomial(1).item()
            # action = random.randint(0,2)
        else:
            with torch.no_grad():
                action_value = self.actor_net(state)
            action = action_value.argmax(dim=1).squeeze().item()

        return action

    def update_agent_from_learner(self):
        learner_state_dict = ray.get(self.parameter_server.send_latest_parameter_to_actor.remote())
        self.actor_net.load_state_dict({k: (1 - self.update_lambda)*v1 + self.update_lambda*v2 for k, v1, v2 in zip(learner_state_dict.keys(), self.actor_net.state_dict().values(), learner_state_dict.values())})

    def transit_to_nstep_return(self, start_index, end_index):
        sum_rewards = 0
        next_obs = None
        dones = False
        for i, index in enumerate(range(start_index, end_index + 1)):
            sum_rewards += (self.gamma**i) * self.local_memory.rewards[index]
            next_obs = self.local_memory.next_obs[index]

            if self.local_memory.dones[index]:
                dones = True
                break

        return [sum_rewards, next_obs, dones]

    def compute_local_priority(self, reward, dones, actions, state, next_state):
        action_value = self.actor_net(state)
        next_state_action_value = self.actor_net(next_state)

        non_terminal_mask = 1 - dones.float()
        terminal_mask = dones.float()
        action_value_target = rescale((reward + (self.gamma ** (self.nstep + 1))*inv_rescale(next_state_action_value.max(dim=1)[0])) * non_terminal_mask + reward * terminal_mask)
        td_error = (action_value_target - action_value.gather(1, actions.view(-1, 1)).view(-1)).abs() + 0.01

        return list(td_error.split(1))

    def run(self):
        obs = self.env.reset()
        step_count = 0

        with torch.no_grad():
            for _ in count():
                state = torch.cat(
                    [
                        torch.from_numpy(obs).float(),
                        torch.from_numpy(self.env.get_action_mask()).float().reshape(1, -1).repeat(obs.shape[0], 1)
                    ],
                    dim=1
                )
                self.local_memory.obs.append(state)

                # sample action
                action = self.epsilon_greedy_policy(state.unsqueeze(0))
                self.local_memory.actions.append(action)
                obs, reward, done, _ = self.env.step(action)

                self.local_memory.rewards.append(reward)
                self.local_memory.dones.append(done)
                next_state = torch.cat(
                    [
                        torch.from_numpy(obs).float(),
                        torch.from_numpy(self.env.get_action_mask()).float().reshape(1, -1).repeat(obs.shape[0], 1)
                    ],
                    dim=1
                )
                self.local_memory.next_obs.append(next_state)

                if done:
                    obs = self.env.reset()
                    self.episode_count += 1
                    step_count = 0

                else:
                    if self.local_memory.__len__() == 25 + self.nstep:
                        # transit to multi-step output
                        actions = []
                        rewards = []
                        states = []
                        next_states = []
                        dones = []
                        for start_index in [i for i in range(self.local_memory.__len__())][:-self.nstep]:
                            end_index = start_index + self.nstep
                            nstep_return_output = self.transit_to_nstep_return(start_index, end_index)
                            rewards.append(nstep_return_output[0])
                            dones.append(nstep_return_output[2])
                            actions.append(self.local_memory.actions[start_index])
                            states.append(self.local_memory.obs[start_index])
                            next_states.append(nstep_return_output[1])

                        # compute priority
                        priorities = self.compute_local_priority(
                            reward=torch.tensor(rewards, dtype=torch.float),
                            dones=torch.tensor(dones, dtype=torch.bool),
                            actions=torch.tensor(actions).long(),
                            state=torch.stack(states, dim=0),
                            next_state=torch.stack(next_states, dim=0)
                        )
                        self.memory_server.receive_sample_from_actor.remote([
                            actions,
                            rewards,
                            states,
                            next_states,
                            dones,
                            priorities
                        ])

                        self.local_memory.reset(self.nstep)

                if self.act_count % self.actor_update_frequency == 0 and self.act_count > 0:
                    self.update_agent_from_learner()

                step_count += 1
                self.act_count += 1

class LocalMemoryR2D2(object):
    def __init__(self):
        self.actions = []
        self.rewards = []
        self.obs = []
        self.hns_buffer = []
        self.cns_buffer = []
        self.dones = []

    def __len__(self):
        return len(self.dones)

    def reset(self, num_instance_to_keep=0):
        if num_instance_to_keep == 0:
            self.actions = []
            self.rewards = []
            self.obs = []
            self.hns_buffer = []
            self.cns_buffer = []
            self.dones = []
        else:
            del self.actions[:-num_instance_to_keep]
            del self.rewards[:-num_instance_to_keep]
            del self.obs[:-num_instance_to_keep]
            del self.hns_buffer[:-num_instance_to_keep]
            del self.cns_buffer[:-num_instance_to_keep]
            del self.dones[:-num_instance_to_keep]

@ray.remote(num_cpus=1)
class ActorR2D2():
    def __init__(self,
                 agent_core_net,
                 actor_id,
                 actor_total_num,
                 gamma,
                 nstep,
                 actor_update_frequency,
                 memory_size_bound,
                 device,
                 memory_server,
                 actor_epsilon,
                 actor_alpha,
                 parameter_server,
                 update_lambda,
                 sequence_length,
                 hidden_state_dim,
                 env_config,
                 trade_data_path,
                 num_layer,
                 burn_in_length=1):
        self.actor_net = deepcopy(agent_core_net).cpu()
        self.actor_net.load_state_dict({k: v.cpu() for k, v in agent_core_net.state_dict().items()})
        self.actor_net.eval()
        self.actor_id = actor_id
        self.actor_total_num = actor_total_num
        self.gamma = gamma
        self.nstep = nstep
        self.actor_update_frequency = actor_update_frequency
        self.memory_size_bound = memory_size_bound
        self.device = device
        self.memory_server = memory_server
        self.parameter_server = parameter_server
        self.update_lambda = update_lambda
        self.sequence_length = sequence_length
        self.hidden_state_dim = hidden_state_dim
        self.burn_in_length = burn_in_length
        self.num_layer = num_layer
        self.epsilon = actor_epsilon ** (1 + (actor_id * actor_alpha) / (self.actor_total_num - 1))
        self.env = BitcoinTradeEnv(trade_data_path, env_config)
        self.local_memory = LocalMemoryR2D2()
        self.episode_count = 0
        self.act_count = 0

    def epsilon_greedy_policy(self, state, hidden_state):
        hn, cn = hidden_state
        with torch.no_grad():
            action_value, (hn, cn) = self.actor_net(state, (hn.unsqueeze(0), cn.unsqueeze(0)))

        if random.random() < self.epsilon:
            action = torch.from_numpy(self.env.get_action_mask()).float().multinomial(1).item()
            # action = random.randint(0,2)
        else:
            action = action_value.argmax(dim=2).squeeze().item()

        return action, (hn, cn)

    def update_agent_from_learner(self):
        learner_state_dict = ray.get(self.parameter_server.send_latest_parameter_to_actor.remote())
        self.actor_net.load_state_dict({k: (1 - self.update_lambda)*v1 + self.update_lambda*v2 for k, v1, v2 in zip(learner_state_dict.keys(), self.actor_net.state_dict().values(), learner_state_dict.values())})

    def transit_to_nstep_return(self, start_index, end_index):
        sum_rewards = 0
        dones = False
        for i, index in enumerate(range(start_index, end_index + 1)):
            sum_rewards += (self.gamma**i) * self.local_memory.rewards[index]

            if self.local_memory.dones[index]:
                dones = True
                break

        return [sum_rewards, dones]

    def compute_local_priority(self, reward, dones, actions, sequential_state_input, sequential_hidden_state):
        hns = sequential_hidden_state[0]
        cns = sequential_hidden_state[1]
        action_value, _ = self.actor_net(sequential_state_input, (torch.stack(hns, dim=0), torch.stack(cns, dim=0)))
        action_value = action_value.squeeze()

        dones_ = dones[self.burn_in_length+self.nstep+1:]
        reward_ = reward[self.burn_in_length:-self.nstep-1]
        non_terminal_mask = 1 - dones_.float()
        terminal_mask = dones_.float()
        action_value_target = rescale((reward_ + (self.gamma ** (self.nstep + 1))*inv_rescale(action_value[self.burn_in_length+self.nstep+1:, :].max(dim=1)[0])) * non_terminal_mask + reward_ * terminal_mask)

        td_error = (action_value_target - action_value[self.burn_in_length+self.nstep+1:, :].gather(1, actions[self.burn_in_length+self.nstep+1:].view(-1, 1)).view(-1)).abs() + 0.01
        priority = (0.9 * td_error.max() + 0.1 * td_error.mean()).view(1)

        return priority

    def send_to_memory_pool(self, done):
        sequential_action = []
        sequential_reward = []
        sequential_state = []
        sequential_dones = []
        sequential_hns = []
        sequential_cns = []
        if done:
            for start_index in [i for i in range(len(self.local_memory))]:
                if self.burn_in_length <= start_index < len(self.local_memory) - self.nstep - 1 :
                    end_index = start_index + self.nstep
                    nstep_return_output = self.transit_to_nstep_return(start_index, end_index)
                    sequential_reward.append(nstep_return_output[0])
                    sequential_dones.append(nstep_return_output[1])
                else:
                    sequential_reward.append(self.local_memory.rewards[start_index])
                    sequential_dones.append(self.local_memory.dones[start_index])

                sequential_action.append(self.local_memory.actions[start_index])
                sequential_state.append(self.local_memory.obs[start_index])
                sequential_hns.append(self.local_memory.hns_buffer[start_index])
                sequential_cns.append(self.local_memory.cns_buffer[start_index])

        else:
            for start_index in [i for i in range(self.sequence_length + self.burn_in_length + self.nstep + 1)]:
                if self.burn_in_length <= start_index < self.sequence_length + self.burn_in_length:
                    end_index = start_index + self.nstep
                    nstep_return_output = self.transit_to_nstep_return(start_index, end_index)
                    sequential_reward.append(nstep_return_output[0])
                    sequential_dones.append(nstep_return_output[1])
                else:
                    sequential_reward.append(self.local_memory.rewards[start_index])
                    sequential_dones.append(self.local_memory.dones[start_index])

                sequential_action.append(self.local_memory.actions[start_index])
                sequential_state.append(self.local_memory.obs[start_index])
                sequential_hns.append(self.local_memory.hns_buffer[start_index])
                sequential_cns.append(self.local_memory.cns_buffer[start_index])

            # compute priority
            sequential_priority = self.compute_local_priority(
                reward=torch.tensor(sequential_reward, dtype=torch.float),
                dones=torch.tensor(sequential_dones, dtype=torch.bool),
                actions=torch.tensor(sequential_action).long(),
                sequential_state_input=torch.stack(sequential_state, dim=0).unsqueeze(0),
                sequential_hidden_state=(self.local_memory.hns_buffer, self.local_memory.cns_buffer)
            )
            self.memory_server.receive_sample_from_actor.remote([
                [torch.tensor(sequential_action).long()],
                [torch.tensor(sequential_reward, dtype=torch.float)],
                [torch.stack(sequential_state, dim=0)],
                [torch.tensor(sequential_dones, dtype=torch.bool)],
                [torch.stack(sequential_hns, dim=0)],
                [torch.stack(sequential_cns, dim=0)],
                [sequential_priority]
            ])

    def run(self):
        obs = self.env.reset()
        hn, cn = (torch.zeros(self.num_layer, 1, self.hidden_state_dim), torch.zeros(self.num_layer, 1, self.hidden_state_dim))
        self.local_memory.hns_buffer.append(hn)
        self.local_memory.cns_buffer.append(cn)
        step_count = 0

        for _ in count():
            with torch.no_grad():
                state = torch.cat(
                    [
                        torch.from_numpy(obs).float(),
                        torch.from_numpy(self.env.get_action_mask()).float().reshape(1, -1).repeat(obs.shape[0], 1)
                    ],
                    dim=1
                )
                self.local_memory.obs.append(state)

                # sample action
                action, (hn, cn) = self.epsilon_greedy_policy(state.unsqueeze(0).unsqueeze(0), (hn, cn))
                self.local_memory.actions.append(action)
                self.local_memory.hns_buffer.append(hn)
                self.local_memory.cns_buffer.append(cn)

                obs, reward, done, _ = self.env.step(action)

                self.local_memory.rewards.append(reward)
                self.local_memory.dones.append(done)

                if done:
                    self.send_to_memory_pool(done=True)
                    obs = self.env.reset()
                    self.episode_count += 1
                    step_count = 0

                    hn, cn = (torch.zeros(self.num_layer, 1, self.hidden_state_dim), torch.zeros(self.num_layer, 1, self.hidden_state_dim))
                    self.local_memory.reset(0)
                    self.local_memory.hns_buffer.append(hn)
                    self.local_memory.cns_buffer.append(cn)

                else:
                    if self.local_memory.__len__() == self.sequence_length + self.burn_in_length + self.nstep + 1:
                        # transit to multi-step output
                        self.send_to_memory_pool(done=False)
                        self.local_memory.reset((self.nstep + self.burn_in_length + self.sequence_length + 1) // 2)

                if self.act_count % self.actor_update_frequency == 0 and self.act_count > 0:
                    self.update_agent_from_learner()

                step_count += 1
                self.act_count += 1