# Final DNC version packaging all the modules

import torch

from torch import nn

from .memory import memory_unit

from .processor import processor

from .controller import backward_controller

class DNC_Module(nn.Module):

    def __init__(self, num_inputs, num_outputs, controller_size, controller_layers, num_read_heads, num_write_heads, N, M):

        # Params:

        # num_inputs : Size of input data

        # num_outputs : Size of output data

        # controller_size : Size of LSTM Controller output/state

        # controller_layers : Number of layers in LSTM Network

        # num_read_heads : Number of Read heads to be created

        # num_write_heads : Number of Write heads to be created

        # N : Number of memory cells

        # M : Size of Each memory cell

        super(DNC_Module, self).__init__()

        self.num_inputs = num_inputs

        self.num_outputs = num_outputs

        self.N = N

        self.M = M

        # Creating NTM modules

        self.memory = memory_unit(self.N, self.M)

        self.processor = processor(self.num_inputs, self.num_outputs, self.M, self.N, num_read_heads, num_write_heads, controller_size, controller_layers)

        # Creating the Reverse Controller

        self.bController = backward_controller(self.num_inputs, controller_size, controller_layers)

    def initialization(self, batch_size):   # Initializing all the Modules

        self.batch_size = batch_size

        self.memory.reset_memory(batch_size)

        self.previous_state = self.processor.create_new_state(batch_size)

        self.previous_backward_states = self.bController.create_hidden_state(batch_size)

    def backward_prediction(self, X):       # Giving the input to the Backward Controller for making DNC Bi-Directional

        # X dim: (seq_len x batch_size x self.num_inputs)

        output, _ = self.bController(X, self.previous_backward_states)  # Output dim: (seq_len x batch x controller_size)

        return output                                                   # Use embiddings from last to first (reverse way)

    def forward(self, X=None, backward_embeddings=None):

        if X is None:

            X = torch.zeros(self.batch_size, self.num_inputs)

        out, self.previous_state = self.processor(X, backward_embeddings, self.previous_state, self.memory)

        return out, self.previous_state

    '''

    def calculate_num_params(self):     # This maybe for model statistics. Adapted

        num_params = 0

        for p in self.parameters():

            num_params += p.data.view(-1).size(0)

        return num_params

    '''