import torch
import torch.nn as nn
import math
import torch.nn.functional as F
from torch.nn import init
from model.DynamicInteraction import DynamicInteraction_Layer0, DynamicInteraction_Layer
import numpy as np

class InteractionModule(nn.Module):
    def __init__(self, opt, num_layer_routing=3, num_cells=4, path_hid=128):
        super(InteractionModule, self).__init__()
        self.opt = opt
        self.num_cells = num_cells = 4
        self.repre_size = 300
        self.num_domain = 6
        self.dynamic_itr_l0 = DynamicInteraction_Layer0(opt, num_cells, num_cells)
        self.dynamic_itr_l1 = DynamicInteraction_Layer(opt, num_cells, num_cells)
        self.dynamic_itr_l2 = DynamicInteraction_Layer(opt, num_cells, 1)
        total_paths = num_cells ** 2 * (num_layer_routing - 1) + num_cells
        self.path_mapping = nn.Linear(total_paths, path_hid)
        self.bn = nn.BatchNorm1d(opt.embed_size)
        # self.pred_proj = nn.Linear(self.repre_size, self.num_domain)
        self.pred_proj = nn.Sequential(
                        nn.Linear(self.repre_size, self.num_domain),   ## 二分类
                        nn.Sigmoid())
        self.loss_BCE = nn.BCELoss()

    # def calSim_i2t(self, ref_imgs, stc, rgn):
    #     ''' ref_imgs--(n_img, n_stc, n_rgn, d), 
    #         stc -- (n_stc, d)
    #     '''
    #     n_img, n_stc, n_rgn, d = ref_imgs.size()
    #     ref_imgs = (self.bn(ref_imgs.view(n_img*n_stc*n_rgn, d))).view(n_img, n_stc, n_rgn, d)
    #     ref_imgs = ref_imgs.mean(2) + ref_imgs.max(2)[0]
    #     ref_imgs = F.normalize(ref_imgs, dim=-1)    # (n_img, n_stc, d)
    #     stc = stc.unsqueeze(0)
    #     sim = (stc * ref_imgs).sum(-1)
    #     return sim

    # def calSim_t2i(self, ref_wrds, img, wrd, stc_lens):
    #     ''' ref_wrds--(n_img, n_stc, n_wrd, d), 
    #         img -- (n_img, d)
    #     '''
    #     n_img, n_stc, n_wrd, d = ref_wrds.size()
    #     ref_wrds = (self.bn(ref_wrds.view(n_img*n_stc*n_wrd, d))).view(n_img, n_stc, n_wrd, d)
    #     ref_wrds = ref_wrds.mean(2) + ref_wrds.max(2)[0]
    #     ref_wrds = F.normalize(ref_wrds, dim=-1)    # (n_img, n_stc, d)
    #     img = img.unsqueeze(1) # (n_img, 1, d)
    #     sim = (img * ref_wrds).sum(-1)
    #     return sim

    def calScores(self, aggr_res, rgn, img, wrd, stc, stc_lens):
        assert len(aggr_res) == 1
        aggr_res = aggr_res[0]
        if self.opt.direction == 'i2t':
            scores = self.calSim_i2t(aggr_res, stc, rgn)
        else:
            scores = self.calSim_t2i(aggr_res, img, wrd, stc_lens)      
        return scores

    def calLoss(self, repre, domain_label):
        pred_y = self.pred_proj(repre)
        # print('pred_y:', pred_y)
        # print('domain_label:', domain_label)
        domain_label = domain_label.float()
        # pred_y1 = torch.log(F.softmax(pred_y0,dim=1))
        # pred_y1 = nn.Sigmoid(pred_y0, dim=1)
        loss = self.loss_BCE(pred_y, domain_label)
        # loss=F.nll_loss(pred_y1, domain_label)
        return loss


    # def forward(self, rgn, img, wrd, stc, stc_lens):
    def forward(self, repre_init, domain_label, know2repre_food, know2repre_sightseeing, know2repre_mall, know2repre_nightlife, know2repre_hotel):
        # pairs_emb_lst, paths_l0 = self.dynamic_itr_l0(rgn, img, wrd, stc, stc_lens)
        # pairs_emb_lst, paths_l1 = self.dynamic_itr_l1(pairs_emb_lst, rgn, img, wrd, stc, stc_lens)
        # pairs_emb_lst, paths_l2 = self.dynamic_itr_l2(pairs_emb_lst, rgn, img, wrd, stc, stc_lens)
        # score = self.calScores(pairs_emb_lst, rgn, img, wrd, stc, stc_lens)
        pairs_emb_lst, paths_l0 = self.dynamic_itr_l0(repre_init, know2repre_food, know2repre_sightseeing, know2repre_mall, know2repre_nightlife, know2repre_hotel)
        pairs_emb_lst, paths_l1 = self.dynamic_itr_l1(pairs_emb_lst, know2repre_food, know2repre_sightseeing, know2repre_mall, know2repre_nightlife, know2repre_hotel)
        pairs_emb_lst, paths_l2 = self.dynamic_itr_l2(pairs_emb_lst, know2repre_food, know2repre_sightseeing, know2repre_mall, know2repre_nightlife, know2repre_hotel)

        # max pooling之后的表征 .cpu().numpy()
        # final_repre_update = torch.max(pairs_emb_lst[0], 1)
        
        # final_repre_update = torch.mean(pairs_emb_lst[0], 1)
        # ==不取mean操作==final_repre_update = pairs_emb_lst[0]
        # print('final_repre_update: ', np.shape(final_repre_update.cpu().numpy()))
        # torch.Size([64, 6, 300])
        # print('pairs_emb_lst[0]: ', np.shape(pairs_emb_lst[0]))
        # exit()
        # final_repre_update1 = torch.mean(pairs_emb_lst[0], 1)
        final_repre_update = torch.mean(pairs_emb_lst[0], 1)
        loss_return = self.calLoss(final_repre_update, domain_label)
        # print('loss_return: ', np.shape(loss_return))
        # print('loss_return: ', loss_return)

        # print('三层卷积完成！')
        # pairs_emb_lst:  1
        # print('pairs_emb_lst: ', len(pairs_emb_lst))
        # print('pairs_emb_lst: ', np.shape(pairs_emb_lst))
        # print('paths_l0: ', np.shape(paths_l0))

        # exit()
        return final_repre_update, loss_return
        # # pairs_emb_lst, paths_l1 = self.dynamic_itr_l1(pairs_emb_lst)
        # # pairs_emb_lst, paths_l2 = self.dynamic_itr_l2(pairs_emb_lst)
        # score = self.calScores(pairs_emb_lst, rgn, img, wrd, stc, stc_lens)

        # n_img, n_stc = paths_l2.size()[:2]
        # if self.opt.direction == 'i2t':
        #     paths_l0 = paths_l0.contiguous().view(n_img, -1).unsqueeze(1).expand(-1, n_stc, -1)
        #     paths_l1 = paths_l1.view(n_img, n_stc, -1)
        #     paths_l2 = paths_l2.view(n_img, n_stc, -1)
        #     paths = torch.cat([paths_l0, paths_l1, paths_l2], dim=-1) # (n_img, n_stc, total_paths)
        #     paths = paths.mean(dim=1) # (n_img, total_paths)
        # else:
        #     paths_l0 = paths_l0.contiguous().view(n_stc, -1).unsqueeze(0).expand(n_img, -1, -1)
        #     paths_l1 = paths_l1.view(n_img, n_stc, -1)
        #     paths_l2 = paths_l2.view(n_img, n_stc, -1)
        #     paths = torch.cat([paths_l0, paths_l1, paths_l2], dim=-1) # (n_img, n_stc, total_paths)
        #     paths = paths.mean(dim=0) # (n_stc, total_paths)
            
        # paths = self.path_mapping(paths)
        # paths = F.normalize(paths, dim=-1)
        # sim_paths = paths.matmul(paths.t())

        # if self.training:
        #     return score, sim_paths
        # else:
        #     return score