import torch
import torch.nn as nn
from torch.nn.init import kaiming_normal_, constant_
from .util import conv, hbonet_block, pointwise_conv_no_relu, predict_flow, deconv, crop_like, depthwise_separable_conv, dds_conv, sandglass


__all__ = [
    'invertedresidualflownet', 'invertedresidualflownet_bn'
]


class InvertedResidualFlowNet(nn.Module):
    expansion = 1

    def __init__(self,batchNorm=True):
        super(InvertedResidualFlowNet,self).__init__()
        self.batchNorm = batchNorm

        self.conv1 = depthwise_separable_conv(self.batchNorm, 6, 48, stride=2)
        self.conv2 = hbonet_block(self.batchNorm, 48, 2, 96)
        self.conv3 = hbonet_block(self.batchNorm, 96, 6, 96)
        self.conv3_1 = hbonet_block(self.batchNorm, 96, 6, 96)
        self.conv4 = hbonet_block(self.batchNorm, 96, 6, 96)
        self.conv4_1 = hbonet_block(self.batchNorm, 96, 6, 96)
        self.conv5 = hbonet_block(self.batchNorm, 96, 6, 96)
        self.conv5_1 = hbonet_block(self.batchNorm, 96, 6, 96)

        self.upconv2 = depthwise_separable_conv(self.batchNorm, 96, 96)
        self.upconv3 = depthwise_separable_conv(self.batchNorm, 96, 96)
        self.upconv3_1 = depthwise_separable_conv(self.batchNorm, 96, 96)
        self.upconv4 = depthwise_separable_conv(self.batchNorm, 96, 96)
        self.upconv4_1 = depthwise_separable_conv(self.batchNorm, 96, 96)
        self.upconv5 = depthwise_separable_conv(self.batchNorm, 96, 96)
        self.upconv5_1 = depthwise_separable_conv(self.batchNorm, 96, 96)

        self.upsample = nn.Upsample(size=(61, 81), mode='bilinear', align_corners=True)

        self.predict_flow2 = predict_flow(96)

        for m in self.modules():
            if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
                kaiming_normal_(m.weight, 0.1)
                if m.bias is not None:
                    constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                constant_(m.weight, 1)
                constant_(m.bias, 0)

    def forward(self, x):
        out_conv1 = self.conv1(x)
        out_conv2 = self.upsample(self.upconv2(self.conv2(out_conv1)))
        out_conv3 = self.upsample(self.upconv3(self.conv3(out_conv2)))
        out_conv3_1 = self.upsample(self.upconv3_1(self.conv3_1(torch.add(out_conv2, out_conv3))))
        out_conv4 = self.upsample(self.upconv4(self.conv4(torch.add(out_conv3, out_conv3_1))))
        out_conv4_1 = self.upsample(self.upconv4_1(self.conv4_1(torch.add(out_conv3_1, out_conv4))))
        out_conv5 = self.upsample(self.upconv5(self.conv5(torch.add(out_conv4, out_conv4_1))))
        out_conv5_1 = self.upsample(self.upconv5_1(self.conv5_1(torch.add(out_conv4_1, out_conv5))))

        out_up2 = torch.add(out_conv5, out_conv5_1)

        flow2 = self.predict_flow2(out_up2)

        if self.training:
            return flow2#,flow3,flow4,flow5,flow6
        else:
            return flow2

    def weight_parameters(self):
        return [param for name, param in self.named_parameters() if 'weight' in name]

    def bias_parameters(self):
        return [param for name, param in self.named_parameters() if 'bias' in name]


def invertedresidualflownet(data=None):
    """FlowNetS model architecture from the
    "Learning Optical Flow with Convolutional Networks" paper (https://arxiv.org/abs/1504.06852)

    Args:
        data : pretrained weights of the network. will create a new one if not set
    """
    model = InvertedResidualFlowNet(batchNorm=False)
    if data is not None:
        model.load_state_dict(data['state_dict'])
    return model


def invertedresidualflownet_bn(data=None):
    """FlowNetS model architecture from the
    "Learning Optical Flow with Convolutional Networks" paper (https://arxiv.org/abs/1504.06852)

    Args:
        data : pretrained weights of the network. will create a new one if not set
    """
    model = InvertedResidualFlowNet(batchNorm=True)
    if data is not None:
        model.load_state_dict(data['state_dict'])
    return model

import torch
import torch.nn as nn
from torch.nn.init import kaiming_normal_, constant_
from .util import conv, hbonet_block, pointwise_conv_no_relu, predict_flow, deconv, crop_like, depthwise_separable_conv, dds_conv, sandglass


__all__ = [
    'invertedresidualflownet', 'invertedresidualflownet_bn'
]


class InvertedResidualFlowNet(nn.Module):
    expansion = 1

    def __init__(self,batchNorm=True):
        super(InvertedResidualFlowNet,self).__init__()
        self.batchNorm = batchNorm

        self.conv1 = depthwise_separable_conv(self.batchNorm, 6, 48, stride=2)
        self.conv2 = hbonet_block(self.batchNorm, 48, 2, 96)
        self.conv3 = hbonet_block(self.batchNorm, 96, 6, 96)
        self.conv3_1 = hbonet_block(self.batchNorm, 96, 6, 96)
        self.conv4 = hbonet_block(self.batchNorm, 96, 6, 96)
        self.conv4_1 = hbonet_block(self.batchNorm, 96, 6, 96)
        self.conv5 = hbonet_block(self.batchNorm, 96, 6, 96)
        self.conv5_1 = hbonet_block(self.batchNorm, 96, 6, 96)

        self.upconv2 = depthwise_separable_conv(self.batchNorm, 96, 96)
        self.upconv3 = depthwise_separable_conv(self.batchNorm, 96, 96)
        self.upconv3_1 = depthwise_separable_conv(self.batchNorm, 96, 96)
        self.upconv4 = depthwise_separable_conv(self.batchNorm, 96, 96)
        self.upconv4_1 = depthwise_separable_conv(self.batchNorm, 96, 96)
        self.upconv5 = depthwise_separable_conv(self.batchNorm, 96, 96)
        self.upconv5_1 = depthwise_separable_conv(self.batchNorm, 96, 96)

        self.upsample = nn.Upsample(size=(61, 81), mode='bilinear', align_corners=True)

        self.predict_flow2 = predict_flow(96)

        for m in self.modules():
            if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
                kaiming_normal_(m.weight, 0.1)
                if m.bias is not None:
                    constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                constant_(m.weight, 1)
                constant_(m.bias, 0)

    def forward(self, x):
        out_conv1 = self.conv1(x)
        out_conv2 = self.upsample(self.upconv2(self.conv2(out_conv1)))
        out_conv3 = self.upsample(self.upconv3(self.conv3(out_conv2)))
        out_conv3_1 = self.upsample(self.upconv3_1(self.conv3_1(torch.add(out_conv2, out_conv3))))
        out_conv4 = self.upsample(self.upconv4(self.conv4(torch.add(out_conv3, out_conv3_1))))
        out_conv4_1 = self.upsample(self.upconv4_1(self.conv4_1(torch.add(out_conv3_1, out_conv4))))
        out_conv5 = self.upsample(self.upconv5(self.conv5(torch.add(out_conv4, out_conv4_1))))
        out_conv5_1 = self.upsample(self.upconv5_1(self.conv5_1(torch.add(out_conv4_1, out_conv5))))

        out_up2 = torch.add(out_conv5, out_conv5_1)

        flow2 = self.predict_flow2(out_up2)

        if self.training:
            return flow2#,flow3,flow4,flow5,flow6
        else:
            return flow2

    def weight_parameters(self):
        return [param for name, param in self.named_parameters() if 'weight' in name]

    def bias_parameters(self):
        return [param for name, param in self.named_parameters() if 'bias' in name]


def invertedresidualflownet(data=None):
    """FlowNetS model architecture from the
    "Learning Optical Flow with Convolutional Networks" paper (https://arxiv.org/abs/1504.06852)

    Args:
        data : pretrained weights of the network. will create a new one if not set
    """
    model = InvertedResidualFlowNet(batchNorm=False)
    if data is not None:
        model.load_state_dict(data['state_dict'])
    return model


def invertedresidualflownet_bn(data=None):
    """FlowNetS model architecture from the
    "Learning Optical Flow with Convolutional Networks" paper (https://arxiv.org/abs/1504.06852)

    Args:
        data : pretrained weights of the network. will create a new one if not set
    """
    model = InvertedResidualFlowNet(batchNorm=True)
    if data is not None:
        model.load_state_dict(data['state_dict'])
    return model

#,flow3,flow4,flow5,flow6 FlowNetS model architecture from the     "Learning Optical Flow with Convolutional Networks" paper (https://arxiv.org/abs/1504.06852)      Args:         data : pretrained weights of the network. will create a new one if not set FlowNetS model architecture from the     "Learning Optical Flow with Convolutional Networks" paper (https://arxiv.org/abs/1504.06852)      Args:         data : pretrained weights of the network. will create a new one if not set