# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.

# modified version of WMG.
# whenever agent sees a new state, add a new memo and store into it

import torch
import torch.nn as nn
import numpy as np
import pdb

from agents.networks.shared.transformer import Transformer
from agents.networks.shared.general import LinearLayer
from agents.networks.shared.general import SeparateActorCriticLayers
from agents.networks.shared.general import SharedActorCriticLayers
from utils.graph import Graph


# The WMG code was refactored in minor ways before the Sokoban experiments.
#   V1:  Used for Pathfinding and BabyAI.
#        Supports a single factor type.
#        Actor and critic networks are separate.
#        The Memo matrix is initialized full of zero vectors.
#        Pre-embedded age vectors are 1-hot.
#   V2:  Used for Sokoban.
#        Supports multiple factor types, one of which is the Core vector.
#        Actor and critic networks share their initial layer.
#        The Memo matrix is initialized empty.
#        Pre-embedded age vectors are 1-hot, concatenated with the normalized age scalar.


class WMG_State_Network(nn.Module):
    """ Modified version of Working Memory Graph """
    def __init__(self, observation_space, action_space, spec):
        super().__init__()
        self.spec = spec
        self.size_state_matrix = max(spec.val("WMG_MAX_MEMOS"), spec.val("WMG_MAX_OBS"))
        self.factored_observations = isinstance(observation_space, tuple) or isinstance(observation_space, Graph)
        self.tfm_vec_size = spec.val("WMG_NUM_ATTENTION_HEADS") * spec.val("WMG_ATTENTION_HEAD_SIZE")
        if spec.val("WMG_MAX_OBS"):
            assert not self.factored_observations
            self.state_vector_len = observation_space
        else:
            self.state_vector_len = spec.val("WMG_MEMO_SIZE")
        self.prepare_vector_embedding_layers(observation_space)
        self.prepare_age_encodings()
        self.prepare_state_encodings()
        self.tfm = Transformer(
            spec.val("WMG_NUM_ATTENTION_HEADS"),
            spec.val("WMG_ATTENTION_HEAD_SIZE"),
            spec.val("WMG_NUM_LAYERS"),
            spec.val("WMG_HIDDEN_SIZE"),
            spec.val("LAYER_TYPE")
        )
        if spec.val("V2"):
            self.actor_critic_layers = SharedActorCriticLayers(
                self.tfm_vec_size, 2, spec.val("AC_HIDDEN_LAYER_SIZE"), action_space
            )
        else:
            self.actor_critic_layers = SeparateActorCriticLayers(
                self.tfm_vec_size, 2, spec.val("AC_HIDDEN_LAYER_SIZE"), action_space
            )
        if spec.val("WMG_MAX_MEMOS") > 0:
            self.memo_creation_layer = LinearLayer(self.tfm_vec_size, spec.val("WMG_MEMO_SIZE"))
            self.memo_erase_layer = LinearLayer(2*spec.val("WMG_MEMO_SIZE"), spec.val("WMG_MEMO_SIZE"))
            self.memo_weight = LinearLayer(2*spec.val("WMG_MEMO_SIZE"), 1)
        self.net_state = None
        self.prev_state_id = 0
        self.tmp_ctr = 0

    def forward(self, x, old_matrix):
        embedded_vec_list = []

        # Embed the Core and any Factors.
        if self.factored_observations:
            if self.spec.val("V2"):
                for entity in x.entities:
                    embedded_vec_list.append(self.embedding_layers[entity.type](torch.tensor(entity.data)))
            else:
                embedded_vec_list.append(self.core_embedding_layer(torch.tensor(np.float32(x[0]))))
                for factor_vec in x[1:]:
                    embedded_vec_list.append(self.factor_embedding_layer(torch.tensor(np.float32(factor_vec))))
        else:
            x_tens = torch.tensor(np.float32(x))
            embedded_vec_list.append(self.core_embedding_layer(x_tens))

        # Embed the state vectors.
        self.tmp_ctr += 1
        for age, state_id in enumerate(old_matrix.age_index):
            if self.spec.val("V2"):
                embedded_vector = self.add_age_embedding(self.state_embedding_layer(old_matrix.get_vector(age)), age)
            else:
                embedded_vector = self.state_embedding_layer(
                    torch.cat((old_matrix.get_vector(state_id),
                               self.age[age],
                               self.state_encodings[state_id]))
                )
            embedded_vec_list.append(embedded_vector)

        # Apply the Transformer.
        tfm_input = torch.cat(embedded_vec_list).view(1, len(embedded_vec_list), self.tfm_vec_size)
        tfm_output, attention_probs = self.tfm(tfm_input)
        tfm_output = tfm_output[:, 0, :]  # Take only the Core column's output vector.

        # determine the current state
        curr_state_id = x[0]

        # Update the state.
        if self.size_state_matrix:
            if self.spec.val("WMG_MAX_MEMOS"):
                # Calculate the new memo to write into the State Matrix
                new_vector = torch.tanh(self.memo_creation_layer(tfm_output))[0]
                if self.prev_state_id in old_matrix.vectors.keys():
                    old_vector = old_matrix.get_vector(self.prev_state_id)
                else:
                    old_vector = torch.zeros_like(new_vector)
                erase_vector = torch.tanh(self.memo_erase_layer(
                    torch.cat((new_vector, old_vector), 0)))
                memo_write_weight = torch.sigmoid(self.memo_weight(torch.cat((new_vector, old_vector), 0)))
                new_vector = (1.-memo_write_weight) * old_vector * erase_vector\
                    + memo_write_weight * new_vector
            else:
                # Store a new observation.
                new_vector = torch.tanh(x_tens)
            # update the memo at previous state
            old_matrix.add_vector(new_vector, self.prev_state_id)

        else:
            new_matrix = None
        self.prev_state_id = curr_state_id
        # Return the actor-critic logits, and the state matrix.
        policy, value_est = self.actor_critic_layers(tfm_output)

        # return policy, value_est, new_matrix
        return policy, value_est, old_matrix, attention_probs

    def prepare_vector_embedding_layers(self, observation_space):
        if self.spec.val("V2"):
            assert self.factored_observations
            self.factor_sizes = observation_space.entity_type_sizes
            self.embedding_layers = nn.ModuleList()
            for factor_size in self.factor_sizes:
                self.embedding_layers.append(LinearLayer(factor_size, self.tfm_vec_size))
            if self.size_state_matrix:
                self.state_embedding_layer = LinearLayer(self.state_vector_len, self.tfm_vec_size)
        else:
            if self.factored_observations:
                self.core_vec_size = observation_space[0]
                self.factor_vec_size = observation_space[1]
            else:
                self.core_vec_size = observation_space
            self.core_embedding_layer = LinearLayer(self.core_vec_size, self.tfm_vec_size)
            if self.factored_observations:
                self.factor_embedding_layer = LinearLayer(self.factor_vec_size, self.tfm_vec_size)
            if self.size_state_matrix:
                # one S for age matrix, the other for state index matrix
                self.state_embedding_layer = LinearLayer(self.state_vector_len + 2*S, self.tfm_vec_size)

    def prepare_age_encodings(self):
        self.age = []
        if self.spec.val("V2"):
            # self.age contains a total of S+2 vectors, each of size S+2.
            # self.age[0] is for the current time step.
            # self.age[S] is for the oldest time step in the history.
            # self.age[S+1] is for out-of-range ages.
            # One additional element is copied into any age vector to hold the raw age value itself.
            # Each vector is of size S+2: a 1-hot vector of size S+1, and the age scalar.
            self.max_age_index = self.size_state_matrix + 1
            age_vector_size = self.size_state_matrix + 2
            for i in range(self.max_age_index):
                pos = np.zeros(age_vector_size, np.float32)
                pos[i] = 1.
                self.age.append(pos)
            self.age.append(np.zeros(age_vector_size, np.float32))
            self.age_embedding_layer = LinearLayer(age_vector_size, self.tfm_vec_size)
        else:
            for i in range(self.size_state_matrix):
                pos = np.zeros(self.size_state_matrix, np.float32)
                pos[i] = 1.
                self.age.append(torch.tensor(pos))

    def add_age_embedding(self, vector_in, age):
        age_vec = self.age[min(age, self.max_age_index)]
        age_vec[-1] = age / self.size_state_matrix  # Include the normalized age scalar to convey order.
        age_embedding = self.age_embedding_layer(torch.tensor(age_vec))
        return vector_in + age_embedding

    def prepare_state_encodings(self):
        """prepare the state one-hot vectors."""
        self.state_encodings = []
        if self.spec.val("V2"):
            print("prepare_state_encodings() not implemented yet.")
        else:
            for i in range(self.size_state_matrix):
                pos = np.zeros(self.size_state_matrix, np.float32)
                pos[i] = 1.
                self.state_encodings.append(torch.tensor(pos))

    def init_state(self):
        if self.size_state_matrix:
            if not self.net_state:
                self.net_state = StateMatrix(self.size_state_matrix, self.state_vector_len, self.spec)
            else:
                for k in self.net_state.vectors.keys():
                    self.net_state.vectors[k].zero_()
            self.net_state.detach_from_history()
            return self.net_state
        else:
            return None

    def detach_from_history(self, matrix):
        if matrix is not None:
            matrix.detach_from_history()
        return matrix


class StateMatrix(object):
    """Used to store either Memos, or past observations (as factors)."""
    def __init__(self, max_vectors, vector_len, spec):
        self.vector_len = vector_len
        self.max_vectors = max_vectors
        self.spec = spec
        self.age_index = list(range(self.spec.val("WMG_INIT_MEMOS")))  # now keeps unused vectors at all zeros
        if self.spec.val("V2"):
            self.num_vectors = 0  # The matrix starts out empty.
        else:
            self.num_vectors = self.spec.val("WMG_INIT_MEMOS")  # The matrix starts out full.
        self.vectors = {n: torch.tensor(np.zeros(self.vector_len, np.float32)) for n in self.age_index}

    def add_vector(self, vector, state_id):
        if state_id not in self.vectors.keys():
            self.vectors[state_id] = torch.tensor(np.zeros(self.vector_len, np.float32))
            self.age_index.append(state_id)
            self.num_vectors = len(self.age_index)

        self.vectors[state_id] = vector
        # rearranges age_index here
        self.age_index.remove(state_id)
        self.age_index.insert(0, state_id)

    def get_vector(self, state_id):
        assert (state_id >= 0) and (state_id < self.max_vectors)

        # there used to be memory leak here. kept here to check for bugs.
        if state_id not in self.vectors.keys():
            print('inconsistent vector keys and state_ids')
            self.vectors[state_id] = torch.tensor(np.zeros(self.vector_len, np.float32))
        if self.spec.val("V2"):
            state_id += 1  # Map i=0 to the most recent Memo.
        if state_id not in self.age_index:
            print(f'wrong input! i is {state_id}, available states is {self.age_index}')
        return self.vectors[state_id]

    def detach_from_history(self):
        for i in self.vectors.keys():
            self.vectors[i] = self.vectors[i].detach()

# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.

# modified version of WMG.
# whenever agent sees a new state, add a new memo and store into it

import torch
import torch.nn as nn
import numpy as np
import pdb

from agents.networks.shared.transformer import Transformer
from agents.networks.shared.general import LinearLayer
from agents.networks.shared.general import SeparateActorCriticLayers
from agents.networks.shared.general import SharedActorCriticLayers
from utils.graph import Graph


# The WMG code was refactored in minor ways before the Sokoban experiments.
#   V1:  Used for Pathfinding and BabyAI.
#        Supports a single factor type.
#        Actor and critic networks are separate.
#        The Memo matrix is initialized full of zero vectors.
#        Pre-embedded age vectors are 1-hot.
#   V2:  Used for Sokoban.
#        Supports multiple factor types, one of which is the Core vector.
#        Actor and critic networks share their initial layer.
#        The Memo matrix is initialized empty.
#        Pre-embedded age vectors are 1-hot, concatenated with the normalized age scalar.


class WMG_State_Network(nn.Module):
    """ Modified version of Working Memory Graph """
    def __init__(self, observation_space, action_space, spec):
        super().__init__()
        self.spec = spec
        self.size_state_matrix = max(spec.val("WMG_MAX_MEMOS"), spec.val("WMG_MAX_OBS"))
        self.factored_observations = isinstance(observation_space, tuple) or isinstance(observation_space, Graph)
        self.tfm_vec_size = spec.val("WMG_NUM_ATTENTION_HEADS") * spec.val("WMG_ATTENTION_HEAD_SIZE")
        if spec.val("WMG_MAX_OBS"):
            assert not self.factored_observations
            self.state_vector_len = observation_space
        else:
            self.state_vector_len = spec.val("WMG_MEMO_SIZE")
        self.prepare_vector_embedding_layers(observation_space)
        self.prepare_age_encodings()
        self.prepare_state_encodings()
        self.tfm = Transformer(
            spec.val("WMG_NUM_ATTENTION_HEADS"),
            spec.val("WMG_ATTENTION_HEAD_SIZE"),
            spec.val("WMG_NUM_LAYERS"),
            spec.val("WMG_HIDDEN_SIZE"),
            spec.val("LAYER_TYPE")
        )
        if spec.val("V2"):
            self.actor_critic_layers = SharedActorCriticLayers(
                self.tfm_vec_size, 2, spec.val("AC_HIDDEN_LAYER_SIZE"), action_space
            )
        else:
            self.actor_critic_layers = SeparateActorCriticLayers(
                self.tfm_vec_size, 2, spec.val("AC_HIDDEN_LAYER_SIZE"), action_space
            )
        if spec.val("WMG_MAX_MEMOS") > 0:
            self.memo_creation_layer = LinearLayer(self.tfm_vec_size, spec.val("WMG_MEMO_SIZE"))
            self.memo_erase_layer = LinearLayer(2*spec.val("WMG_MEMO_SIZE"), spec.val("WMG_MEMO_SIZE"))
            self.memo_weight = LinearLayer(2*spec.val("WMG_MEMO_SIZE"), 1)
        self.net_state = None
        self.prev_state_id = 0
        self.tmp_ctr = 0

    def forward(self, x, old_matrix):
        embedded_vec_list = []

        # Embed the Core and any Factors.
        if self.factored_observations:
            if self.spec.val("V2"):
                for entity in x.entities:
                    embedded_vec_list.append(self.embedding_layers[entity.type](torch.tensor(entity.data)))
            else:
                embedded_vec_list.append(self.core_embedding_layer(torch.tensor(np.float32(x[0]))))
                for factor_vec in x[1:]:
                    embedded_vec_list.append(self.factor_embedding_layer(torch.tensor(np.float32(factor_vec))))
        else:
            x_tens = torch.tensor(np.float32(x))
            embedded_vec_list.append(self.core_embedding_layer(x_tens))

        # Embed the state vectors.
        self.tmp_ctr += 1
        for age, state_id in enumerate(old_matrix.age_index):
            if self.spec.val("V2"):
                embedded_vector = self.add_age_embedding(self.state_embedding_layer(old_matrix.get_vector(age)), age)
            else:
                embedded_vector = self.state_embedding_layer(
                    torch.cat((old_matrix.get_vector(state_id),
                               self.age[age],
                               self.state_encodings[state_id]))
                )
            embedded_vec_list.append(embedded_vector)

        # Apply the Transformer.
        tfm_input = torch.cat(embedded_vec_list).view(1, len(embedded_vec_list), self.tfm_vec_size)
        tfm_output, attention_probs = self.tfm(tfm_input)
        tfm_output = tfm_output[:, 0, :]  # Take only the Core column's output vector.

        # determine the current state
        curr_state_id = x[0]

        # Update the state.
        if self.size_state_matrix:
            if self.spec.val("WMG_MAX_MEMOS"):
                # Calculate the new memo to write into the State Matrix
                new_vector = torch.tanh(self.memo_creation_layer(tfm_output))[0]
                if self.prev_state_id in old_matrix.vectors.keys():
                    old_vector = old_matrix.get_vector(self.prev_state_id)
                else:
                    old_vector = torch.zeros_like(new_vector)
                erase_vector = torch.tanh(self.memo_erase_layer(
                    torch.cat((new_vector, old_vector), 0)))
                memo_write_weight = torch.sigmoid(self.memo_weight(torch.cat((new_vector, old_vector), 0)))
                new_vector = (1.-memo_write_weight) * old_vector * erase_vector\
                    + memo_write_weight * new_vector
            else:
                # Store a new observation.
                new_vector = torch.tanh(x_tens)
            # update the memo at previous state
            old_matrix.add_vector(new_vector, self.prev_state_id)

        else:
            new_matrix = None
        self.prev_state_id = curr_state_id
        # Return the actor-critic logits, and the state matrix.
        policy, value_est = self.actor_critic_layers(tfm_output)

        # return policy, value_est, new_matrix
        return policy, value_est, old_matrix, attention_probs

    def prepare_vector_embedding_layers(self, observation_space):
        if self.spec.val("V2"):
            assert self.factored_observations
            self.factor_sizes = observation_space.entity_type_sizes
            self.embedding_layers = nn.ModuleList()
            for factor_size in self.factor_sizes:
                self.embedding_layers.append(LinearLayer(factor_size, self.tfm_vec_size))
            if self.size_state_matrix:
                self.state_embedding_layer = LinearLayer(self.state_vector_len, self.tfm_vec_size)
        else:
            if self.factored_observations:
                self.core_vec_size = observation_space[0]
                self.factor_vec_size = observation_space[1]
            else:
                self.core_vec_size = observation_space
            self.core_embedding_layer = LinearLayer(self.core_vec_size, self.tfm_vec_size)
            if self.factored_observations:
                self.factor_embedding_layer = LinearLayer(self.factor_vec_size, self.tfm_vec_size)
            if self.size_state_matrix:
                # one S for age matrix, the other for state index matrix
                self.state_embedding_layer = LinearLayer(self.state_vector_len + 2*S, self.tfm_vec_size)

    def prepare_age_encodings(self):
        self.age = []
        if self.spec.val("V2"):
            # self.age contains a total of S+2 vectors, each of size S+2.
            # self.age[0] is for the current time step.
            # self.age[S] is for the oldest time step in the history.
            # self.age[S+1] is for out-of-range ages.
            # One additional element is copied into any age vector to hold the raw age value itself.
            # Each vector is of size S+2: a 1-hot vector of size S+1, and the age scalar.
            self.max_age_index = self.size_state_matrix + 1
            age_vector_size = self.size_state_matrix + 2
            for i in range(self.max_age_index):
                pos = np.zeros(age_vector_size, np.float32)
                pos[i] = 1.
                self.age.append(pos)
            self.age.append(np.zeros(age_vector_size, np.float32))
            self.age_embedding_layer = LinearLayer(age_vector_size, self.tfm_vec_size)
        else:
            for i in range(self.size_state_matrix):
                pos = np.zeros(self.size_state_matrix, np.float32)
                pos[i] = 1.
                self.age.append(torch.tensor(pos))

    def add_age_embedding(self, vector_in, age):
        age_vec = self.age[min(age, self.max_age_index)]
        age_vec[-1] = age / self.size_state_matrix  # Include the normalized age scalar to convey order.
        age_embedding = self.age_embedding_layer(torch.tensor(age_vec))
        return vector_in + age_embedding

    def prepare_state_encodings(self):
        """prepare the state one-hot vectors."""
        self.state_encodings = []
        if self.spec.val("V2"):
            print("prepare_state_encodings() not implemented yet.")
        else:
            for i in range(self.size_state_matrix):
                pos = np.zeros(self.size_state_matrix, np.float32)
                pos[i] = 1.
                self.state_encodings.append(torch.tensor(pos))

    def init_state(self):
        if self.size_state_matrix:
            if not self.net_state:
                self.net_state = StateMatrix(self.size_state_matrix, self.state_vector_len, self.spec)
            else:
                for k in self.net_state.vectors.keys():
                    self.net_state.vectors[k].zero_()
            self.net_state.detach_from_history()
            return self.net_state
        else:
            return None

    def detach_from_history(self, matrix):
        if matrix is not None:
            matrix.detach_from_history()
        return matrix


class StateMatrix(object):
    """Used to store either Memos, or past observations (as factors)."""
    def __init__(self, max_vectors, vector_len, spec):
        self.vector_len = vector_len
        self.max_vectors = max_vectors
        self.spec = spec
        self.age_index = list(range(self.spec.val("WMG_INIT_MEMOS")))  # now keeps unused vectors at all zeros
        if self.spec.val("V2"):
            self.num_vectors = 0  # The matrix starts out empty.
        else:
            self.num_vectors = self.spec.val("WMG_INIT_MEMOS")  # The matrix starts out full.
        self.vectors = {n: torch.tensor(np.zeros(self.vector_len, np.float32)) for n in self.age_index}

    def add_vector(self, vector, state_id):
        if state_id not in self.vectors.keys():
            self.vectors[state_id] = torch.tensor(np.zeros(self.vector_len, np.float32))
            self.age_index.append(state_id)
            self.num_vectors = len(self.age_index)

        self.vectors[state_id] = vector
        # rearranges age_index here
        self.age_index.remove(state_id)
        self.age_index.insert(0, state_id)

    def get_vector(self, state_id):
        assert (state_id >= 0) and (state_id < self.max_vectors)

        # there used to be memory leak here. kept here to check for bugs.
        if state_id not in self.vectors.keys():
            print('inconsistent vector keys and state_ids')
            self.vectors[state_id] = torch.tensor(np.zeros(self.vector_len, np.float32))
        if self.spec.val("V2"):
            state_id += 1  # Map i=0 to the most recent Memo.
        if state_id not in self.age_index:
            print(f'wrong input! i is {state_id}, available states is {self.age_index}')
        return self.vectors[state_id]

    def detach_from_history(self):
        for i in self.vectors.keys():
            self.vectors[i] = self.vectors[i].detach()

# Copyright (c) Microsoft Corporation. # Licensed under the MIT license. # modified version of WMG. # whenever agent sees a new state, add a new memo and store into it # The WMG code was refactored in minor ways before the Sokoban experiments. #   V1:  Used for Pathfinding and BabyAI. #        Supports a single factor type. #        Actor and critic networks are separate. #        The Memo matrix is initialized full of zero vectors. #        Pre-embedded age vectors are 1-hot. #   V2:  Used for Sokoban. #        Supports multiple factor types, one of which is the Core vector. #        Actor and critic networks share their initial layer. #        The Memo matrix is initialized empty. #        Pre-embedded age vectors are 1-hot, concatenated with the normalized age scalar. Modified version of Working Memory Graph # Embed the Core and any Factors. # Embed the state vectors. # Apply the Transformer. # Take only the Core column's output vector. # determine the current state # Update the state. # Calculate the new memo to write into the State Matrix # Store a new observation. # update the memo at previous state # Return the actor-critic logits, and the state matrix. # return policy, value_est, new_matrix # one S for age matrix, the other for state index matrix # self.age contains a total of S+2 vectors, each of size S+2. # self.age[0] is for the current time step. # self.age[S] is for the oldest time step in the history. # self.age[S+1] is for out-of-range ages. # One additional element is copied into any age vector to hold the raw age value itself. # Each vector is of size S+2: a 1-hot vector of size S+1, and the age scalar. # Include the normalized age scalar to convey order. prepare the state one-hot vectors. Used to store either Memos, or past observations (as factors). # now keeps unused vectors at all zeros # The matrix starts out empty. # The matrix starts out full. # rearranges age_index here # there used to be memory leak here. kept here to check for bugs. # Map i=0 to the most recent Memo.