from __future__ import absolute_import, division, print_function

import numpy as np
import torch
import torch.nn as nn

from collections import OrderedDict
from layers import *
from .md2_swin import BasicLayer


class DepthDecoder(nn.Module):
    def __init__(self, num_ch_enc, scales=range(4), num_output_channels=1, use_skips=True):
        super(DepthDecoder, self).__init__()

        self.num_output_channels = num_output_channels
        self.use_skips = use_skips
        self.upsample_mode = 'nearest'
        self.scales = scales

        self.num_ch_enc = num_ch_enc
        self.num_ch_dec = np.array([16, 32, 64, 128, 256])

        # decoder
        self.convs = OrderedDict()
        for i in range(4, -1, -1):
            # upconv_0
            num_ch_in = self.num_ch_enc[-1] if i == 4 else self.num_ch_dec[i + 1]
            num_ch_out = self.num_ch_dec[i]
            self.convs[("upconv", i, 0)] = ConvBlock(num_ch_in, num_ch_out)

            # upconv_1
            num_ch_in = self.num_ch_dec[i]
            if self.use_skips and i > 0:
                num_ch_in += self.num_ch_enc[i - 1]
            num_ch_out = self.num_ch_dec[i]
            self.convs[("upconv", i, 1)] = ConvBlock(num_ch_in, num_ch_out)

        for s in self.scales:
            self.convs[("dispconv", s)] = Conv3x3(self.num_ch_dec[s], self.num_output_channels)

        self.decoder = nn.ModuleList(list(self.convs.values()))
        self.sigmoid = nn.Sigmoid()

    def forward(self, input_features):
        self.outputs = {}

        # decoder
        x = input_features[-1]
        for i in range(4, -1, -1):
            x = self.convs[("upconv", i, 0)](x)
            x = [upsample(x)]
            if self.use_skips and i > 0:
                x += [input_features[i - 1]]
            x = torch.cat(x, 1)
            x = self.convs[("upconv", i, 1)](x)
            if i in self.scales:
                self.outputs[("disp", i)] = self.sigmoid(self.convs[("dispconv", i)](x))

        return self.outputs





class CD_DispNet(nn.Module):
    def __init__(self, num_ch_enc, scales=range(4), num_output_channels=1, use_skips=True):
        super(CD_DispNet, self).__init__()

        self.num_output_channels = num_output_channels
        self.use_skips = use_skips
        self.upsample_mode = 'nearest'
        self.scales = scales

        self.num_ch_enc = num_ch_enc
        self.num_ch_dec = np.array([16, 32, 64, 128, 256])


        self.attn_feat1 = BasicLayer(16, (96, 320), 2, 2, 32)
        self.attn_feat2 = BasicLayer(32, (48, 160), 2, 4, 16)
        self.attn_feat3 = BasicLayer(64, (24, 80), 2, 8, 8)
        self.attn_feat4 = BasicLayer(128, (12, 40), 6, 16, 4)
        self.attn_feat5 = BasicLayer(256, (6, 20), 2, 32, 2)

        self.attn_real1 = BasicLayer(16, (96, 320), 2, 2, 32)
        self.attn_real2 = BasicLayer(32, (48, 160), 2, 4, 16)
        self.attn_real3 = BasicLayer(64, (24, 80), 2, 8, 8)
        self.attn_real4 = BasicLayer(128, (12, 40), 6, 16, 4)
        self.attn_real5 = BasicLayer(256, (6, 20), 2, 32, 2)        
        
        # decoder
        self.convs = OrderedDict()
        for i in range(4, -1, -1):
            # upconv_0
            num_ch_in = self.num_ch_enc[-1] if i == 4 else self.num_ch_dec[i + 1]
            num_ch_out = self.num_ch_dec[i]
            self.convs[("real_upconv", i, 0)] = ConvBlock(num_ch_in, num_ch_out)
            self.convs[("feat_upconv", i, 0)] = ConvBlock(num_ch_in, num_ch_out)

            # upconv_1
            num_ch_in = self.num_ch_dec[i]
            if self.use_skips and i > 0:
                num_ch_in += self.num_ch_enc[i - 1]
            num_ch_out = self.num_ch_dec[i]
            self.convs[("real_upconv", i, 1)] = ConvBlock(num_ch_in, num_ch_out)
            self.convs[("feat_upconv", i, 1)] = ConvBlock(num_ch_in, num_ch_out)

        for s in self.scales:
            self.convs[("real_dispconv", s)] = Conv3x3(self.num_ch_dec[s], self.num_output_channels)
            self.convs[("feat_dispconv", s)] = Conv3x3(self.num_ch_dec[s], self.num_output_channels)


        self.decoder = nn.ModuleList(list(self.convs.values()))
        self.sigmoid = nn.Sigmoid()

    def forward(self, feat_features, real_features):
        self.outputs_feat = {}
        self.outputs_real = {}

        # decoder
        attn_feat = [self.attn_feat1, self.attn_feat2, self.attn_feat3, self.attn_feat4, self.attn_feat5]
        attn_real = [self.attn_real1, self.attn_real2, self.attn_real3, self.attn_real4, self.attn_real5]

        x_feat = feat_features[-1]
        x_real = real_features[-1]

        for i in range(4, -1, -1):
            x_feat = self.convs[("feat_upconv", i, 0)](x_feat)
            x_real = self.convs[("real_upconv", i, 0)](x_real)

            
            x_feat_out = x_feat + self.attention(x_feat, x_real, attn_feat[i])
            x_real_out = x_real + self.attention(x_real, x_feat, attn_real[i])


            x_feat = x_feat_out
            x_real = x_real_out


            x_real = [upsample(x_real)]
            x_feat = [upsample(x_feat)]

            if self.use_skips and i > 0:
                x_feat += [feat_features[i - 1]]
                x_real += [real_features[i - 1]]

            x_feat = torch.cat(x_feat, 1)
            x_real = torch.cat(x_real, 1)

            x_feat = self.convs[("feat_upconv", i, 1)](x_feat)
            x_real = self.convs[("real_upconv", i, 1)](x_real)

            if i in self.scales:
                self.outputs_feat[("disp", i)] = self.sigmoid(self.convs[("feat_dispconv", i)](x_feat))
                self.outputs_real[("disp", i)] = self.sigmoid(self.convs[("real_dispconv", i)](x_real))
                
        return self.outputs_feat, self.outputs_real

    def attention(self, x, y, attn_fn):

        x = x.permute(0,2,3,1)
        y = y.permute(0,2,3,1)
        B, H, W, C = x.shape
        x = x.view(B, H*W, C)
        y = y.view(B, H*W, C)

        out = attn_fn(x, y)

        out = out.view((len(out), H, W, C))
        out = out.permute(0,3,1,2)
        
        return out
