#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File    :   meta-sr.py    
@Contact :   fh.li@foxmail.com
@Modify Time  2020/4/9 10:46 AM  
'''
# Residual Dense Network for Image Super-Resolution
# https://arxiv.org/abs/1802.08797

# Residual Dense Network for Image Super-Resolution
# https://arxiv.org/abs/1802.08797

from model import common
import time
import torch
import torch.nn as nn
import math


def make_model(args, parent=False):
    return MetaRDN(args)


class RDB_Conv(nn.Module):
    def __init__(self, inChannels, growRate, kSize=3):
        super(RDB_Conv, self).__init__()
        Cin = inChannels
        G = growRate
        self.conv = nn.Sequential(*[
            nn.Conv2d(Cin, G, kSize, padding=(kSize - 1) // 2, stride=1),
            nn.ReLU()
        ])

    def forward(self, x):
        out = self.conv(x)
        return torch.cat((x, out), 1)


class RDB(nn.Module):
    def __init__(self, growRate0, growRate, nConvLayers, kSize=3):
        super(RDB, self).__init__()
        G0 = growRate0
        G = growRate
        C = nConvLayers

        convs = []
        for c in range(C):
            convs.append(RDB_Conv(G0 + c * G, G))
        self.convs = nn.Sequential(*convs)

        # Local Feature Fusion
        self.LFF = nn.Conv2d(G0 + C * G, G0, 1, padding=0, stride=1)

    def forward(self, x):
        return self.LFF(self.convs(x)) + x


class Pos2Weight(nn.Module):
    def __init__(self, inC, kernel_size=3, outC=3):
        super(Pos2Weight, self).__init__()
        self.inC = inC
        self.kernel_size = kernel_size
        self.outC = outC
        self.meta_block = nn.Sequential(
            nn.Linear(3, 256),
            nn.ReLU(inplace=True),
            nn.Linear(256, self.kernel_size * self.kernel_size * self.inC * self.outC)
        )

    def forward(self, x):
        output = self.meta_block(x)
        return output


class MetaRDN(nn.Module):
    def __init__(self, args):
        super(MetaRDN, self).__init__()
        r = args.scale[0]
        G0 = args.G0
        kSize = args.RDNkSize
        self.scale = 1
        self.args = args
        self.scale_idx = 0
        # number of RDB blocks, conv layers, out channels
        self.D, C, G = {
            'A': (20, 6, 32),
            'B': (16, 8, 64),
        }[args.RDNconfig]

        rgb_mean = (0.4488, 0.4371, 0.4040)
        rgb_std = (1.0, 1.0, 1.0)
        self.sub_mean = common.MeanShift(args.rgb_range, rgb_mean, rgb_std)
        self.add_mean = common.MeanShift(args.rgb_range, rgb_mean, rgb_std, 1)

        # Shallow feature extraction net
        self.SFENet1 = nn.Conv2d(args.n_colors, G0, kSize, padding=(kSize - 1) // 2, stride=1)
        self.SFENet2 = nn.Conv2d(G0, G0, kSize, padding=(kSize - 1) // 2, stride=1)

        # Redidual dense blocks and dense feature fusion
        self.RDBs = nn.ModuleList()
        for i in range(self.D):
            self.RDBs.append(
                RDB(growRate0=G0, growRate=G, nConvLayers=C)
            )
        # Global Feature Fusion
        self.GFF = nn.Sequential(*[
            nn.Conv2d(self.D * G0, G0, 1, padding=0, stride=1),
            nn.Conv2d(G0, G0, kSize, padding=(kSize - 1) // 2, stride=1)
        ])

        ## position to weight
        self.P2W = Pos2Weight(inC=G0)
        # self.weight_conv = nn.Conv2d(G0 + 2, 3, 3, 1, 1)

    def repeat_x(self, x):
        scale_int = math.ceil(self.scale)
        N, C, H, W = x.size()
        x = x.view(N, C, H, 1, W, 1)

        x = torch.cat([x] * scale_int, 3)
        x = torch.cat([x] * scale_int, 5).permute(0, 3, 5, 1, 2, 4)

        return x.contiguous().view(-1, C, H, W)

    def repeat_weight(self, weight, scale, inw, inh):
        k = int(math.sqrt(weight.size(0)))
        outw = inw * scale
        outh = inh * scale
        weight = weight.view(k, k, -1)
        scale_w = (outw + k - 1) // k
        scale_h = (outh + k - 1) // k
        weight = torch.cat([weight] * scale_h, 0)
        weight = torch.cat([weight] * scale_w, 1)

        weight = weight[0:outh, 0:outw, :]

        return weight

    def forward(self, x, scale=2):
        # # d1 =time.time()
        # print(x)
        # x = x * 255

        x = self.sub_mean(x)
        # print(x)
        f__1 = self.SFENet1(x)
        x = self.SFENet2(f__1)
        #
        RDBs_out = []
        for i in range(self.D):
            x = self.RDBs[i](x)
            RDBs_out.append(x)
        #
        x = self.GFF(torch.cat(RDBs_out, 1))
        x += f__1
        # print(x)

        # # local_weight = self.P2W(pos_mat.view(pos_mat.size(1), -1))  ###   (outH*outW, outC*inC*kernel_size*kernel_size)
        # # # print(d2)
        # # up_x = self.repeat_x(x)  ### the output is (N*r*r,inC,inH,inW)
        # #
        # # cols = nn.functional.unfold(up_x, 3, padding=1)
        # # scale_int = math.ceil(self.scale)
        # # local_weight = self.repeat_weight(local_weight, scale_int, x.size(2), x.size(3))
        # #
        # # cols = cols.contiguous().view(cols.size(0) // (scale_int ** 2), scale_int ** 2, cols.size(1), cols.size(2),
        # #                               1).permute(0, 1, 3, 4, 2).contiguous()
        # #
        # # local_weight = local_weight.contiguous().view(x.size(2), scale_int, x.size(3), scale_int, -1, 3).permute(1, 3,
        # #                                                                                                          0, 2,
        # #                                                                                                          4,
        # #                                                                                                          5).contiguous()
        # # local_weight = local_weight.contiguous().view(scale_int ** 2, x.size(2) * x.size(3), -1, 3)
        # #
        # # out = torch.matmul(cols, local_weight).permute(0, 1, 4, 2, 3)
        # # out = out.contiguous().view(x.size(0), scale_int, scale_int, 3, x.size(2), x.size(3)).permute(0, 3, 4, 1, 5, 2)
        # # out = out.contiguous().view(x.size(0), 3, scale_int * x.size(2), scale_int * x.size(3))
        # # out = self.add_mean(out)
        # up_x, up_input = upsampling(x,input_sub,scale)
        # # print(up_input.shape)
        # # print(up_x.shape)
        # # w_x = self.weight_conv(up_x)
        #
        # # x = up_input
        #
        # x = self.add_mean(x)
        # h, w = x.size()[2:]
        # hr_h, hr_w = int(scale * h), int(scale * w)
        # up_x = nn.functional.upsample(x, [hr_h, hr_w], mode='bilinear')

        return x

    def set_scale(self, scale_idx):
        self.scale_idx = scale_idx
        self.scale = self.args.scale[scale_idx]

def upsampling(x, input, scale=2):
    ###here is our meta learning upsampling function
    # the scale
    h, w = x.size()[2:]
    hr_h, hr_w = int(scale * h), int(scale * w)

    pos_h = torch.arange(0, hr_h / scale, 1 / scale).repeat(hr_w)
    pos_h = pos_h.view(1, 1, hr_w, hr_h).transpose(3, 2)

    pos_w = torch.arange(0, hr_w / scale, 1 / scale).repeat(hr_h)
    pos_w = pos_w.view(1, 1, hr_h, hr_w)

    pos = torch.cat((pos_h, pos_w), 1)

    int_pos = torch.floor(pos)
    res_pos = pos - int_pos
    # res_pos = res_pos.cuda()
    res_pos = res_pos.expand(x.size(0), res_pos.size(1), hr_h, hr_w)

    up_x = nn.functional.upsample(x, [hr_h, hr_w], mode='bilinear')
    up_input = nn.functional.upsample(input, [hr_h, hr_w], mode='bilinear')
    # print(up_x.size())
    # print(res_pos.size())
    pos_x = torch.cat((up_x, res_pos), 1)

    return pos_x, up_input