import torch
import torch.nn as nn
import torch.nn.functional as F
import math
import numpy as np
from .SelfAttention import SelfAttention
from .Router import Router
from model.Refinement import Refinement


class RectifiedIdentityCell(nn.Module):
    def __init__(self, opt, num_out_path):
        super(RectifiedIdentityCell, self).__init__()
        self.keep_mapping = nn.ReLU()
        self.router = Router(num_out_path, opt.embed_size, opt.hid_router)

    def forward(self, x):
        path_prob = self.router(x)
        emb = self.keep_mapping(x)

        return emb, path_prob

class IntraModelReasoningCell(nn.Module):
    def __init__(self, opt, num_out_path):
        super(IntraModelReasoningCell, self).__init__()
        self.opt = opt
        self.router = Router(num_out_path, opt.embed_size, opt.hid_router)
        self.sa = SelfAttention(opt.embed_size, opt.hid_IMRC, opt.num_head_IMRC)

    def forward(self, inp, stc_lens=None):
        path_prob = self.router(inp)
        if inp.dim() == 4:
            n_img, n_stc, n_local, dim = inp.size()
            x = inp.view(-1, n_local, dim)
        else:
            x = inp

        sa_emb = self.sa(x)
        # exit()
        # if inp.dim() == 4:
        #     sa_emb = sa_emb.view(n_img, n_stc, n_local, -1)
        return sa_emb, path_prob

class CrossModalRefinementCell(nn.Module):
    def __init__(self, opt, num_out_path):
        super(CrossModalRefinementCell, self).__init__()
        self.direction = opt.direction
        # self.refine = Refinement(opt.embed_size, opt.raw_feature_norm_CMRC, opt.lambda_softmax_CMRC, opt.direction) 
        self.router = Router(num_out_path, opt.embed_size, opt.hid_router)

    def know_repre(self, repre_domain, know_repre):
        know_repre_tran = know_repre.transpose(1,0)
        context_know_score = torch.softmax(torch.matmul(repre_domain, know_repre_tran), dim=-1)
        context_know_repre = torch.matmul(context_know_score, know_repre) 
        return context_know_repre

    def forward(self, repre_init, know2repre_food, know2repre_sightseeing, know2repre_mall, know2repre_nightlife, know2repre_hotel):
        # if self.direction == 'i2t':
        # l_emb = rgn
        # else:
            # l_emb = wrd

        path_prob = self.router(repre_init)
        
        # repre_food:  torch.Size([64, 300])
        # repre_sightseeing:  torch.Size([64, 300])
        # repre_mall:  torch.Size([64, 300])
        # repre_nightlife:  torch.Size([64, 300])
        # repre_hotel:  torch.Size([64, 300])
        repre_food = repre_init[:,0,:].unsqueeze(1)
        repre_sightseeing = repre_init[:,1,:].unsqueeze(1)
        repre_mall = repre_init[:,2,:].unsqueeze(1)
        repre_nightlife = repre_init[:,3,:].unsqueeze(1)
        repre_hotel = repre_init[:,4,:].unsqueeze(1)
        repre_without_know = repre_init[:,5,:].unsqueeze(1)
        
        know_food_repre = self.know_repre(repre_food, know2repre_food)
        know_sightseeing_repre = self.know_repre(repre_sightseeing, know2repre_sightseeing)
        know_mall_repre = self.know_repre(repre_mall, know2repre_mall)
        know_nightlife_repre = self.know_repre(repre_nightlife, know2repre_nightlife)
        know_hotel_repre = self.know_repre(repre_hotel, know2repre_hotel)
        # paths = torch.cat([paths_l0, paths_l1, paths_l2], dim=-1) # (n_img, n_stc, total_paths)

        # repre_cat:  torch.Size([64, 6, 300])
        repre_cat = torch.cat([know_food_repre, know_sightseeing_repre, know_mall_repre, know_nightlife_repre, know_hotel_repre, repre_without_know], dim=1)

        # print('repre_cat: ', np.shape(repre_cat))
        # exit()
        # print('repre_food: ', np.shape(repre_food))
        # print('repre_sightseeing: ',  np.shape(repre_sightseeing))
        # print('repre_mall: ',  np.shape(repre_mall))
        # print('repre_nightlife: ',  np.shape(repre_nightlife))
        # print('repre_hotel: ',  np.shape(repre_hotel))
        # exit()

        
        # rf_pairs_emb = self.refine(rgn, img, wrd, stc, stc_lens)

        return repre_cat, path_prob

class GlobalLocalGuidanceCell(nn.Module):
    def __init__(self, opt, num_out_path):
        super(GlobalLocalGuidanceCell, self).__init__()
        self.opt = opt
        self.direction = self.opt.direction
        self.router = Router(num_out_path, opt.embed_size, opt.hid_router)
        self.fc_1 = nn.Linear(opt.embed_size, opt.embed_size)
        self.fc_2 = nn.Linear(opt.embed_size, opt.embed_size)

    def regulate(self, l_emb, g_emb_expand):
        l_emb_mid = self.fc_1(l_emb)
        x = l_emb_mid * g_emb_expand
        x = F.normalize(x, dim=-2)
        ref_l_emb = (1 + x) * l_emb
        return ref_l_emb

    def forward_i2t(self, repre_init):
        batch_size = repre_init.size(0)
        repre1 = self.fc_1(repre_init)
        repre2 = torch.transpose(self.fc_2(repre_init), 1, 2)
        score_pair = F.normalize(torch.matmul(repre1, repre2), dim=1)
        pair_weight_repre = torch.matmul(score_pair, repre_init)

        # origin_code==
        # n_img = rgn.size(0)
        # n_rgn = rgn.size(-2)
        # n_stc = stc.size(0) 
        # ref_rgns = []
        # for i in range(n_stc):
        #     if rgn.dim() == 4:
        #         query = rgn[:, i, :, :]
        #     else:
        #         query = rgn
            
        #     stc_i = stc[i].unsqueeze(0).unsqueeze(1).contiguous()
        #     stc_i_expand = stc_i.expand(n_img, n_rgn, -1)
        #     ref_rgn = self.regulate(query, stc_i_expand)
        #     ref_rgn = ref_rgn.unsqueeze(1)
        #     ref_rgns.append(ref_rgn)
        
        # ref_rgns = torch.cat(ref_rgns, dim=1)  
        return pair_weight_repre

    def forward(self, repre_init):
        # if self.direction == 'i2t':
        path_prob = self.router(repre_init)

        ref_emb = self.forward_i2t(repre_init)
            
        # else:
        #     path_prob = self.router(wrd)

        #     ref_emb = self.forward_t2i(rgn, img, wrd, stc, stc_lens)
        return ref_emb, path_prob

    # def forward_t2i(self, rgn, img, wrd, stc, stc_lens):
    #     # print(rgn.size())
    #     n_img = rgn.size(0)
    #     n_rgn = rgn.size(-2)
    #     n_stc = stc.size(0) 
    #     n_wrd = wrd.size(-2)
    #     ref_wrds = []
    #     for i in range(n_stc):
    #         if wrd.dim() == 4:
    #             wrd_i = wrd[:, i, :, :]
    #             wrd_i_expand = wrd_i
    #         else:
    #             wrd_i = wrd[i]
    #             wrd_i_expand = wrd_i.unsqueeze(0).expand(n_img, -1, -1)
    #         img_expand = img.unsqueeze(1).expand(-1, n_wrd, -1)


    #         ref_wrd = self.regulate(wrd_i_expand, img_expand)
    #         ref_wrd = ref_wrd.unsqueeze(1)
    #         ref_wrds.append(ref_wrd)
        
    #     ref_wrds = torch.cat(ref_wrds, dim=1)   # (n_img, n_stc, n_wrd, d) 
    #     return ref_wrds





