import torch
from torch import nn
import torch.nn.functional as F
import horovod.torch as hvd


class ResBlock3d(nn.Module):
    def __init__(self, n_in, n_out, parameters, stride=1, ):
        super(ResBlock3d, self).__init__()
        self.conv1 = nn.Conv3d(n_in, n_out, kernel_size=3, stride=stride, padding=1)
        self.bn1 = nn.BatchNorm3d(n_out, **parameters)
        # self.bn1 = hvd.SyncBatchNorm(n_out, **parameters)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv3d(n_out, n_out, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm3d(n_out, **parameters)
        # self.bn2 = hvd.SyncBatchNorm(n_out, **parameters)

        if stride != 1 or n_out != n_in:
            self.shortcut = nn.Sequential(
                nn.Conv3d(n_in, n_out, kernel_size=1, stride=stride),
                nn.BatchNorm3d(n_out, **parameters))
            # hvd.SyncBatchNorm(n_out, **parameters))
        else:
            self.shortcut = None

    def forward(self, x):
        residual = x
        if self.shortcut is not None:
            residual = self.shortcut(x)
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.bn2(out)

        out += residual
        out = self.relu(out)
        return out


class ResUNet(nn.Module):
    def __init__(self, cfg, in_channels, out_channels):
        super(ResUNet, self).__init__()
        parameters = {
            "momentum": cfg.model['backbone']['bn_momentum'],
            "eps": 1e-3,
        }
        self.fpn = cfg.model['backbone']['fpn']
        self.preBlock = nn.Sequential(
            nn.Conv3d(in_channels, 24, kernel_size=3, padding=1, stride=2),
            nn.BatchNorm3d(24, **parameters),
            # hvd.SyncBatchNorm(24, **parameters),
            nn.ReLU(inplace=True),
            nn.Conv3d(24, 24, kernel_size=3, padding=1),
            nn.BatchNorm3d(24, **parameters),
            # hvd.SyncBatchNorm(24,**parameters),
            nn.ReLU(inplace=True))

        self.forw1 = nn.Sequential(
            ResBlock3d(24, 32, parameters),
            ResBlock3d(32, 32, parameters))

        self.forw2 = nn.Sequential(
            ResBlock3d(32, 64, parameters),
            ResBlock3d(64, 64, parameters))

        self.forw3 = nn.Sequential(
            ResBlock3d(64, 64, parameters),
            ResBlock3d(64, 64, parameters),
            ResBlock3d(64, 64, parameters))

        self.forw4 = nn.Sequential(
            ResBlock3d(64, 64, parameters),
            ResBlock3d(64, 64, parameters),
            ResBlock3d(64, 64, parameters))

        # skip connection in U-net
        self.back2 = nn.Sequential(
            # 64 + 64 + 3, where 3 is the channeld dimension of coord
            ResBlock3d(128 + 3, 128, parameters),
            ResBlock3d(128, 128, parameters),
            ResBlock3d(128, 128, parameters))

        # skip connection in U-net
        self.back3 = nn.Sequential(
            ResBlock3d(128, 64, parameters),
            ResBlock3d(64, 64, parameters),
            ResBlock3d(64, 64, parameters))

        self.maxpool1 = nn.MaxPool3d(kernel_size=2, stride=2,
                                     return_indices=True)
        self.maxpool2 = nn.MaxPool3d(kernel_size=2, stride=2,
                                     return_indices=True)
        self.maxpool3 = nn.MaxPool3d(kernel_size=2, stride=2,
                                     return_indices=True)
        self.maxpool4 = nn.MaxPool3d(kernel_size=2, stride=2,
                                     return_indices=True)

        # upsampling in U-net
        self.path1 = nn.Sequential(
            nn.ConvTranspose3d(64, 64, kernel_size=2, stride=2),
            nn.BatchNorm3d(64, **parameters),
            # hvd.SyncBatchNorm(64),
            nn.ReLU(inplace=True))

        # upsampling in U-net
        self.path2 = nn.Sequential(
            nn.ConvTranspose3d(64, 64, kernel_size=2, stride=2),
            nn.BatchNorm3d(64, **parameters),
            # hvd.SyncBatchNorm(64),
            nn.ReLU(inplace=True))

    # can add coord here
    def forward(self, x, coord):
        out = self.preBlock(x)  # 16
        out_pool = out
        out1 = self.forw1(out_pool)  # 32
        out1_pool, _ = self.maxpool2(out1)
        out2 = self.forw2(out1_pool)  # 64
        # out2 = self.drop(out2)
        out2_pool, _ = self.maxpool3(out2)
        out3 = self.forw3(out2_pool)  # 96
        out3_pool, _ = self.maxpool4(out3)
        out4 = self.forw4(out3_pool)  # 96
        # out4 = self.drop(out4)

        rev3 = self.path1(out4)
        comb3 = self.back3(torch.cat((rev3, out3), 1))  # 96+96
        rev2 = self.path2(comb3)
        comb2 = self.back2(torch.cat((rev2, out2, coord), 1))  # 64+64

        return [x, out1, comb2], comb2  # out2


class ResUNet_FPN(nn.Module):
    def __init__(self, cfg, in_channels, out_channels):
        super(ResUNet_FPN, self).__init__()
        parameters = {
            "momentum": cfg.model['backbone']['bn_momentum']
        }
        self.preBlock = nn.Sequential(
            # nn.Conv3d(in_channels, 24, kernel_size=3, padding=1, stride=2),
            nn.Conv3d(in_channels, 24, kernel_size=3, padding=1, stride=1),
            nn.BatchNorm3d(24, **parameters),
            # hvd.SyncBatchNorm(24, **parameters),
            nn.ReLU(inplace=True),
            nn.Conv3d(24, 24, kernel_size=3, padding=1),
            nn.BatchNorm3d(24, **parameters),
            # hvd.SyncBatchNorm(24, **parameters),
            nn.ReLU(inplace=True))

        self.forw1 = nn.Sequential(
            ResBlock3d(24, 32, parameters),
            ResBlock3d(32, 32, parameters))

        self.forw2 = nn.Sequential(
            ResBlock3d(32, 64, parameters),
            ResBlock3d(64, 64, parameters))

        self.forw3 = nn.Sequential(
            ResBlock3d(64, 64, parameters),
            ResBlock3d(64, 64, parameters),
            ResBlock3d(64, 64, parameters))

        self.forw4 = nn.Sequential(
            ResBlock3d(64, 64, parameters),
            ResBlock3d(64, 64, parameters),
            ResBlock3d(64, 64, parameters))

        # skip connection in U-net
        self.back1 = nn.Sequential(
            ResBlock3d(96, 128, parameters),
            ResBlock3d(128, 128, parameters),
            ResBlock3d(128, 128, parameters))

        # skip connection in U-net
        self.back2 = nn.Sequential(
            # 64 + 64 + 3, where 3 is the channeld dimension of coord
            ResBlock3d(128, 64, parameters),
            ResBlock3d(64, 64, parameters),
            ResBlock3d(64, 64, parameters))

        # skip connection in U-net
        self.back3 = nn.Sequential(
            ResBlock3d(128, 64, parameters),
            ResBlock3d(64, 64, parameters),
            ResBlock3d(64, 64, parameters))

        self.maxpool1 = nn.MaxPool3d(kernel_size=2, stride=2,
                                     return_indices=True)
        self.maxpool2 = nn.MaxPool3d(kernel_size=2, stride=2,
                                     return_indices=True)
        self.maxpool3 = nn.MaxPool3d(kernel_size=2, stride=2,
                                     return_indices=True)
        self.maxpool4 = nn.MaxPool3d(kernel_size=2, stride=2,
                                     return_indices=True)

        # upsampling in U-net
        self.path1 = nn.Sequential(
            nn.ConvTranspose3d(64, 64, kernel_size=2, stride=2),
            nn.BatchNorm3d(64),
            # hvd.SyncBatchNorm(64),
            nn.ReLU(inplace=True))

        # upsampling in U-net
        self.path2 = nn.Sequential(
            nn.ConvTranspose3d(64, 64, kernel_size=2, stride=2),
            nn.BatchNorm3d(64),
            # hvd.SyncBatchNorm(64),
            nn.ReLU(inplace=True))

        # upsampling in U - net
        self.path3 = nn.Sequential(
            nn.ConvTranspose3d(64, 64, kernel_size=2, stride=2),
            nn.BatchNorm3d(64),
            # hvd.SyncBatchNorm(64),
            nn.ReLU(inplace=True))

        # upsampling in U - net
        self.path4 = nn.Sequential(
            nn.ConvTranspose3d(64, 64, kernel_size=2, stride=2),
            nn.BatchNorm3d(64),
            # hvd.SyncBatchNorm(64),
            nn.ReLU(inplace=True))

    def forward(self, x):
        out = self.preBlock(x)
        # out_pool = out
        out_pool, _ = self.maxpool1(out)
        out1 = self.forw1(out_pool)
        out1_pool, _ = self.maxpool2(out1)
        out2 = self.forw2(out1_pool)
        # out2 = self.drop(out2)
        out2_pool, _ = self.maxpool3(out2)
        out3 = self.forw3(out2_pool)
        out3_pool, _ = self.maxpool4(out3)
        out4 = self.forw4(out3_pool)
        # out4 = self.drop(out4)

        rev3 = self.path1(out4)
        comb3 = self.back3(torch.cat((rev3, out3), 1))
        rev2 = self.path2(comb3)
        comb2 = self.back2(torch.cat((rev2, out2), 1))
        rev1 = self.path3(comb2)
        comb1 = self.back1(torch.cat((rev1, out1), 1))

        return [comb3, comb2, comb1], out2