import torch
from torch import nn
import numpy as np
import torch.nn.functional as F
import torchvision.models as models
import torch.utils.model_zoo as model_zoo
from .depth_decoder import DepthDecoder


class feat_depth(nn.Module):
    def __init__(self, num_layers=50):
        super(feat_depth, self).__init__()

        self.feat_extractor = Feat_encoder(50)


        self.encoder = ResnetEncoder(50, False, 4)
        self.decoder = DepthDecoder(self.encoder.num_ch_enc)

 
    def forward(self, x, is_feat=False):
        feat_x = self.feat_extractor(x, is_feat)

        if is_feat:
            return feat_x
        #x = torch.cat((x, feat_x), dim=1)
        return self.decoder(self.encoder(feat_x))


class feat_full_depth(nn.Module):
    def __init__(self, num_layers=50):
        super(feat_full_depth, self).__init__()

        self.feat_extractor = Feat_full_encoder(50)  

        self.encoder = ResnetEncoder(num_layers, False, 6)
        self.decoder = DepthDecoder(self.encoder.num_ch_enc)

    def forward(self, x, is_feat=False):
        feat_x = self.feat_extractor(x, is_feat)

        if is_feat:
            return feat_x
        #x = torch.cat((x, feat_x), dim=1)
        return self.decoder(self.encoder(feat_x))


class feat_classifier(nn.Module):
    def __init__(self, num_layers=50):
        super(feat_classifier, self).__init__()

        self.feat_extractor = Feat_encoder(50)
        ckpt = torch.load('./checkpoints/vkitti2kitti_encoder2/best_net_G_Encoder.pth')
        self.feat_extractor.load_state_dict(ckpt, strict=False)

        for param in self.feat_extractor.parameters():
            param.requires_grad = False

        for param in self.feat_extractor.out.parameters():
            param.requires_grad = True

        self.encoder = ResnetEncoder(num_layers, False, 4)
        self.avgpool = nn.AdaptiveAvgPool2d((1,1))
        self.fc = nn.Linear(2048, 1000)

    def forward(self, x):
        x = self.feat_extractor(x)
        x = self.encoder(x)
        x = self.avgpool(x[-1]).flatten(1)
        out = self.fc(x)
        return out


class feat_full_classifier(nn.Module):
    def __init__(self, num_layers=50):
        super(feat_full_classifier, self).__init__()

        self.feat_extractor = Feat_full_encoder(50)
        ckpt = torch.load('./checkpoints/vkitti2kitti_encoder_full/25_net_G_Encoder.pth')
        self.feat_extractor.load_state_dict(ckpt, strict=False)

        for param in self.feat_extractor.parameters():
            param.requires_grad = False

        for param in self.feat_extractor.out.parameters():
            param.requires_grad = True
        for param in self.feat_extractor.feat2_out.parameters():
            param.requires_grad = True
        for param in self.feat_extractor.feat3_out.parameters():
            param.requires_grad = True
        for param in self.feat_extractor.feat4_out.parameters():
            param.requires_grad = True
        for param in self.feat_extractor.feat5_out.parameters():
            param.requires_grad = True

        self.encoder = ResnetEncoder(num_layers, False, 6)
        self.avgpool = nn.AdaptiveAvgPool2d((1,1))
        self.fc = nn.Linear(2048, 1000)

    def forward(self, x):
        x = self.feat_extractor(x)
        x = self.encoder(x)
        x = self.avgpool(x[-1]).flatten(1)
        out = self.fc(x)
        return out




class feat_pose(nn.Module):
    def __init__(self, num_layers=18):
        super(feat_pose, self).__init__()

        self.feat_extractor = Feat_full_encoder(50)
        ckpt = torch.load('./checkpoints/vkitti2kitti_encoder_full/25_net_G_Encoder.pth')

        if ckpt:
            ckpt['conv1.weight'] = torch.cat([ckpt['encoder.conv1.weight']] * 2, 1) / 2

        for param in self.feat_extractor.parameters():
            param.requires_grad = False

        for param in self.feat_extractor.out.parameters():
            param.requires_grad = True
        for param in self.feat_extractor.feat2_out.parameters():
            param.requires_grad = True
        for param in self.feat_extractor.feat3_out.parameters():
            param.requires_grad = True
        for param in self.feat_extractor.feat4_out.parameters():
            param.requires_grad = True
        for param in self.feat_extractor.feat5_out.parameters():
            param.requires_grad = True


        self.encoder = ResnetEncoder(num_layers, True, 6)
        self.decoder = PoseDecoder(self.encoder.num_ch_enc, 1, 2)

    def forward(self, x):

        x = self.feat_extractor(x, is_feat=False)
        x = self.encoder(x)

        x = self.decoder([x])

        return x




class Feat_encoder(nn.Module):
    def __init__(self, num_layers, pretrained_path=None, pose=False):
        super(Feat_encoder, self).__init__()

        self.num_ch_enc = np.array([64, 64, 128, 256, 512])

        resnets = {18: models.resnet18,
                   34: models.resnet34,
                   50: models.resnet50,
                   101: models.resnet101,
                   152: models.resnet152}

        if num_layers not in resnets:
            raise ValueError("{} is not a valid number of resnet layers".format(num_layers))

        if pose:
            self.encoder = resnet_multiimage_input(num_layers, num_input_images=2)
        else:
            self.encoder = resnets[num_layers](pretrained=True)

        if pretrained_path is not None:
            checkpoint = torch.load(pretrained_path)
            self.encoder.load_state_dict(checkpoint)

        if num_layers > 34:
            self.num_ch_enc[1:] *= 4

        self.out = nn.Sequential(nn.Conv2d(64, 12, 3, 1, 1), nn.ReLU())


    def forward(self, input_image, is_feat=False):
        

        input_image = (input_image - 0.45) / 0.225
        
        feat1 = self.encoder.relu(self.encoder.bn1(self.encoder.conv1(input_image)))
        feat1 = F.interpolate(feat1, scale_factor=2)

        if is_feat:
            return feat1

        out = self.out(feat1)
        
        return out




class Feat_full_encoder(nn.Module):
    def __init__(self, num_layers=50, pretrained_path=None):
        super(Feat_full_encoder, self).__init__()

        self.num_ch_enc = np.array([64, 64, 128, 256, 512])

        resnets = {18: models.resnet18,
                   34: models.resnet34,
                   50: models.resnet50,
                   101: models.resnet101,
                   152: models.resnet152}

        if num_layers not in resnets:
            raise ValueError("{} is not a valid number of resnet layers".format(num_layers))


        self.encoder = resnets[num_layers](pretrained=True)
        if pretrained_path is not None:
            checkpoint = torch.load(pretrained_path)
            self.encoder.load_state_dict(checkpoint)

        if num_layers > 34:
            self.num_ch_enc[1:] *= 4

        self.feat2_out = nn.Sequential(nn.Conv2d(256, 64, 3, 1, 1), nn.ReLU())
        self.feat3_out = nn.Sequential(nn.Conv2d(512, 64, 3, 1, 1), nn.ReLU())
        self.feat4_out = nn.Sequential(nn.Conv2d(1024, 64, 3, 1, 1), nn.ReLU())
        self.feat5_out = nn.Sequential(nn.Conv2d(2048, 64, 3, 1, 1), nn.ReLU())

        self.out = nn.Sequential(nn.Conv2d(64*5, 18, 3, 1, 1), nn.ReLU())


    def forward(self, input_image, is_feat):
        
        input_image = (input_image - 0.45) / 0.225

        self.features = []
        feat1 = self.encoder.relu(self.encoder.bn1(self.encoder.conv1(input_image)))

        if is_feat:
            return F.interpolate(feat1, scale_factor=2, mode="bilinear", align_corners=False)

        feat2 = self.encoder.layer1(self.encoder.maxpool(feat1))

        feat2_out = self.feat2_out(feat2)

        feat3 = self.encoder.layer2(feat2)

        feat3_out = self.feat3_out(feat3)

        feat4 = self.encoder.layer3(feat3)

        feat4_out = self.feat4_out(feat4)

        feat5 = self.encoder.layer4(feat4)

        feat5_out = self.feat5_out(feat5)

        feat = [feat1]
        feat.append(F.interpolate(feat2_out, scale_factor=2, mode="bilinear", align_corners=False))
        feat.append(F.interpolate(feat3_out, scale_factor=4, mode="bilinear", align_corners=False))
        feat.append(F.interpolate(feat4_out, scale_factor=8, mode="bilinear", align_corners=False))
        feat.append(F.interpolate(feat5_out, scale_factor=16, mode="bilinear", align_corners=False))

        feat = torch.cat(feat, dim=1)
        feat = F.interpolate(feat, scale_factor=2, mode="bilinear", align_corners=False)
        out = self.out(feat)
        return out




class Encoder(nn.Module):
    def __init__(self, num_layers, pretrained_path=None):
        super(Encoder, self).__init__()

        self.num_ch_enc = np.array([64, 64, 128, 256, 512])

        resnets = {18: models.resnet18,
                   34: models.resnet34,
                   50: models.resnet50,
                   101: models.resnet101,
                   152: models.resnet152}

        if num_layers not in resnets:
            raise ValueError("{} is not a valid number of resnet layers".format(num_layers))


        self.encoder = resnets[num_layers](pretrained=True)
        if pretrained_path is not None:
            checkpoint = torch.load(pretrained_path)
            self.encoder.load_state_dict(checkpoint)

        if num_layers > 34:
            self.num_ch_enc[1:] *= 4

        self.out1 = nn.Conv2d(64, 12, 3, 3, 1)
        
    def forward(self, input_image):
        
        outputs = []

        input_image = (input_image + 1.0)/2.0
        input_image = (input_image - 0.45) / 0.225

        self.features = []
        feat1 = self.encoder.relu(self.encoder.bn1(self.encoder.conv1(input_image)))
        feat2 = self.encoder.layer1(self.encoder.maxpool(feat1))
        feat3 = self.encoder.layer2(feat2)
        feat4 = self.encoder.layer3(feat3)
        feat5 = self.encoder.layer4(feat4)

        outputs.append(feat1)
        outputs.append(feat2)
        outputs.append(feat3)
        outputs.append(feat4)
        outputs.append(feat5)
        
        return outputs





class Encoderfull(nn.Module):
    def __init__(self, num_layers=50, pretrained_path=None):
        super(Encoderfull, self).__init__()

        self.num_ch_enc = np.array([64, 64, 128, 256, 512])

        resnets = {18: models.resnet18,
                   34: models.resnet34,
                   50: models.resnet50,
                   101: models.resnet101,
                   152: models.resnet152}

        if num_layers not in resnets:
            raise ValueError("{} is not a valid number of resnet layers".format(num_layers))


        self.encoder = resnets[num_layers](pretrained=True)
        if pretrained_path is not None:
            checkpoint = torch.load(pretrained_path)
            self.encoder.load_state_dict(checkpoint)

        if num_layers > 34:
            self.num_ch_enc[1:] *= 4

        self.feat2_out = nn.Sequential(nn.Conv2d(256, 64, 3, 3, 1), nn.ReLU())
        self.feat3_out = nn.Sequential(nn.Conv2d(512, 64, 3, 3, 1), nn.ReLU())
        self.feat4_out = nn.Sequential(nn.Conv2d(1024, 64, 3, 3, 1), nn.ReLU())
        self.feat5_out = nn.Sequential(nn.Conv2d(2048, 64, 3, 3, 1), nn.ReLU())

        self.out = nn.Sequential(nn.Conv2d(64, 12, 3, 3, 1), nn.Tanh())
        
    def forward(self, input_image):
        

        outputs = []

        input_image = (input_image + 1.0)/2.0
        input_image = (input_image - 0.45) / 0.225

        self.features = []
        feat1 = self.encoder.relu(self.encoder.bn1(self.encoder.conv1(input_image)))
        feat2 = self.encoder.layer1(self.encoder.maxpool(feat1))

        feat2_out = self.feat2_out(feat2)

        feat3 = self.encoder.layer2(feat2)

        feat3_out = self.feat3_out(feat3)

        feat4 = self.encoder.layer3(feat3)

        feat4_out = self.feat4_out(feat4)

        feat5 = self.encoder.layer4(feat4)

        feat5_out = self.feat5_out(feat5)

        outputs.append(feat1)
        outputs.append(feat2_out)
        outputs.append(feat3_out)
        outputs.append(feat4_out)
        outputs.append(feat5_out)
        outputs.append(feat5)

        
        return outputs




class Encoderfull3(nn.Module):
    def __init__(self, num_layers=50, pretrained_path=None):
        super(Encoderfull3, self).__init__()

        self.num_ch_enc = np.array([64, 64, 128, 256, 512])

        resnets = {18: models.resnet18,
                   34: models.resnet34,
                   50: models.resnet50,
                   101: models.resnet101,
                   152: models.resnet152}

        if num_layers not in resnets:
            raise ValueError("{} is not a valid number of resnet layers".format(num_layers))


        self.encoder = resnets[num_layers](pretrained=True)
        if pretrained_path is not None:
            checkpoint = torch.load(pretrained_path)
            self.encoder.load_state_dict(checkpoint)

        if num_layers > 34:
            self.num_ch_enc[1:] *= 4

        self.out = nn.Sequential(nn.Conv2d(64*5, 18, 3, 1, 1), nn.ReLU())
        self.feat2_out = nn.Sequential(nn.Conv2d(256, 64, 3, 1, 1), nn.ReLU())
        self.feat3_out = nn.Sequential(nn.Conv2d(512, 64, 3, 1, 1), nn.ReLU())
        self.feat4_out = nn.Sequential(nn.Conv2d(1024, 64, 3, 1, 1), nn.ReLU())
        self.feat5_out = nn.Sequential(nn.Conv2d(2048, 64, 3, 1, 1), nn.ReLU())
        
    def forward(self, input_image):
        

        outputs = []

        input_image = (input_image + 1.0)/2.0
        input_image = (input_image - 0.45) / 0.225

        self.features = []
        feat1 = self.encoder.relu(self.encoder.bn1(self.encoder.conv1(input_image)))
        feat2 = self.encoder.layer1(self.encoder.maxpool(feat1))
        feat3 = self.encoder.layer2(feat2)
        feat4 = self.encoder.layer3(feat3)
        feat5 = self.encoder.layer4(feat4)


        outputs.append(feat1)
        outputs.append(feat2)
        outputs.append(feat3)
        outputs.append(feat4)
        outputs.append(feat5)
        outputs.append(feat5)

        
        return outputs




class ResNetMultiImageInput(models.ResNet):
    """Constructs a resnet model with varying number of input images.
    Adapted from https://github.com/pytorch/vision/blob/master/torchvision/models/resnet.py
    """
    def __init__(self, block, layers, num_classes=1000, num_input_images=1):
        super(ResNetMultiImageInput, self).__init__(block, layers)
        self.inplanes = 64
        self.conv1 = nn.Conv2d(
            num_input_images * 3, 64, kernel_size=7, stride=2, padding=3, bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        self.layer1 = self._make_layer(block, 64, layers[0])
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)


def resnet_multiimage_input(num_layers, pretrained=False, num_input_images=1):
    """Constructs a ResNet model.
    Args:
        num_layers (int): Number of resnet layers. Must be 18 or 50
        pretrained (bool): If True, returns a model pre-trained on ImageNet
        num_input_images (int): Number of frames stacked as input
    """
    assert num_layers in [18, 50], "Can only run with 18 or 50 layer resnet"
    blocks = {18: [2, 2, 2, 2], 50: [3, 4, 6, 3]}[num_layers]
    block_type = {18: models.resnet.BasicBlock, 50: models.resnet.Bottleneck}[num_layers]
    model = ResNetMultiImageInput(block_type, blocks, num_input_images=num_input_images)

    if pretrained:
        loaded = model_zoo.load_url(models.resnet.model_urls['resnet{}'.format(num_layers)])
        loaded['conv1.weight'] = torch.cat(
            [loaded['conv1.weight']] * num_input_images, 1) / num_input_images
        model.load_state_dict(loaded)
    return model





class Decoder(nn.Module):
    def __init__(self, num_ch_enc, num_output_channels=3):
        super(Decoder, self).__init__()

        num_ch_dec = [16, 32, 64, 128, 256]

        # upconv
        self.upconv5 = ConvBlock(num_ch_enc[4], num_ch_dec[4])
        self.upconv4 = ConvBlock(num_ch_dec[4], num_ch_dec[3])
        self.upconv3 = ConvBlock(num_ch_dec[3], num_ch_dec[2])
        self.upconv2 = ConvBlock(num_ch_dec[2], num_ch_dec[1])
        self.upconv1 = ConvBlock(num_ch_dec[1], num_ch_dec[0])

        # iconv
        self.iconv5 = ConvBlock(num_ch_dec[4], num_ch_dec[4])
        self.iconv4 = ConvBlock(num_ch_dec[3], num_ch_dec[3])
        self.iconv3 = ConvBlock(num_ch_dec[2], num_ch_dec[2])
        self.iconv2 = ConvBlock(num_ch_dec[1], num_ch_dec[1])
        self.iconv1 = ConvBlock(num_ch_dec[0], num_ch_dec[0])

        # disp
        self.disp4 = Conv3x3(num_ch_dec[3], num_output_channels)
        self.disp3 = Conv3x3(num_ch_dec[2], num_output_channels)
        self.disp2 = Conv3x3(num_ch_dec[1], num_output_channels)
        self.disp1 = Conv3x3(num_ch_dec[0], num_output_channels)

        self.tanh = nn.Tanh()


    def forward(self, input_features, frame_id=0):
        self.outputs = {}
        econv5 = input_features
        # (64,64,128,256,512)*4

        upconv5 = upsample(self.upconv5(econv5))
        iconv5 = self.iconv5(upconv5)

        upconv4 = upsample(self.upconv4(iconv5))
        iconv4 = self.iconv4(upconv4)

        upconv3 = upsample(self.upconv3(iconv4))
        iconv3 = self.iconv3(upconv3)

        upconv2 = upsample(self.upconv2(iconv3))
        iconv2 = self.iconv2(upconv2)

        upconv1 = upsample(self.upconv1(iconv2))
        iconv1 = self.iconv1(upconv1)

        self.outputs[("recon", frame_id, 3)] = self.tanh(self.disp4(iconv4))
        self.outputs[("recon", frame_id, 2)] = self.tanh(self.disp3(iconv3))
        self.outputs[("recon", frame_id, 1)] = self.tanh(self.disp2(iconv2))
        self.outputs[("recon", frame_id, 0)] = self.tanh(self.disp1(iconv1))

        return self.outputs




class ResNetMultiImageInput(models.ResNet):
    """Constructs a resnet model with varying number of input images.
    Adapted from https://github.com/pytorch/vision/blob/master/torchvision/models/resnet.py
    """
    def __init__(self, block, layers, num_classes=1000, num_input_images=1):
        super(ResNetMultiImageInput, self).__init__(block, layers)
        self.inplanes = 64
        self.conv1 = nn.Conv2d(
            num_input_images * 3, 64, kernel_size=7, stride=2, padding=3, bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        self.layer1 = self._make_layer(block, 64, layers[0])
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)


def resnet_multiimage_input(num_layers, pretrained=False, num_input_images=1):
    """Constructs a ResNet model.
    Args:
        num_layers (int): Number of resnet layers. Must be 18 or 50
        pretrained (bool): If True, returns a model pre-trained on ImageNet
        num_input_images (int): Number of frames stacked as input
    """
    assert num_layers in [18, 50], "Can only run with 18 or 50 layer resnet"
    blocks = {18: [2, 2, 2, 2], 50: [3, 4, 6, 3]}[num_layers]
    block_type = {18: models.resnet.BasicBlock, 50: models.resnet.Bottleneck}[num_layers]
    model = ResNetMultiImageInput(block_type, blocks, num_input_images=num_input_images)

    if pretrained:
        loaded = model_zoo.load_url(models.resnet.model_urls['resnet{}'.format(num_layers)])
        loaded['conv1.weight'] = torch.cat(
            [loaded['conv1.weight']] * num_input_images, 1) / num_input_images
        model.load_state_dict(loaded)
    return model


class ResnetEncoder(nn.Module):
    """Pytorch module for a resnet encoder
    """
    def __init__(self, num_layers, pretrained, num_input_images=1):
        super(ResnetEncoder, self).__init__()

        self.num_ch_enc = np.array([64, 64, 128, 256, 512])

        resnets = {18: models.resnet18,
                   34: models.resnet34,
                   50: models.resnet50,
                   101: models.resnet101,
                   152: models.resnet152}
        

        if num_layers not in resnets:
            raise ValueError("{} is not a valid number of resnet layers".format(num_layers))

        if num_input_images > 1:
            self.encoder = resnet_multiimage_input(num_layers, pretrained, num_input_images)
        else:
            self.encoder = resnets[num_layers](pretrained)

        if num_layers > 34:
            self.num_ch_enc[1:] *= 4

    def forward(self, input_image):
        self.features = []

        x = self.encoder.conv1(input_image)
        x = self.encoder.bn1(x)
        self.features.append(self.encoder.relu(x))
        self.features.append(self.encoder.layer1(self.encoder.maxpool(self.features[-1])))
        self.features.append(self.encoder.layer2(self.features[-1]))
        self.features.append(self.encoder.layer3(self.features[-1]))
        self.features.append(self.encoder.layer4(self.features[-1]))

        return self.features



def upsample(x):
    return F.interpolate(x, scale_factor=2, mode="nearest")


class ConvBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(ConvBlock, self).__init__()
        self.conv = Conv3x3(in_channels, out_channels)
        self.nonlin = nn.ELU(inplace=True)
    def forward(self, x):
        out = self.conv(x)
        out = self.nonlin(out)
        return out


class Conv1x1(nn.Module):
    def __init__(self, in_channels, out_channels, bias=False):
        super(Conv1x1, self).__init__()
        self.conv = nn.Conv2d(int(in_channels), int(out_channels), kernel_size=1, stride=1, bias=bias)
    def forward(self, x):
        out = self.conv(x)
        return out


class Conv3x3(nn.Module):
    def __init__(self, in_channels, out_channels, use_refl=True):
        super(Conv3x3, self).__init__()
        if use_refl:
            self.pad = nn.ReflectionPad2d(1)
        else:
            self.pad = nn.ZeroPad2d(1)
        self.conv = nn.Conv2d(int(in_channels), int(out_channels), 3)
    def forward(self, x):
        out = self.pad(x)
        out = self.conv(out)
        return out


class Conv5x5(nn.Module):
    def __init__(self, in_channels, out_channels, use_refl=True):
        super(Conv5x5, self).__init__()
        if use_refl:
            self.pad = nn.ReflectionPad2d(2)
        else:
            self.pad = nn.ZeroPad2d(2)
        self.conv = nn.Conv2d(int(in_channels), int(out_channels), 5)
    def forward(self, x):
        out = self.pad(x)
        out = self.conv(out)
        return out


class CRPBlock(nn.Module):
    def __init__(self, in_planes, out_planes, n_stages):
        super(CRPBlock, self).__init__()
        for i in range(n_stages):
            setattr(self, '{}_{}'.format(i + 1, 'pointwise'), Conv1x1(in_planes if (i == 0) else out_planes, out_planes, False))
        self.stride = 1
        self.n_stages = n_stages
        self.maxpool = nn.MaxPool2d(kernel_size=5, stride=1, padding=2)

    def forward(self, x):
        top = x
        for i in range(self.n_stages):
            top = self.maxpool(top)
            top = getattr(self, '{}_{}'.format(i + 1, 'pointwise'))(top)
            x = top + x
        return x







import os.path as osp
import torch
import torch.nn as nn
from torch.nn import BatchNorm2d as bn
def conv3x3(in_planes, out_planes, stride=1):
    """3x3 convolution with padding"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False)


def conv1x1(in_planes, out_planes, stride=1):
    """1x1 convolution"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)


class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, inplanes, planes, stride=1, downsample=None):
        super(BasicBlock, self).__init__()
        self.conv1 = conv3x3(inplanes, planes, stride)
        self.bn1 = bn(planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3(planes, planes)
        self.bn2 = bn(planes)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        residual = x

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

        out = self.conv2(out)
        out = self.bn2(out)

        if self.downsample is not None:
            residual = self.downsample(x)

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

        return out


class Bottleneck(nn.Module):
    expansion = 4

    def __init__(self, inplanes, planes, stride=1, downsample=None):
        super(Bottleneck, self).__init__()
        self.conv1 = conv1x1(inplanes, planes)
        self.bn1 = bn(planes)
        self.conv2 = conv3x3(planes, planes, stride)
        self.bn2 = bn(planes)
        self.conv3 = conv1x1(planes, planes * self.expansion)
        self.bn3 = bn(planes * self.expansion)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        residual = x

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

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.bn3(out)

        if self.downsample is not None:
            residual = self.downsample(x)

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

        return out


class ResNet(nn.Module):

    def __init__(self, block, layers, num_classes=1000):
        super(ResNet, self).__init__()
        self.inplanes = 64
        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
        self.bn1 = bn(64)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        self.layer1 = self._make_layer(block, 64, layers[0])
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(512 * block.expansion, num_classes)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, bn):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

    def _make_layer(self, block, planes, blocks, stride=1):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                conv1x1(self.inplanes, planes * block.expansion, stride),
                bn(planes * block.expansion),
            )

        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample))
        self.inplanes = planes * block.expansion
        for _ in range(1, blocks):
            layers.append(block(self.inplanes, planes))

        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        return x


def resnet18(pretrained_path=None):
    """Constructs a ResNet-18 model.
    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    model = ResNet(BasicBlock, [2, 2, 2, 2])
    if pretrained_path is not None:
        model.load_state_dict(torch.load(pretrained_path))
        print('Loaded pre-trained weights')
    return model


def resnet34(pretrained_path=None, **kwargs):
    """Constructs a ResNet-34 model.
    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs)
    if pretrained_path is not None:
        model.load_state_dict(torch.load(osp.join(pretrained_path, 'resnet34.pth')))
        print('Loaded pre-trained weights')
    return model


def resnet50(pretrained_path=None, **kwargs):
    """Constructs a ResNet-50 model.
    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
    if pretrained_path is not None:
        model.load_state_dict(torch.load(osp.join(pretrained_path, 'resnet50.pth')))
        print('Loaded pre-trained weights')
    return model


def resnet101(pretrained_path=None, **kwargs):
    """Constructs a ResNet-101 model.
    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs)
    if pretrained_path is not None:
        model.load_state_dict(torch.load(osp.join(pretrained_path, 'resnet101.pth')))
        print('Loaded pre-trained weights')
    return model




class autoencoder(nn.Module):
    def __init__(self, options):
        super(autoencoder, self).__init__()
        self.opt = options

        self.Encoder = Encoder(50)
        #self.Decoder = Decoder(self.Encoder.num_ch_enc)

        self.ssim = SSIM()
        self.count = 0

    def forward(self, inputs):
        features, out = self.Encoder(inputs)
        #outputs = self.Decoder(features, 0)
        if self.training:
            loss_dict = self.compute_losses(inputs, features)
            return out, loss_dict
        return out

    def robust_l1(self, pred, target):
        eps = 1e-3
        return torch.sqrt(torch.pow(target - pred, 2) + eps ** 2)

    def compute_reprojection_loss(self, pred, target):
        photometric_loss = self.robust_l1(pred, target).mean(1, True)
        ssim_loss = self.ssim(pred, target).mean(1, True)
        reprojection_loss = (0.85 * ssim_loss + 0.15 * photometric_loss)
        return reprojection_loss

    def compute_losses(self, inputs, target, features):
        
        total_loss = 0

        for i in range(5):
            f=features[i]
            smooth_loss = self.get_smooth_loss(f, target)
            total_loss += smooth_loss/ (2 ** i)/5

        total_loss += self.robust_l1(inputs, target)

        return total_loss

    def get_smooth_loss(self, disp, img):
        b, _, h, w = disp.size()
        img = F.interpolate(img, (h, w), mode='area')

        disp_dx, disp_dy = self.gradient(disp)
        img_dx, img_dy = self.gradient(img)

        disp_dxx, disp_dxy = self.gradient(disp_dx)
        disp_dyx, disp_dyy = self.gradient(disp_dy)

        img_dxx, img_dxy = self.gradient(img_dx)
        img_dyx, img_dyy = self.gradient(img_dy)

        smooth1 = torch.mean(disp_dx.abs() * torch.exp(-img_dx.abs().mean(1, True))) + \
                  torch.mean(disp_dy.abs() * torch.exp(-img_dy.abs().mean(1, True)))

        smooth2 = torch.mean(disp_dxx.abs() * torch.exp(-img_dxx.abs().mean(1, True))) + \
                  torch.mean(disp_dxy.abs() * torch.exp(-img_dxy.abs().mean(1, True))) + \
                  torch.mean(disp_dyx.abs() * torch.exp(-img_dyx.abs().mean(1, True))) + \
                  torch.mean(disp_dyy.abs() * torch.exp(-img_dyy.abs().mean(1, True)))

        return -self.opt.dis * smooth1+ self.opt.cvt * smooth2

    def gradient(self, D):
        dy = D[:, :, 1:] - D[:, :, :-1]
        dx = D[:, :, :, 1:] - D[:, :, :, :-1]
        return dx, dy




