import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F

from .wordebd import WORDEBD
from .auxiliary.factory import get_embedding
from collections import OrderedDict


class CNN(nn.Module):
    '''
        An aggregation method that encodes every document through different
        convolution filters (followed by max-over-time pooling).
    '''

    def __init__(self, ebd, args):
        super(CNN, self).__init__()
        self.args = args

        self.ebd = ebd
        self.aux = get_embedding(args)

        self.input_dim = self.ebd.embedding_dim + self.aux.embedding_dim

        # Convolution
        self.convs = nn.ModuleList([nn.Conv1d(
            in_channels=self.input_dim,
            out_channels=args.cnn_num_filters,
            kernel_size=K) for K in args.cnn_filter_sizes])

        # used for visualization
        if args.mode == 'visualize':
            self.scores = [[] for _ in args.cnn_filter_sizes]

        self.ebd_dim = args.cnn_num_filters * len(args.cnn_filter_sizes)

    def _conv_max_pool(self, x, conv_filter=None, weights=None):
        '''
        Compute sentence level convolution
        Input:
            x:      batch_size, max_doc_len, embedding_dim
        Output:     batch_size, num_filters_total
        '''
        assert(len(x.size()) == 3)

        x = x.permute(0, 2, 1)  # batch_size, embedding_dim, doc_len
        x = x.contiguous()

        # Apply the 1d conv. Resulting dimension is
        # [batch_size, num_filters, doc_len-filter_size+1] * len(filter_size)
        assert(not ((conv_filter is None) and (weights is None)))
        if conv_filter is not None:
            x = [conv(x) for conv in conv_filter]

        elif weights is not None:
            x = [F.conv1d(x, weight=weights['convs.{}.weight'.format(i)],
                          bias=weights['convs.{}.bias'.format(i)])
                 for i in range(len(self.args.cnn_filter_sizes))]

        # max pool over time. Resulting dimension is
        # [batch_size, num_filters] * len(filter_size)
        x = [F.max_pool1d(sub_x, sub_x.size(2)).squeeze(2) for sub_x in x]

        # concatenate along all filters. Resulting dimension is
        # [batch_size, num_filters_total]
        x = torch.cat(x, 1)
        x = F.relu(x)

        return x

    def forward(self, data, weights=None):
        '''
            @param data dictionary
                @key text: batch_size * max_text_len
            @param weights placeholder used for maml

            @return output: batch_size * embedding_dim
        '''
        # Apply the word embedding, result:  batch_size, doc_len, embedding_dim
        ebd = self.ebd(data, weights)

        # add augmented embedding if applicable
        aux = self.aux(data, weights)

        ebd = torch.cat([ebd, aux], dim=2)

        # apply 1d conv + max pool, result:  batch_size, num_filters_total
        if weights is None:
            ebd = self._conv_max_pool(ebd, conv_filter=self.convs)
        else:
            ebd = self._conv_max_pool(ebd, weights=weights)

        # update max scores
        if self.args.mode == 'visualize':
            for i, s in enumerate(self.compute_score(data)):
                self.scores[i].append(torch.max(s).item())

        return ebd

    def compute_score(self, data, normalize=False):
        # preparing the input
        ebd = self.ebd(data)
        aux = self.aux(data)
        # (batch_size, doc_len, input_dim)
        x = torch.cat([ebd, aux], dim=-1).detach()

        # (out_channels, in_channels, kernel_size)
        w = [c.weight.data for c in self.convs]
        # (kernel_size, out_channels, in_channels)
        w = [c.permute(2, 0, 1) for c in w]
        # (out_channels * kernel_size, in_channels)
        w = [c.reshape(-1, self.input_dim) for c in w]

        # (batch_size, doc_len, out_channels * kernel_size)
        x = [x @ c.t() for c in w]
        # (batch_size, doc_len)
        x = [F.max_pool1d(z, z.shape[-1]).squeeze(-1) for z in x]

        if normalize:
            x = [x / np.mean(s) for x, s in zip(x, self.scores)]

        return x
