"""Implementation of Neural Turing Machine"""

# global
import ivy
import math
import collections

# local
from ivy import Linear, LSTM


NTMControllerState = collections.namedtuple(
    "NTMControllerState",
    ("controller_state", "read_vector_list", "w_list", "usage_indicator", "M"),
)


def _expand(x, dim, N):
    return ivy.concat([ivy.expand_dims(x, axis=dim) for _ in range(N)], axis=dim)


class NTMCell(ivy.Module):
    def __init__(
        self,
        controller,
        controller_proj,
        output_proj,
        output_dim,
        ctrl_input_size,
        ctrl_output_size,
        total_parameter_num,
        memory_size,
        memory_vector_dim,
        read_head_num,
        write_head_num,
        v=None,
        usage=None,
        addressing_mode="content_and_location",
        shift_range=1,
        clip_value=20,
        init_value=1e-6,
        sequential_writing=False,
        retroactive_updates=False,
        retroactive_discount=0.96,
        with_erase=True,
        seed=0,
    ):
        # vanilla ntm
        self._memory_size = memory_size
        self._memory_vector_dim = memory_vector_dim
        self._read_head_num = read_head_num
        self._write_head_num = write_head_num
        self._init_value = init_value
        self._addressing_mode = addressing_mode
        self._clip_value = clip_value
        self._output_dim = output_dim
        self._shift_range = shift_range
        self._num_parameters_per_head = (
            self._memory_vector_dim + 1 + 1 + (self._shift_range * 2 + 1) + 1
        )
        self._num_heads = self._read_head_num + self._write_head_num
        ivy.seed(seed_value=seed)

        # fns + classes
        self._controller = controller
        self._controller_proj = controller_proj
        self._output_proj = output_proj

        # usage
        if usage is not None:
            self._usage = usage
        else:
            self._usage = ivy.zeros(
                [
                    memory_size,
                ]
            )

        # step
        self._step = 0

        # MERLIN changes
        self._sequential_writing = sequential_writing
        self._retroactive_updates = retroactive_updates
        self._retroactive_discount = retroactive_discount
        self._with_erase = with_erase

        # variables
        ivy.Module.__init__(self, device="cpu")

    def _create_variables(self, device=None, dtype=None):
        vars_dict = dict()
        wlim = (6 / (2 * self._memory_vector_dim)) ** 0.5
        vars_dict["read_weights"] = dict(
            zip(
                ["w_" + str(i) for i in range(self._read_head_num)],
                [
                    ivy.random_uniform(
                        low=-wlim,
                        high=wlim,
                        shape=[
                            self._memory_vector_dim,
                        ],
                        device=device,
                    )
                    for _ in range(self._read_head_num)
                ],
            )
        )
        wlim = (6 / (2 * self._memory_size)) ** 0.5
        vars_dict["write_weights"] = dict(
            zip(
                [
                    "w_" + str(i)
                    for i in range(self._read_head_num + self._write_head_num)
                ],
                [
                    ivy.random_uniform(
                        low=-wlim,
                        high=wlim,
                        shape=[
                            self._memory_size,
                        ],
                        device=device,
                    )
                    for _ in range(self._read_head_num + self._write_head_num)
                ],
            )
        )
        vars_dict["memory"] = (
            ivy.ones([self._memory_size, self._memory_vector_dim], device=device)
            * self._init_value
        )
        return vars_dict

    def _addressing(self, k, beta, g, s, gamma, prev_M, prev_w):
        # Sec 3.3.1 Focusing by Content

        # Cosine Similarity

        k = ivy.expand_dims(k, axis=2)
        inner_product = ivy.matmul(prev_M, k)
        k_norm = ivy.sum(k**2, axis=1, keepdims=True) ** 0.5
        M_norm = ivy.sum(prev_M**2, axis=2, keepdims=True) ** 0.5
        norm_product = M_norm * k_norm
        K = ivy.squeeze(inner_product / (norm_product + 1e-8), axis=-1)  # eq (6)

        # Calculating w^c

        K_amplified = ivy.exp(ivy.expand_dims(beta, axis=1) * K)
        w_c = K_amplified / ivy.sum(K_amplified, axis=1, keepdims=True)  # eq (5)

        if self._addressing_mode == "content":  # Only focus on content
            return w_c

        # Sec 3.3.2 Focusing by Location

        g = ivy.expand_dims(g, axis=1)
        w_g = g * w_c + (1 - g) * prev_w  # eq (7)

        s = ivy.concat(
            [
                s[:, : self._shift_range + 1],
                ivy.zeros(
                    [s.shape[0], self._memory_size - (self._shift_range * 2 + 1)]
                ),
                s[:, -self._shift_range :],
            ],
            axis=1,
        )
        t = ivy.concat([ivy.flip(s, axis=[1]), ivy.flip(s, axis=[1])], axis=1)
        s_matrix = ivy.stack(
            [
                t[:, self._memory_size - i - 1 : self._memory_size * 2 - i - 1]
                for i in range(self._memory_size)
            ],
            axis=1,
        )
        w_ = ivy.sum(ivy.expand_dims(w_g, axis=1) * s_matrix, axis=2)  # eq (8)
        w_sharpen = w_ ** ivy.expand_dims(gamma, axis=1)
        w = w_sharpen / ivy.sum(w_sharpen, axis=1, keepdims=True)  # eq (9)

        return w

    # Public Methods #
    # ---------------#

    def get_start_state(self, _=None, batch_size=None, dtype_str=None, v=None):
        if v is None:
            v = self.v
        else:
            v = ivy.Container(v)
        read_vector_list = [
            _expand(ivy.tanh(var), dim=0, N=batch_size)
            for _, var in v.read_weights.cont_to_iterator()
        ]
        w_list = [
            _expand(ivy.softmax(var), dim=0, N=batch_size)
            for _, var in v.write_weights.cont_to_iterator()
        ]
        usage_indicator = _expand(self._usage, dim=0, N=batch_size)
        M = _expand(v.memory, dim=0, N=batch_size)
        return NTMControllerState(
            controller_state=self._controller.get_initial_state(
                batch_shape=(batch_size,)
            ),
            read_vector_list=read_vector_list,
            w_list=w_list,
            usage_indicator=usage_indicator,
            M=M,
        )

    def _forward(self, x, prev_state):
        prev_read_vector_list = prev_state[1]

        controller_input = ivy.concat([x] + prev_read_vector_list, axis=1)
        controller_output, controller_state = self._controller(
            ivy.expand_dims(controller_input, axis=-2), initial_state=prev_state[0]
        )
        controller_output = controller_output[..., -1, :]

        parameters = self._controller_proj(controller_output)
        parameters = ivy.clip(parameters, -self._clip_value, self._clip_value)
        head_parameter_list = ivy.split(
            parameters[:, : self._num_parameters_per_head * self._num_heads],
            num_or_size_splits=self._num_heads,
            axis=1,
        )
        erase_add_list = ivy.split(
            parameters[:, self._num_parameters_per_head * self._num_heads :],
            num_or_size_splits=2 * self._write_head_num,
            axis=1,
        )

        prev_w_list = prev_state[2]
        prev_M = prev_state[4]
        w_list = []
        for i, head_parameter in enumerate(head_parameter_list):
            k = ivy.tanh(head_parameter[:, 0 : self._memory_vector_dim])
            beta = ivy.softplus(head_parameter[:, self._memory_vector_dim])
            g = ivy.sigmoid(head_parameter[:, self._memory_vector_dim + 1])
            s = ivy.softmax(
                head_parameter[
                    :,
                    self._memory_vector_dim
                    + 2 : self._memory_vector_dim
                    + 2
                    + (self._shift_range * 2 + 1),
                ]
            )
            gamma = ivy.softplus(head_parameter[:, -1]) + 1
            w = self._addressing(k, beta, g, s, gamma, prev_M, prev_w_list[i])
            w_list.append(w)

        # Reading (Sec 3.1)

        read_w_list = w_list[: self._read_head_num]
        if self._step == 0:
            usage_indicator = ivy.zeros_like(w_list[0])
        else:
            usage_indicator = prev_state[3] + ivy.sum(ivy.concat(read_w_list, axis=0))
        read_vector_list = []
        for i in range(self._read_head_num):
            read_vector = ivy.sum(
                ivy.expand_dims(read_w_list[i], axis=2) * prev_M, axis=1
            )
            read_vector_list.append(read_vector)

        # Writing (Sec 3.2)

        prev_wrtie_w_list = prev_w_list[self._read_head_num :]
        w_wr_size = (
            math.ceil(self._memory_size / 2)
            if self._retroactive_updates
            else self._memory_size
        )
        if self._sequential_writing:
            batch_size = ivy.shape(x)[0]
            if self._step < w_wr_size:
                w_wr_list = [
                    ivy.tile(
                        ivy.astype(
                            ivy.one_hot(ivy.array([self._step]), w_wr_size), "float32"
                        ),
                        (batch_size, 1),
                    )
                ] * self._write_head_num
            else:
                batch_idxs = ivy.expand_dims(ivy.arange(batch_size, stop=0), axis=-1)
                mem_idxs = ivy.expand_dims(
                    ivy.argmax(usage_indicator[..., :w_wr_size], axis=-1), axis=-1
                )
                total_idxs = ivy.concat((batch_idxs, mem_idxs), axis=-1)
                w_wr_list = [
                    ivy.scatter_nd(
                        total_idxs,
                        ivy.ones((batch_size,)),
                        shape=(batch_size, w_wr_size),
                    )
                ] * self._write_head_num
        else:
            w_wr_list = w_list[self._read_head_num :]
        if self._retroactive_updates:
            w_ret_list = [
                self._retroactive_discount * prev_wrtie_w[..., w_wr_size:]
                + (1 - self._retroactive_discount) * prev_wrtie_w[..., :w_wr_size]
                for prev_wrtie_w in prev_wrtie_w_list
            ]
            w_wrtie_list = [
                ivy.concat((w_wr, w_ret), axis=-1)
                for w_wr, w_ret in zip(w_wr_list, w_ret_list)
            ]
        else:
            w_wrtie_list = w_wr_list
        M = prev_M
        for i in range(self._write_head_num):
            w = ivy.expand_dims(w_wrtie_list[i], axis=2)
            if self._with_erase:
                erase_vector = ivy.expand_dims(
                    ivy.sigmoid(erase_add_list[i * 2]), axis=1
                )
                M = M * 1 - ivy.matmul(w, erase_vector)
            add_vector = ivy.expand_dims(ivy.tanh(erase_add_list[i * 2 + 1]), axis=1)
            M = M + ivy.matmul(w, add_vector)

        NTM_output = self._output_proj(
            ivy.concat([controller_output] + read_vector_list, axis=1)
        )
        NTM_output = ivy.clip(NTM_output, -self._clip_value, self._clip_value)

        self._step += 1
        return NTM_output, NTMControllerState(
            controller_state=controller_state,
            read_vector_list=read_vector_list,
            w_list=w_list,
            usage_indicator=usage_indicator,
            M=M,
        )

    # Properties #
    # -----------#

    @property
    def _state_size(self):
        return NTMControllerState(
            controller_state=self._controller.state_size[0],
            read_vector_list=[
                self._memory_vector_dim for _ in range(self._read_head_num)
            ],
            w_list=[
                self._memory_size
                for _ in range(self._read_head_num + self._write_head_num)
            ],
            usage_indicator=[self._memory_size],
            M=[self._memory_size, self._memory_vector_dim],
        )

    @property
    def _output_size(self):
        return self._output_dim


class NTM(ivy.Module):
    def __init__(
        self,
        input_dim,
        output_dim,
        ctrl_output_size,
        ctrl_layers,
        memory_size,
        memory_vector_dim,
        read_head_num,
        write_head_num,
        v=None,
        usage=None,
        addressing_mode="content_and_location",
        shift_range=1,
        clip_value=20,
        init_value=1e-6,
        sequential_writing=False,
        retroactive_updates=False,
        retroactive_discount=0.96,
        with_erase=True,
    ):
        ctrl_input_size = read_head_num * memory_vector_dim + input_dim
        num_heads = read_head_num + write_head_num
        num_parameters_per_head = memory_vector_dim + 1 + 1 + (shift_range * 2 + 1) + 1
        total_parameter_num = (
            num_parameters_per_head * num_heads + memory_vector_dim * 2 * write_head_num
        )
        ctrl_v = (
            v.ntm_cell.controller
            if v is not None and "ntm_cell" in v and "controller" in v.ntm_cell
            else None
        )
        ctrl = LSTM(ctrl_input_size, ctrl_output_size, num_layers=ctrl_layers, v=ctrl_v)
        ctrl_proj_v = (
            v.ntm_cell.controller_proj
            if v is not None and "ntm_cell" in v and "controller_proj" in v.ntm_cell
            else None
        )
        ctrl_proj = Linear(ctrl_output_size, total_parameter_num, v=ctrl_proj_v)
        out_proj_v = (
            v.ntm_cell.output_proj
            if v is not None and "ntm_cell" in v and "output_proj" in v.ntm_cell
            else None
        )
        out_proj = Linear(
            ctrl_output_size + read_head_num * memory_vector_dim,
            output_dim,
            v=out_proj_v,
        )

        ntm_v = v.ntm if v is not None and "ntm" in v else None
        ntm_cell = NTMCell(
            ctrl,
            ctrl_proj,
            out_proj,
            output_dim,
            ctrl_input_size,
            ctrl_output_size,
            total_parameter_num,
            memory_size,
            memory_vector_dim,
            read_head_num,
            write_head_num,
            ntm_v,
            usage,
            addressing_mode,
            shift_range,
            clip_value,
            init_value,
            sequential_writing,
            retroactive_updates,
            retroactive_discount,
            with_erase,
        )
        self._ntm_cell = ntm_cell
        ivy.Module.__init__(self, device="cpu", v=v)

    def _forward(self, inputs, hidden=None):
        inputs_shape = list(inputs.shape)
        batch_shape = inputs_shape[:-2]
        time_dim = inputs_shape[-2]
        inputs = ivy.reshape(inputs, [-1] + list(inputs.shape[-2:]))
        if hidden is None:
            hidden = self._ntm_cell.get_start_state(
                inputs, inputs.shape[0], inputs.dtype, v=self.v.ntm_cell
            )
        outputs = []
        for x in ivy.unstack(inputs, axis=1):
            output, hidden = self._ntm_cell(x, hidden)
            outputs.append(output)
        ret = ivy.stack(outputs, axis=1)
        return ivy.reshape(ret, batch_shape + [time_dim, -1])
