import torch
from torch import nn

from util.util import get_positional_encoding
from model.encoder_layer import EncoderLayer

def log(string, verbose):
    if verbose:
        print(string)

class EncoderModel(nn.Module):
    def __init__(self, input_dim, hidden_dim, d_ff, output_dim, n_layers,
                 num_heads, use_attention=True, use_feedforward=True,
                 use_positional=True, device='cpu'):
        super(EncoderModel, self).__init__()
        self._device = device
        self._use_positional = use_positional
        self.embedding_layer = nn.Embedding(input_dim, hidden_dim)
        self.layers = nn.ModuleList([
            EncoderLayer(
                hidden_dim,
                d_ff,
                num_heads,
                use_attention,
                use_feedforward,
                device=device
            ) for i in range(n_layers)
        ])
        self.output_layer = nn.Linear(hidden_dim, output_dim)

    def forward(self, x, return_att_weights=False, verbose=False):
        log(f'Handling {x}', verbose)
        # x shape: (seqlen, batch)
        x = x.to(self._device)
        hidden = self.embedding_layer(x)
        # hidden shape: (seqlen, batch, hiddendim)

        if self._use_positional:
            positional_encoding = get_positional_encoding(
                    n_positions=hidden.shape[0],
                    n_dimensions=hidden.shape[-1],
                    device=self._device
                )
            # reshaping to (seqlen, 1, hiddendim)
            positional_encoding = torch.reshape(
                    positional_encoding,
                    (hidden.shape[0], 1, hidden.shape[-1])
            )
            hidden = hidden + positional_encoding

        list_att_weights = []
        for layer in self.layers:
            hidden, att_weights = layer(hidden)
            list_att_weights.append(att_weights)

        result = self.output_layer(hidden)

        log('Result: {result}', verbose)

        if return_att_weights:
            return result, list_att_weights
        else:
            return result