import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.nn.init as init
from torch.nn.utils.rnn import pack_padded_sequence
import numpy as np

import misc.config as config


class Net(nn.Module):
    def __init__(self, embedding_tokens):
        super(Net, self).__init__()
        question_features = 1024
        vision_features = config.output_features
        glimpses = 2

        self.text = TextProcessor(embedding_tokens=embedding_tokens,embedding_features=300,lstm_features=question_features,drop=0.5,)
        self.attention = Attention(v_features=vision_features,q_features=question_features,mid_features=512,glimpses=2,drop=0.5,)
        self.classifier = Classifier(in_features=512,mid_features=1024,out_features=config.max_answers,drop=0.5,)

        for m in self.modules():
            if isinstance(m, nn.Linear) or isinstance(m, nn.Conv2d):
                init.xavier_uniform_(m.weight)
                if m.bias is not None:
                    m.bias.data.zero_()

    def forward(self, v, q, q_len):
        q = self.text(q, list(q_len.data))

        v = v / (v.norm(p=2, dim=1, keepdim=True).expand_as(v) + 1e-8)
        att,att_prob = self.attention(v, q)
        # v = apply_attention(v, a)
        # combined = torch.cat([v, q], dim=1)

        answer = self.classifier(att)
        return answer,att_prob
class Net_grad(nn.Module):
    def __init__(self, att_dim):
        super(Net_grad, self).__init__()
        self.attention_dim=att_dim
        self.att_lin = nn.Linear(self.attention_dim, self.attention_dim)
        init.xavier_uniform_(self.att_lin.weight)
    def forward(self, p_att):
        final_att_prob = self.att_lin(p_att)
        return final_att_prob

class Net_Discriminitor(nn.Module):
    def __init__(self, att_dim):
        super(Net_Discriminitor, self).__init__()
        self.attention_dim=att_dim
        self.discriminitor = Discriminitor(in_features=self.attention_dim,mid_features=self.attention_dim,out_features=1,drop=0.5,)
        # init.xavier_uniform_(self.att_lin.weight)
    def forward(self, p_att):
        final_att_prob = self.discriminitor(p_att)
        return final_att_prob

class Discriminitor(nn.Sequential):
    def __init__(self, in_features, mid_features, out_features, drop=0.0):
        super(Discriminitor, self).__init__()
        self.add_module('drop1', nn.Dropout(drop))
        self.add_module('lin1', nn.Linear(in_features, mid_features))
        self.add_module('relu', nn.ReLU())
        self.add_module('drop2', nn.Dropout(drop))
        self.add_module('lin2', nn.Linear(mid_features, out_features))
        self.add_module('sigmoid', nn.Sigmoid())

class Classifier(nn.Sequential):
    def __init__(self, in_features, mid_features, out_features, drop=0.0):
        super(Classifier, self).__init__()
        self.add_module('drop1', nn.Dropout(drop))
        self.add_module('lin1', nn.Linear(in_features, mid_features))
        self.add_module('relu', nn.ReLU())
        self.add_module('drop2', nn.Dropout(drop))
        self.add_module('lin2', nn.Linear(mid_features, out_features))

class TextProcessor(nn.Module):
    def __init__(self, embedding_tokens, embedding_features, lstm_features, drop=0.0):
        super(TextProcessor, self).__init__()
        self.embedding = nn.Embedding(embedding_tokens, embedding_features, padding_idx=0)
        self.drop = nn.Dropout(drop)
        self.tanh = nn.Tanh()
        self.lstm = nn.LSTM(input_size=embedding_features,hidden_size=lstm_features,num_layers=1)
        self.features = lstm_features

        self._init_lstm(self.lstm.weight_ih_l0)
        self._init_lstm(self.lstm.weight_hh_l0)
        self.lstm.bias_ih_l0.data.zero_()
        self.lstm.bias_hh_l0.data.zero_()

        init.xavier_uniform_(self.embedding.weight)

    def _init_lstm(self, weight):
        for w in weight.chunk(4, 0):
            init.xavier_uniform_(w)

    def forward(self, q, q_len):
        embedded = self.embedding(q.type(torch.cuda.LongTensor))
        tanhed = self.tanh(self.drop(embedded))
        packed = pack_padded_sequence(tanhed, q_len, batch_first=True)
        _, (_, c) = self.lstm(packed)
        return c.squeeze(0)

class Attention(nn.Module):
    def __init__(self, v_features, q_features, mid_features, glimpses, drop=0.0):
        super(Attention, self).__init__()
        self.mid_features=mid_features
        self.imgSpatialSize = 196

        self.v_conv = nn.Conv2d(v_features, mid_features, 1, bias=False)  # let self.lin take care of bias
        self.q_lin = nn.Linear(q_features, mid_features)
        # self.x_conv = nn.Conv2d(mid_features, glimpses, 1)
        self.hidden_common = nn.Linear(self.mid_features, 1)    

        self.dropout = nn.Dropout(drop)
        self.relu = nn.ReLU(inplace=True)
        # self.dropout = nn.Dropout(p=self.drop_ratio)
        # self.tanh=nn.Tanh()
        self.softmax=nn.Softmax()

    def forward(self, v, q):
        v = self.v_conv(self.dropout(v))#v torch.Size([200, 512, 14, 14])
        q_lin = self.q_lin(self.dropout(q))
        q = tile_2d_over_nd(q_lin, v)#q torch.Size([200, 512, 14, 14])
        # print('q',q.shape)
        # print('v',v.shape)
        x = self.dropout(self.relu(v + q))#x torch.Size([200, 512, 14, 14])
        # print('x',x.shape)
        img_ques_common=x.view(-1,self.mid_features)#img_ques_common torch.Size([200*14*14, 512])
        h=self.hidden_common(img_ques_common)
        h=h.view(-1,self.imgSpatialSize)
        p=self.softmax(h)# p torch.Size([60, 196])
        
        # -- weighted sum of image features
        img_att=torch.bmm(p.view(-1, 1, self.imgSpatialSize),v.view(-1, self.imgSpatialSize , self.mid_features))# img_att torch.Size([60, 1, 512])
        u = q_lin + img_att.view(-1, self.mid_features)# u torch.Size([60, 512])

        return u,p

def apply_attention(input, attention):
    """ Apply any number of attention maps over the input.
        The attention map has to have the same size in all dimensions except dim=1.
    """
    n, c = input.size()[:2]
    glimpses = attention.size(1)

    # flatten the spatial dims into the third dim, since we don't need to care about how they are arranged
    input = input.view(n, c, -1)
    attention = attention.view(n, glimpses, -1)
    s = input.size(2)

    # apply a softmax to each attention map separately
    # since softmax only takes 2d inputs, we have to collapse the first two dimensions together
    # so that each glimpse is normalized separately
    attention = attention.view(n * glimpses, -1)
    print('attention',attention.shape)
    attention = F.softmax(attention)
    print('attention',attention.shape)

    # apply the weighting by creating a new dim to tile both tensors over
    target_size = [n, glimpses, c, s]
    input = input.view(n, 1, c, s).expand(*target_size)
    attention = attention.view(n, glimpses, 1, s).expand(*target_size)
    weighted = input * attention
    # sum over only the spatial dimension
    weighted_mean = weighted.sum(dim=3)
    # the shape at this point is (n, glimpses, c, 1)
    return weighted_mean.view(n, -1)


def tile_2d_over_nd(feature_vector, feature_map):
    """ Repeat the same feature vector over all spatial positions of a given feature map.
        The feature vector should have the same batch size and number of features as the feature map.
    """
    n, c = feature_vector.size()
    spatial_size = feature_map.dim() - 2
    tiled = feature_vector.view(n, c, *([1] * spatial_size)).expand_as(feature_map)
    return tiled
