import torch
import torch.nn as nn
import torch.nn.functional as F
import copy
import pickle
import numpy as np
from config.global_config import GlobalConfig


from model.Cells import RectifiedIdentityCell, IntraModelReasoningCell, GlobalLocalGuidanceCell, CrossModalRefinementCell

def unsqueeze2d(x):
    return x.unsqueeze(-1).unsqueeze(-1)

def unsqueeze3d(x):
    return x.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)

def clones(module, N):
    '''Produce N identical layers.'''
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])

class DynamicInteraction_Layer0(nn.Module):
    def __init__(self, opt, num_cell, num_out_path):
        super(DynamicInteraction_Layer0, self).__init__()
        self.opt = opt
        self.threshold = 0.0001
        self.eps = 1e-8
        self.num_cell = num_cell
        self.num_out_path = num_out_path
        self.ric = RectifiedIdentityCell(opt, num_out_path)
        self.imrc = IntraModelReasoningCell(opt, num_out_path)
        self.glgc = GlobalLocalGuidanceCell(opt, num_out_path)
        self.cmrc = CrossModalRefinementCell(opt, num_out_path)
        # triplet_vec1 = torch.tensor(triplet_vec)
        self.temperature = 1.0

    def forward(self, repre_init, know2repre_food, know2repre_sightseeing, know2repre_mall, know2repre_nightlife, know2repre_hotel):
        # if self.opt.direction == 'i2t':
        aggr_res_lst, all_path_prob = self.forward_i2t(repre_init, know2repre_food, know2repre_sightseeing, know2repre_mall, know2repre_nightlife, know2repre_hotel)
        # else:
        #     aggr_res_lst = self.forward_t2i(rgn, img, wrd, stc, stc_lens)
        return aggr_res_lst, all_path_prob
         
    def forward_i2t(self, repre_init, know2repre_food, know2repre_sightseeing, know2repre_mall, know2repre_nightlife, know2repre_hotel):
        path_prob = [None] * self.num_cell
        emb_lst = [None] * self.num_cell
        # emb_lst[0]: torch.Size([64, 6, 300]) ==更新后的表征
        # path_prob[0]: torch.Size([64, 4]) ==cell0对于下一层四个cell四条path的概率 
        emb_lst[0], path_prob[0] = self.ric(repre_init)
        emb_lst[1], path_prob[1] = self.imrc(repre_init)
        emb_lst[2], path_prob[2] = self.glgc(repre_init)
        emb_lst[3], path_prob[3] = self.cmrc(repre_init, know2repre_food, know2repre_sightseeing, know2repre_mall, know2repre_nightlife, know2repre_hotel)

        # print('path_prob[0]: ', path_prob[0])
        # print('path_prob[1]: ', path_prob[1])
        # print('path_prob[2]: ', path_prob[2])
        # print('path_prob[3]: ', path_prob[3])

        # gate_mask_shape:  torch.Size([64, 4])
        # all_path_prob_shape:  torch.Size([64, 4, 4])
        # all_path_prob_shape:  torch.Size([64, 4, 4])
        gate_mask = (sum(path_prob) < self.threshold).float()
        # print('gate_mask_shape: ', np.shape(gate_mask))
        
        # 64, 4, 4---(-1) 维指的是上一层4个cell对于下一层的第一个cell的权重
        all_path_prob = torch.stack(path_prob, dim=2)  
        # 4*4 其中第0行代表前边4个cell对下一层第0个cell的输入
        # print('all_path_prob1: ', all_path_prob)
        # exit()     
        # 对即将输入cell0的上层4种cell的权重做了归一化
        all_path_prob = all_path_prob / (all_path_prob.sum(dim=-1, keepdim=True) + self.eps)
        # print('all_path_prob2: ', all_path_prob)
        # all_path_prob[:, :, i]:  torch.Size([64, 4])--前边一层4个cell对下一层第i个cell的概率
        # print('all_path_prob[:, :, i]: ', np.shape(all_path_prob[:, :, 1]))
        # 转化成上一层每条path对于下一层的4个cell的权重
        path_prob = [all_path_prob[:, :, i] for i in range(all_path_prob.size(2))] 
        # print('path_prob: ', path_prob)
        # exit()
        # print('path_prob_shape: ', np.shape(path_prob.data.cpu().numpy()))
        # print(path_prob)
        # exit()
        # emb_lst[0]_shape: torch.Size([64, 6, 300])
        # emb_lst[2]_shape: torch.Size([64, 6, 300])
        # emb_lst[1]_shape: torch.Size([64, 6, 300])
        # emb_lst[3]_shape: torch.Size([64, 6, 300])
        # path_prob[0]_shape: torch.Size([64, 4])
        # path_prob[2]_shape: torch.Size([64, 4])
        # path_prob[1]_shape: torch.Size([64, 4])
        # path_prob[3]_shape: torch.Size([64, 4])       
        # print('emb_lst[0]_shape:', np.shape(emb_lst[0]))
        # print('emb_lst[2]_shape:', np.shape(emb_lst[1]))
        # print('emb_lst[1]_shape:', np.shape(emb_lst[2]))
        # print('emb_lst[3]_shape:', np.shape(emb_lst[3]))
        # print('path_prob[0]_shape:', np.shape(path_prob[0]))
        # print('path_prob[2]_shape:', np.shape(path_prob[1]))
        # print('path_prob[1]_shape:', np.shape(path_prob[2]))
        # print('path_prob[3]_shape:', np.shape(path_prob[3]))
        # exit()
        # emb_lst[3], path_prob[3] = self.cmrc(repre_init)

        aggr_res_lst = []
        for i in range(self.num_out_path):
            # skip_emb:  torch.Size([64, 6, 300])
            # cur_path:  torch.Size([64, 1, 1, 1])
            skip_emb = unsqueeze2d(gate_mask[:, i]) * emb_lst[0]
            # print('skip_emb: ', skip_emb)
            # print('skip_emb: ', np.shape(skip_emb))
            res = 0
            # j对应上一次四个cell的输出；对应权重相加
            for j in range(self.num_cell):
                # cur_path:  torch.Size([64, 1, 1, 1])
                # cur_path = unsqueeze3d(path_prob[j][:, i])
                cur_path = unsqueeze2d(path_prob[j][:, i])
                # print('cur_path: ', np.shape(cur_path))
                # exit()
                cur_emb = emb_lst[j]
                # if emb_lst[j].dim() == 3:
                #     cur_emb = emb_lst[j].unsqueeze(1)
                # else:   # 4
                #     cur_emb = emb_lst[j]
                res = res + cur_path * cur_emb
            # res = res + skip_emb.unsqueeze(1)
            res = res + skip_emb
            aggr_res_lst.append(res)
        # print('aggr_res_lst_shape: ', np.shape(aggr_res_lst.data.cpu().numpy()))
        # print('all_path_prob_shape: ', np.shape(all_path_prob.data.cpu().numpy()))

        # exit()
        return aggr_res_lst, all_path_prob

    # def forward_i2t(self, repre_init, know2repre_food, know2repre_sightseeing, know2repre_mall, know2repre_nightlife, know2repre_hotel):
    #     path_prob = [None] * self.num_cell
    #     emb_lst = [None] * self.num_cell
    #     # emb_lst[0]: torch.Size([64, 6, 300]) ==更新后的表征
    #     # path_prob[0]: torch.Size([64, 4]) ==cell0对于下一层四个cell四条path的概率 
    #     emb_lst[0], path_prob[0] = self.ric(repre_init)
    #     emb_lst[1], path_prob[1] = self.imrc(repre_init)
    #     emb_lst[2], path_prob[2] = self.glgc(repre_init)
    #     emb_lst[3], path_prob[3] = self.cmrc(repre_init, know2repre_food, know2repre_sightseeing, know2repre_mall, know2repre_nightlife, know2repre_hotel)

    #     # print('path_prob[0]: ', path_prob[0])
    #     # print('path_prob[1]: ', path_prob[1])
    #     # print('path_prob[2]: ', path_prob[2])
    #     # print('path_prob[3]: ', path_prob[3])

    #     # gate_mask_shape:  torch.Size([64, 4])
    #     # all_path_prob_shape:  torch.Size([64, 4, 4])
    #     # all_path_prob_shape:  torch.Size([64, 4, 4])
    #     gate_mask = (sum(path_prob) < self.threshold).float()
    #     # print('gate_mask_shape: ', np.shape(gate_mask))
        
    #     # 64, 4, 4---(-1) 维指的是上一层4个cell对于下一层的第一个cell的权重
    #     all_path_prob = torch.stack(path_prob, dim=2)  
    #     # 4*4 其中第0行代表前边4个cell对下一层第0个cell的输入
    #     # print('all_path_prob1: ', all_path_prob)
    #     # exit()     
    #     # 对即将输入cell0的上层4种cell的权重做了归一化
    #     all_path_prob = all_path_prob / (all_path_prob.sum(dim=-1, keepdim=True) + self.eps)
    #     # print('all_path_prob2: ', all_path_prob)
    #     # all_path_prob[:, :, i]:  torch.Size([64, 4])--前边一层4个cell对下一层第i个cell的概率
    #     # print('all_path_prob[:, :, i]: ', np.shape(all_path_prob[:, :, 1]))
    #     # 转化成上一层每条path对于下一层的4个cell的权重
    #     path_prob = [all_path_prob[:, :, i] for i in range(all_path_prob.size(2))] 
    #     # print('path_prob: ', path_prob)
    #     # exit()
    #     # print('path_prob_shape: ', np.shape(path_prob.data.cpu().numpy()))
    #     # print(path_prob)
    #     # exit()
    #     # emb_lst[0]_shape: torch.Size([64, 6, 300])
    #     # emb_lst[2]_shape: torch.Size([64, 6, 300])
    #     # emb_lst[1]_shape: torch.Size([64, 6, 300])
    #     # emb_lst[3]_shape: torch.Size([64, 6, 300])
    #     # path_prob[0]_shape: torch.Size([64, 4])
    #     # path_prob[2]_shape: torch.Size([64, 4])
    #     # path_prob[1]_shape: torch.Size([64, 4])
    #     # path_prob[3]_shape: torch.Size([64, 4])       
    #     # print('emb_lst[0]_shape:', np.shape(emb_lst[0]))
    #     # print('emb_lst[2]_shape:', np.shape(emb_lst[1]))
    #     # print('emb_lst[1]_shape:', np.shape(emb_lst[2]))
    #     # print('emb_lst[3]_shape:', np.shape(emb_lst[3]))
    #     # print('path_prob[0]_shape:', np.shape(path_prob[0]))
    #     # print('path_prob[2]_shape:', np.shape(path_prob[1]))
    #     # print('path_prob[1]_shape:', np.shape(path_prob[2]))
    #     # print('path_prob[3]_shape:', np.shape(path_prob[3]))
    #     # exit()
    #     # emb_lst[3], path_prob[3] = self.cmrc(repre_init)

    #     aggr_res_lst = []
    #     for i in range(self.num_out_path):
    #         # skip_emb:  torch.Size([64, 6, 300])
    #         # cur_path:  torch.Size([64, 1, 1, 1])
    #         skip_emb = unsqueeze2d(gate_mask[:, i]) * emb_lst[0]
    #         # print('skip_emb: ', skip_emb)
    #         # print('skip_emb: ', np.shape(skip_emb))
    #         res = 0
    #         # j对应上一次四个cell的输出；对应权重相加
    #         for j in range(self.num_cell):
    #             # cur_path:  torch.Size([64, 1, 1, 1])
    #             cur_path = unsqueeze3d(path_prob[j][:, i])
    #             # print('cur_path: ', np.shape(cur_path))
    #             # exit()
    #             if emb_lst[j].dim() == 3:
    #                 cur_emb = emb_lst[j].unsqueeze(1)
    #             else:   # 4
    #                 cur_emb = emb_lst[j]
    #             res = res + cur_path * cur_emb
    #         res = res + skip_emb.unsqueeze(1)
    #         aggr_res_lst.append(res)
    #     # print('aggr_res_lst_shape: ', np.shape(aggr_res_lst.data.cpu().numpy()))
    #     # print('all_path_prob_shape: ', np.shape(all_path_prob.data.cpu().numpy()))

    #     # exit()
    #     return aggr_res_lst, all_path_prob




    def forward_t2i(self, rgn, img, wrd, stc, stc_lens):
        path_prob = [None] * self.num_cell
        emb_lst = [None] * self.num_cell
        emb_lst[0], path_prob[0] = self.ric(wrd)
        emb_lst[1], path_prob[1] = self.glgc(rgn, img, wrd, stc, stc_lens)
        emb_lst[2], path_prob[2] = self.imrc(wrd, stc_lens=stc_lens)
        emb_lst[3], path_prob[3] = self.cmrc(rgn, img, wrd, stc, stc_lens)

        gate_mask = (sum(path_prob) < self.threshold).float() 
        all_path_prob = torch.stack(path_prob, dim=2)  
        all_path_prob = all_path_prob / (all_path_prob.sum(dim=-1, keepdim=True) + self.eps)
        path_prob = [all_path_prob[:, :, i] for i in range(all_path_prob.size(2))]

        aggr_res_lst = []
        for i in range(self.num_out_path):
            skip_emb = unsqueeze2d(gate_mask[:, i]) * emb_lst[0] 
            res = 0
            for j in range(self.num_cell):
                cur_path = unsqueeze2d(path_prob[j][:, i]).unsqueeze(0) 
                if emb_lst[j].dim() == 3:
                    cur_emb = emb_lst[j].unsqueeze(0)
                else:  
                    cur_emb = emb_lst[j]
                res = res + cur_path * cur_emb
            res = res + skip_emb.unsqueeze(0)
            aggr_res_lst.append(res)

        return aggr_res_lst, all_path_prob

class DynamicInteraction_Layer(nn.Module):
    def __init__(self, opt, num_cell, num_out_path):
        super(DynamicInteraction_Layer, self).__init__()
        # self.opt = opt
        # self.threshold = 0.0001
        # self.eps = 1e-8
        # self.num_cell = num_cell
        # self.num_out_path = num_out_path

        # self.ric = RectifiedIdentityCell(opt, num_out_path)
        # self.glgc = GlobalLocalGuidanceCell(opt, num_out_path)
        # self.imrc = IntraModelReasoningCell(opt, num_out_path)
        # self.cmrc = CrossModalRefinementCell(opt, num_out_path)
        self.opt = opt
        self.threshold = 0.0001
        self.eps = 1e-8
        self.num_cell = num_cell
        self.num_out_path = num_out_path
        self.ric = RectifiedIdentityCell(opt, num_out_path)
        self.imrc = IntraModelReasoningCell(opt, num_out_path)
        self.glgc = GlobalLocalGuidanceCell(opt, num_out_path)
        self.cmrc = CrossModalRefinementCell(opt, num_out_path)
        # triplet_vec1 = torch.tensor(triplet_vec)
        self.temperature = 1.0

    def forward(self, pairs_emb_lst, know2repre_food, know2repre_sightseeing, know2repre_mall, know2repre_nightlife, know2repre_hotel):
        # if self.opt.direction == 'i2t':
        # aggr_res_lst = self.forward_i2t(ref_emb, rgn, img, wrd, stc, stc_lens)
        aggr_res_lst, all_path_prob = self.forward_i2t(pairs_emb_lst, know2repre_food, know2repre_sightseeing, know2repre_mall, know2repre_nightlife, know2repre_hotel)

        # else:
        #     aggr_res_lst = self.forward_t2i(ref_emb, rgn, img, wrd, stc, stc_lens)
        
        return aggr_res_lst, all_path_prob

    def forward_i2t(self, ref_rgn, know2repre_food, know2repre_sightseeing, know2repre_mall, know2repre_nightlife, know2repre_hotel):
        # assert len(ref_rgn) == self.num_cell and ref_rgn[0].dim() == 4
        # 64, 1, 6, 300
        # ref_rgn = torch.squeeze(ref_rgn)
        # print('ref_rgn_len: ', len(ref_rgn))
        path_prob = [None] * self.num_cell
        emb_lst = [None] * self.num_cell

        # ref_rgn[0]_shape:  torch.Size([64, 6, 300])
        # ref_rgn[1]_shape:  torch.Size([64, 6, 300])
        # ref_rgn[2]_shape:  torch.Size([64, 6, 300])
        # ref_rgn[3]_shape:  torch.Size([64, 6, 300])
        # print('ref_rgn[0]_shape: ', np.shape(ref_rgn[0]))
        # print('ref_rgn[1]_shape: ', np.shape(ref_rgn[1]))
        # print('ref_rgn[2]_shape: ', np.shape(ref_rgn[2]))
        # print('ref_rgn[3]_shape: ', np.shape(ref_rgn[3]))
        # exit()

        emb_lst[0], path_prob[0] = self.ric(ref_rgn[0])
        emb_lst[1], path_prob[1] = self.imrc(ref_rgn[1])
        emb_lst[2], path_prob[2] = self.glgc(ref_rgn[2])
        emb_lst[3], path_prob[3] = self.cmrc(ref_rgn[3], know2repre_food, know2repre_sightseeing, know2repre_mall, know2repre_nightlife, know2repre_hotel)
        # exit()
        if self.num_out_path == 1:
            aggr_res_lst = []
            gate_mask_lst = []
            res = 0
            for j in range(self.num_cell):
                gate_mask = (path_prob[j] < self.threshold / self.num_cell).float() 
                # print('gate_mask: ', np.shape(gate_mask))
                # print('ref_rgn[j]: ', np.shape(ref_rgn[j]))

                gate_mask_lst.append(gate_mask)
                skip_emb = gate_mask.unsqueeze(-1) * ref_rgn[j]
                res += path_prob[j].unsqueeze(-1) * emb_lst[j]                
                # skip_emb = gate_mask.unsqueeze(-1) * ref_rgn[j]
                # res += path_prob[j].unsqueeze(-1) * emb_lst[j]
                res += skip_emb

            res = res / (sum(gate_mask_lst) + sum(path_prob)).unsqueeze(-1)
            # all_path_prob = torch.stack(path_prob, dim=3) 
            all_path_prob = torch.stack(path_prob, dim=2) 
            aggr_res_lst.append(res)
        else:
            gate_mask = (sum(path_prob) < self.threshold).float()  
            # all_path_prob = torch.stack(path_prob, dim=3)   
            all_path_prob = torch.stack(path_prob, dim=2)   

            
            all_path_prob = all_path_prob / (all_path_prob.sum(dim=-1, keepdim=True) + self.eps)
            path_prob = [all_path_prob[:, :, i] for i in range(all_path_prob.size(2))]
            aggr_res_lst = []
            for i in range(self.num_out_path):
                skip_emb = unsqueeze2d(gate_mask[:, i]) * emb_lst[0]
                res = 0
                for j in range(self.num_cell):
                    cur_path = unsqueeze2d(path_prob[j][:, i])
                    res = res + cur_path * emb_lst[j]
                res = res + skip_emb
                aggr_res_lst.append(res)

        return aggr_res_lst, all_path_prob

    def forward_t2i(self, ref_wrd, rgn, img, wrd, stc, stc_lens):
        assert len(ref_wrd) == self.num_cell and ref_wrd[0].dim() == 4
        path_prob = [None] * self.num_cell
        emb_lst = [None] * self.num_cell
        emb_lst[0], path_prob[0] = self.ric(ref_wrd[0])
        emb_lst[1], path_prob[1] = self.glgc(rgn, img, ref_wrd[1], stc, stc_lens)
        emb_lst[2], path_prob[2] = self.imrc(ref_wrd[2], stc_lens)
        emb_lst[3], path_prob[3] = self.cmrc(rgn, img, ref_wrd[3], stc, stc_lens)
        
        if self.num_out_path == 1:
            aggr_res_lst = []
            gate_mask_lst = []
            res = 0
            for j in range(self.num_cell):
                gate_mask = (path_prob[j] < self.threshold / self.num_cell).float() 
                gate_mask_lst.append(gate_mask)
                skip_emb = gate_mask.unsqueeze(-1) * ref_wrd[j]
                res += path_prob[j].unsqueeze(-1) * emb_lst[j]
                res += skip_emb

            res = res / (sum(gate_mask_lst) + sum(path_prob)).unsqueeze(-1)
            all_path_prob = torch.stack(path_prob, dim=3)  
            aggr_res_lst.append(res)
        else:
            gate_mask = (sum(path_prob) < self.threshold).float()   
            all_path_prob = torch.stack(path_prob, dim=3)   
            all_path_prob = all_path_prob / (all_path_prob.sum(dim=-1, keepdim=True) + self.eps)
            path_prob = [all_path_prob[:, :, :, i] for i in range(all_path_prob.size(3))]

            aggr_res_lst = []
            for i in range(self.num_out_path):
                skip_emb = unsqueeze2d(gate_mask[:, :, i]) * emb_lst[0]
                res = 0
                for j in range(self.num_cell):
                    cur_path = unsqueeze2d(path_prob[j][:, :, i])
                    res = res + cur_path * emb_lst[j]
                res = res + skip_emb
                aggr_res_lst.append(res)

        return aggr_res_lst, all_path_prob

    


