from ..builder import BACKBONES
from .convnext import ConvNeXt
from itertools import chain


@BACKBONES.register_module()
class ConvNeXtTextureMix(ConvNeXt):

    def __init__(self, style_stage_idx=0, content_stage_idx=2, *args, **kv):
        super(ConvNeXtTextureMix, self).__init__(*args, **kv)
        self.style_stage_idx = style_stage_idx
        self.content_stage_idx = content_stage_idx

    def freeze_all_stages(self):
        for i in range(self.num_stages):
            downsample_layer = self.downsample_layers[i]
            stage = self.stages[i]
            downsample_layer.eval()
            stage.eval()
            for param in chain(downsample_layer.parameters(),
                               stage.parameters()):
                param.requires_grad = False

    def unfreeze_all_stages(self):
        for i in range(self.num_stages):
            downsample_layer = self.downsample_layers[i]
            stage = self.stages[i]
            downsample_layer.train()
            stage.train()
            for param in chain(downsample_layer.parameters(),
                               stage.parameters()):
                param.requires_grad = True

    def forward(self, x, ret_texture_mix=False):
        outs = []
        style_feat = None
        content_feat = None
        for i, stage in enumerate(self.stages):
            x = self.downsample_layers[i](x)  # 每次进入stage之前，都进行一次下采样，进入0层进行的是stem block（4倍下采样，先down在norm），其他进行2倍下采样（先norm在down）
            x = stage(x)

            if ret_texture_mix:
                if i == self.style_stage_idx:
                    style_feat = x
                if i == self.content_stage_idx:
                    content_feat = x

            if i in self.out_indices:  # 如果是要输出的stage
                norm_layer = getattr(self, f'norm{i}')
                if self.gap_before_final_norm:  # 在最后的LN之前先进行GAP
                    gap = x.mean([-2, -1], keepdim=True)  # 对hw维度求平均，并保持4维张量形状
                    outs.append(norm_layer(gap).flatten(1))  # 进行LN，并将hw展平
                else:
                    outs.append(norm_layer(x))

        if ret_texture_mix:
            return tuple([style_feat, content_feat])
        else:
            return tuple(outs)
