import torch
import torch.nn as nn
import math
import torch.nn.functional as F
from torch.nn import init
from .dftm_v2.Dynamicfusion import DynamicFusion_Layer0, DynamicFusion_Layer
# from .dftm_v2.Dynamicfusion_norouter import DynamicFusion_Layer0, DynamicFusion_Layer

class DynamicFusionModule(nn.Module):
    def __init__(self, num_layer_routing=3, num_cells=4, embed_size=256):
        super(DynamicFusionModule, self).__init__()
        self.num_cells = num_cells = 4
        self.dynamic_fusion_l0 = DynamicFusion_Layer0(num_cells, num_cells, embed_size)
        self.dynamic_fusion_l1 = DynamicFusion_Layer(num_cells, num_cells, embed_size)
        self.dynamic_fusion_l2 = DynamicFusion_Layer(num_cells, 1, embed_size)

    
    def forward(self, x1,x2):
        pairs_emb_lst1 = self.dynamic_fusion_l0(x1,x2)
        pairs_emb_lst2 = self.dynamic_fusion_l1(pairs_emb_lst1)
        pairs_emb_lst3 = self.dynamic_fusion_l2(pairs_emb_lst2)
        
        feat1 = pairs_emb_lst3[0][0] #+ x1
        feat2 = pairs_emb_lst3[0][1] #+ x2
        fusion_feat = feat1 + feat2

        return fusion_feat

# # 消融1：去掉第二层
# class DynamicFusionModule(nn.Module):
#     def __init__(self, num_layer_routing=3, num_cells=4, embed_size=256):
#         super(DynamicFusionModule, self).__init__()
#         self.num_cells = num_cells = 4
#         self.dynamic_fusion_l0 = DynamicFusion_Layer0(num_cells, num_cells, embed_size)
#         self.dynamic_fusion_l2 = DynamicFusion_Layer(num_cells, 1, embed_size)

    
#     def forward(self, x1,x2,vis_weight=False):
#         #input_cat = torch.cat([input_feat_v,input_feat_i],1)
#         pairs_emb_lst1 = self.dynamic_fusion_l0(x1,x2,vis_weight)
#         pairs_emb_lst3 = self.dynamic_fusion_l2(pairs_emb_lst1,vis_weight)
        
#         feat1 = pairs_emb_lst3[0][0] #+ x1
#         feat2 = pairs_emb_lst3[0][1] #+ x2
#         fusion_feat = feat1 + feat2

#         return fusion_feat

# # 消融2：增加一层中间层
# class DynamicFusionModule(nn.Module):
#     def __init__(self, num_layer_routing=3, num_cells=4, embed_size=256):
#         super(DynamicFusionModule, self).__init__()
#         self.num_cells = num_cells = 4
#         self.dynamic_fusion_l0 = DynamicFusion_Layer0(num_cells, num_cells, embed_size)
#         self.dynamic_fusion_l1 = DynamicFusion_Layer(num_cells, num_cells, embed_size)
#         self.dynamic_fusion_l2 = DynamicFusion_Layer(num_cells, num_cells, embed_size)
#         self.dynamic_fusion_l3 = DynamicFusion_Layer(num_cells, 1, embed_size)

    
#     def forward(self, x1,x2,vis_weight=False):
#         #input_cat = torch.cat([input_feat_v,input_feat_i],1)
#         pairs_emb_lst1 = self.dynamic_fusion_l0(x1,x2,vis_weight)
#         pairs_emb_lst2 = self.dynamic_fusion_l1(pairs_emb_lst1,vis_weight)
#         pairs_emb_lst3 = self.dynamic_fusion_l2(pairs_emb_lst2,vis_weight)
#         pairs_emb_lst4 = self.dynamic_fusion_l3(pairs_emb_lst3,vis_weight)
        
#         feat1 = pairs_emb_lst4[0][0] #+ x1
#         feat2 = pairs_emb_lst4[0][1] #+ x2
#         fusion_feat = feat1 + feat2

#         return fusion_feat
    
# # 消融3：移除Router
# class DynamicFusionModule(nn.Module):
#     def __init__(self, num_layer_routing=3, num_cells=4, embed_size=256):
#         super(DynamicFusionModule, self).__init__()
#         self.num_cells = num_cells = 4
#         self.dynamic_fusion_l0 = DynamicFusion_Layer0(num_cells, num_cells, embed_size)
#         self.dynamic_fusion_l1 = DynamicFusion_Layer(num_cells, num_cells, embed_size)
#         self.dynamic_fusion_l2 = DynamicFusion_Layer(num_cells, 1, embed_size)

    
#     def forward(self, x1,x2,vis_weight=False):
#         #input_cat = torch.cat([input_feat_v,input_feat_i],1)
#         pairs_emb_lst1 = self.dynamic_fusion_l0(x1,x2,vis_weight)
#         pairs_emb_lst2 = self.dynamic_fusion_l1(pairs_emb_lst1,vis_weight)
#         pairs_emb_lst3 = self.dynamic_fusion_l2(pairs_emb_lst2,vis_weight)
        
#         feat1 = pairs_emb_lst3[0][0] #+ x1
#         feat2 = pairs_emb_lst3[0][1] #+ x2
#         fusion_feat = feat1 + feat2

#         return fusion_feat