import sys
# sys.path.append("/Users/momo/human_seg_train")
# print(sys.path)
import torch
import torch.nn as nn
import torch.nn.functional as F
from bodyseg.backbone.backbone import get_norm

class ConvBNReLU(nn.Sequential):
    def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, groups=1, norm_layer=None):
        padding = (kernel_size - 1) // 2
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        super(ConvBNReLU, self).__init__(
            nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, groups=groups, bias=False),
            norm_layer(out_planes),
            nn.ReLU6(inplace=True)
        )

class InvertedResidual(nn.Module):
    def __init__(self, inp, oup, stride, expand_ratio, norm_layer=None):
        super(InvertedResidual, self).__init__()
        self.stride = stride
        assert stride in [1, 2]

        if norm_layer is None:
            norm_layer = nn.BatchNorm2d

        hidden_dim = int(round(inp * expand_ratio))
        self.use_res_connect = self.stride == 1 and inp == oup

        layers = []
        if expand_ratio != 1:
            # pw
            layers.append(ConvBNReLU(inp, hidden_dim, kernel_size=1, norm_layer=norm_layer))
        layers.extend([
            # dw
            ConvBNReLU(hidden_dim, hidden_dim, stride=stride, groups=hidden_dim, norm_layer=norm_layer),
            # pw-linear
            nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
            norm_layer(oup),
        ])
        self.conv = nn.Sequential(*layers)

    def forward(self, x):
        if self.use_res_connect:
            return x + self.conv(x)
        else:
            return self.conv(x)

class UpSampleBlock(nn.Module):
    def __init__(self, in_channels, out_channels, expand_ratio=6):
        super(UpSampleBlock, self).__init__()
        self.refine = InvertedResidual(in_channels, out_channels, 1, expand_ratio)

    def forward(self, x0, x1):
        x = torch.cat([x0, x1], dim=1)
        x = self.refine(x)
        return x

class BodySegNet_32_thin_sigmod(nn.Module):
    def __init__(self, input_channels=3, class_nums=1, output_onnx=False):
        super(BodySegNet_32_thin_sigmod, self).__init__()
        self.class_nums = class_nums
        self.output_onnx = output_onnx

        self.stage_1 = nn.Sequential(
            ConvBNReLU(input_channels, 16, kernel_size=3, stride=2)
        )
        self.stage_2 = nn.Sequential(
            ConvBNReLU(16, 16, kernel_size=3, stride=2, groups=16),
            ConvBNReLU(16, 16, kernel_size=1, stride=1),
        )
        self.stage_3 = nn.Sequential(
            InvertedResidual(16, 24, stride=2, expand_ratio=6),
            InvertedResidual(24, 24, stride=1, expand_ratio=6),
            InvertedResidual(24, 24, stride=1, expand_ratio=6),
        )
        self.stage_4 = nn.Sequential(
            InvertedResidual(24, 32, stride=2, expand_ratio=6),
            InvertedResidual(32, 32, stride=1, expand_ratio=6),
            InvertedResidual(32, 32, stride=1, expand_ratio=6),
            InvertedResidual(32, 32, stride=1, expand_ratio=6),
        )
        self.stage_5 = nn.Sequential(
            InvertedResidual(32, 48, stride=2, expand_ratio=6),
            InvertedResidual(48, 48, stride=1, expand_ratio=6),
            InvertedResidual(48, 48, stride=1, expand_ratio=6),
            InvertedResidual(48, 48, stride=1, expand_ratio=6)
        )
        self.up_to_4 = UpSampleBlock(48 + 32, 16)
        self.up_to_3 = UpSampleBlock(16 + 24, 16)
        self.up_to_2 = UpSampleBlock(16 + 16, 16)
        self.up_to_1 = UpSampleBlock(16 + 16, 16)
        self.last_layer = nn.Sequential(
            ConvBNReLU(16, 16, kernel_size=1, stride=1),
            nn.Conv2d(16, self.class_nums, kernel_size=1, stride=1)
        )
        self._initialize_weights()

    def forward(self, x):
        feature_S = []
        x1 = self.stage_1(x)
        x2 = self.stage_2(x1)
        x3 = self.stage_3(x2)
        x4 = self.stage_4(x3)
        feature = self.stage_5(x4)

        feature = F.interpolate(feature, size=x4.size()[2:], mode='bilinear', align_corners=True)
        feature = self.up_to_4(x4, feature)
        feature = F.interpolate(feature, size=x3.size()[2:], mode='bilinear', align_corners=True)
        feature = self.up_to_3(x3, feature)
        feature = F.interpolate(feature, size=x2.size()[2:], mode='bilinear', align_corners=True)
        feature = self.up_to_2(x2, feature)
        feature = F.interpolate(feature, size=x1.size()[2:], mode='bilinear', align_corners=True)
        feature = self.up_to_1(x1, feature)
        feature_S.append(feature)
        feature = self.last_layer(feature)
        feature_S.append(feature)
        output = F.interpolate(feature, size=x.size()[2:], mode='bilinear', align_corners=True)
        # output = torch.sigmoid(output)
        if self.output_onnx:
            output = torch.argmax(output, dim=1).to(torch.float32)

        return output, feature_S

    def _initialize_weights(self):
        for name, m in self.named_modules():
            if isinstance(m, nn.Conv2d):
                if 'first' in name:
                    nn.init.normal_(m.weight, 0, 0.01)
                else:
                    nn.init.normal_(m.weight, 0, 1.0 / m.weight.shape[1])
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0.0001)
                nn.init.constant_(m.running_mean, 0)
            elif isinstance(m, nn.BatchNorm1d):
                nn.init.constant_(m.weight, 1)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0.0001)
                nn.init.constant_(m.running_mean, 0)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)

class _ASPPModule(nn.Module):
    def __init__(self, inplanes, planes, kernel_size, padding, dilation, BatchNorm):
        super(_ASPPModule, self).__init__()
        self.atrous_conv = nn.Conv2d(inplanes, planes, kernel_size=kernel_size,
                                            stride=1, padding=padding, dilation=dilation, bias=False)
        self.bn = BatchNorm(planes)
        self.relu = nn.ReLU()

        self._init_weight()

    def forward(self, x):
        x = self.atrous_conv(x)
        x = self.bn(x)

        return self.relu(x)

    def _init_weight(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                torch.nn.init.kaiming_normal_(m.weight)

class ASPP(nn.Module):
    def __init__(self, backbone, output_stride, BatchNorm):
        super(ASPP, self).__init__()
        if backbone == 'drn':
            inplanes = 512
        elif backbone == 'mobilenet':
            inplanes = 320
        elif backbone == 'resnet_fpn':
            inplanes = 256
        else:
            inplanes = 2048
        if output_stride == 16:
            dilations = [1, 6, 12, 18]
        elif output_stride == 8:
            dilations = [1, 12, 24, 36]
        else:
            raise NotImplementedError

        self.aspp1 = _ASPPModule(inplanes, 256, 1, padding=0, dilation=dilations[0], BatchNorm=BatchNorm)
        self.aspp2 = _ASPPModule(inplanes, 256, 3, padding=dilations[1], dilation=dilations[1], BatchNorm=BatchNorm)
        self.aspp3 = _ASPPModule(inplanes, 256, 3, padding=dilations[2], dilation=dilations[2], BatchNorm=BatchNorm)
        self.aspp4 = _ASPPModule(inplanes, 256, 3, padding=dilations[3], dilation=dilations[3], BatchNorm=BatchNorm)

        self.global_avg_pool = nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)),
                                             nn.Conv2d(inplanes, 256, 1, stride=1, bias=False),
                                             BatchNorm(256),
                                             nn.ReLU())
        self.conv1 = nn.Conv2d(1280, 256, 1, bias=False)
        self.bn1 = BatchNorm(256)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.5)
        self._init_weight()

    def forward(self, x):
        x1 = self.aspp1(x)
        x2 = self.aspp2(x)
        x3 = self.aspp3(x)
        x4 = self.aspp4(x)
        x5 = self.global_avg_pool(x)
        x5 = F.interpolate(x5, size=x4.size()[2:], mode='bilinear', align_corners=True)
        # x5 = F.interpolate(x5, size=x4.size()[2:], mode='nearest')
        x = torch.cat((x1, x2, x3, x4, x5), dim=1)

        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)

        return self.dropout(x)

    def _init_weight(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                torch.nn.init.kaiming_normal_(m.weight)

class Decoder(nn.Module):
    def __init__(self,num_classes, backbone, BatchNorm):
        super(Decoder, self).__init__()
        if backbone == 'resnet_fpn' or backbone == 'drn':
            low_level_inplanes = 256
        elif backbone == 'xception':
            low_level_inplanes = 128
        elif backbone == 'mobilenet':
            low_level_inplanes = 24
        else:
            raise NotImplementedError

        self.conv1 = nn.Conv2d(low_level_inplanes, 16, 1, bias=False)
        self.bn1 = BatchNorm(16)
        self.relu = nn.ReLU()
        self.last_conv = nn.Sequential(nn.Conv2d(304, 256, kernel_size=3, stride=1, padding=1, bias=False),
                                       BatchNorm(256),
                                       nn.ReLU(),
                                       nn.Dropout(0.5),
                                       nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1, bias=False),
                                       BatchNorm(256),
                                       nn.ReLU(),
                                       nn.Dropout(0.1),
                                       nn.Conv2d(256, num_classes, kernel_size=1, stride=1))

        self.up_to_1 = UpSampleBlock(16 + 16, 16)
        self.last_layer = nn.Sequential(
            ConvBNReLU(16, 16, kernel_size=1, stride=1),
            nn.Conv2d(16, 1, kernel_size=1, stride=1)
        )
        self._init_weight()

    #  x(1,256,8,6) low(1,256,32,24) -》 x(1,1,32,24)
    def forward(self, x, low_level_feat):
        feature_T = []
        # deeplab part
        #(1,256,32,24) -> (1,16,64,48)
        low_level_feat = self.conv1(low_level_feat)
        low_level_feat = self.bn1(low_level_feat)
        low_level_feat = self.relu(low_level_feat)

        # x(1, 256, 8, 6)-> (1,16,32,24)
        x = F.interpolate(x, size=low_level_feat.size()[2:], mode='bilinear', align_corners=True)
        low_level_feat = F.interpolate(low_level_feat,
                                       size=[low_level_feat.size()[2] * 2, low_level_feat.size()[3] * 2],
                                       mode='bilinear', align_corners=True)
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)

        # bodyseg part
        feature = F.interpolate(x, size=[x.size()[2] * 2,x.size()[3] * 2], mode='bilinear', align_corners=True)
        # 输入up_to_1 x(low)(1,16,64,48),上采样2倍后的feature(1,16,64,48)
        feature = self.up_to_1(low_level_feat, feature)
        feature_T.append(feature)
        # (1,1,64,48)
        feature = self.last_layer(feature)
        feature_T.append(feature)
        # (1,1,128,96)
        output = F.interpolate(feature, size=[128,96], mode='bilinear', align_corners=True)
        # output = torch.sigmoid(output)

        return output, feature_T

    def _init_weight(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                torch.nn.init.kaiming_normal_(m.weight)

def build_backbone(backbone, output_stride, BatchNorm, input_channel=3):
    from bodyseg.backbone.fpn import build_resnet_fpn_backbone
    from bodyseg.backbone.xception import AlignedXception

    return build_resnet_fpn_backbone(input_channel)

def build_aspp(backbone, output_stride, BatchNorm):
    return ASPP(backbone, output_stride, BatchNorm)

def build_decoder(num_classes, backbone, BatchNorm):
    return Decoder(num_classes, backbone, BatchNorm)

class DeepLab(nn.Module):
    def __init__(self, input_channel=3, class_num=1):
        super(DeepLab, self).__init__()

        BatchNorm = get_norm("BN")

        self.backbone = build_backbone("resnet_fpn", 16, BatchNorm, input_channel=input_channel)
        self.aspp = build_aspp("resnet_fpn", 16, BatchNorm)
        self.decoder = build_decoder(class_num, "resnet_fpn", BatchNorm)

    def forward(self, input):
        # input(1,3,128,96)
        #output: "p2"(1,256,32,24), "p3"(1,256,16,12), "p4"(1,256,8,6)
        output = self.backbone(input)
        # "p4"(1,256,8,6)  "p2"(1,256,32,24)
        x, low_level_feat = output['p4'], output['p2']
        # x(1,256,8,6)-》(1,256,8,6)
        x = self.aspp(x)
        # x(1,256,8,6) low(1,256,32,24) -》 x(1,1,32,24)
        x, feature_T = self.decoder(x, low_level_feat)
        # x(1,1,128,96)
        x = F.interpolate(x, size=input.size()[2:], mode='bilinear', align_corners=True)
        # x = F.interpolate(x, size=input.size()[2:], mode='nearest')
        return x, feature_T