'''
architecture for sftmd
'''
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from utils.util import PCAEncoder

class CALayer(nn.Module):
    def __init__(self, nf, reduction=16):
        super(CALayer, self).__init__()
        self.body = nn.Sequential(
            nn.Conv2d(nf, nf//reduction, 1, 1, 0),
            nn.LeakyReLU(0.2), 
            nn.Conv2d(nf//reduction, nf, 1, 1, 0),
            nn.Sigmoid()
        )
        self.avg = nn.AdaptiveAvgPool2d(1)
    
    def forward(self, x):
        y = self.avg(x)
        y = self.body(y)
        return torch.mul(x, y)

class CRB_Layer(nn.Module):
    def __init__(self, nf1, nf2):
        super(CRB_Layer, self).__init__()

        body = [
            nn.Conv2d(nf1+nf2, nf1+nf2, 3, 1, 1),
            nn.LeakyReLU(0.2, True),
            nn.Conv2d(nf1+nf2, nf1, 3, 1, 1), 
            CALayer(nf1)
        ]

        self.body = nn.Sequential(*body)

    def forward(self, x):
        f1, f2 = x
        f1 = self.body(torch.cat(x, 1)) + f1
        return [f1, f2]


class Estimator(nn.Module):
    def __init__(self, in_nc=3, nf=64, num_blocks=5, scale=4, kernel_size=4):
        super(Estimator, self).__init__()

        self.ksize = kernel_size

        self.head_LR = nn.Conv2d(in_nc, nf//2, 1, 1, 0)
        self.head_HR = nn.Conv2d(in_nc, nf//2, 9, scale, 4)

        body = [CRB_Layer(nf//2, nf//2) for _ in range(num_blocks)]
        self.body = nn.Sequential(*body)

        self.out = nn.Conv2d(nf//2, 10, 3, 1, 1)
        self.globalPooling = nn.AdaptiveAvgPool2d((1,1))
       

    def forward(self,GT, LR):

        lrf = self.head_LR(LR)
        hrf = self.head_HR(GT)

        f = [lrf, hrf]
        f, _ = self.body(f)
        f = self.out(f)
        f = self.globalPooling(f)
        f = f.view(f.size()[:2])

        return  f

class Restorer(nn.Module):
    def __init__(self, in_nc=3, out_nc=3, nf=64, nb=8, scale=4, input_para=10, min=0.0, max=1.0):
        super(Restorer, self).__init__()
        self.min = min
        self.max = max
        self.para = input_para
        self.num_blocks = nb

        self.head = nn.Conv2d(in_nc, nf, 3, stride=1, padding=1)

        body = [CRB_Layer(nf, input_para) for _ in range(nb)]
        self.body = nn.Sequential(*body)

        self.fusion = nn.Conv2d(nf, nf, 3, 1, 1)
        
        if scale == 4: #x4
            self.upscale = nn.Sequential(
                nn.Conv2d(in_channels=nf, out_channels=nf * scale, kernel_size=3, stride=1, padding=1, bias=True),
                nn.PixelShuffle(scale // 2),
                nn.Conv2d(in_channels=nf, out_channels=nf * scale, kernel_size=3, stride=1, padding=1, bias=True),
                nn.PixelShuffle(scale // 2),
                nn.Conv2d(nf, 3, 3, 1, 1)
            )
        else: #x2, x3
            self.upscale = nn.Sequential(
            nn.Conv2d(in_channels=nf, out_channels=nf*scale**2, kernel_size=3, stride=1, padding=1, bias=True),
            nn.PixelShuffle(scale),
            nn.Conv2d(nf, nf, 3, 1, 1)
            )

    def forward(self, input, ker_code):
        B, C, H, W = input.size() # I_LR batch
        B_h, C_h = ker_code.size() # Batch, Len=10
        ker_code_exp = ker_code.view((B_h, C_h, 1, 1)).expand((B_h, C_h, H, W)) #kernel_map stretch

        f = self.head(input)
        inputs = [f, ker_code_exp]
        f, _ = self.body(inputs)
        f = self.fusion(f)
        out = self.upscale(f)
        
        return out #torch.clamp(out, min=self.min, max=self.max)

class DAN(nn.Module):
    def __init__(self,nf=64, nb=16, scale=4, input_para=10, kernel_size=21, loop=8):
        super(DAN, self).__init__()

        self.ksize = kernel_size
        self.loop = loop
        self.scale = scale

        self.Restorer = Restorer(nf=nf, nb=nb, scale=scale, input_para=input_para)
        self.Predictor = Estimator(kernel_size=kernel_size, scale=self.scale)

        self.encoder = nn.Parameter(torch.load('/data/IKC/pca_matrix.pth')[None], requires_grad=False)

        kernel = torch.zeros(1, self.ksize, self.ksize)
        kernel[:,self.ksize//2, self.ksize//2] = 1
        self.init_kernel = nn.Parameter(kernel, requires_grad=False)
        self.init_ker_map = nn.Parameter(
            self.init_kernel.view(1, 1, self.ksize**2).matmul(self.encoder)[:,0], requires_grad=False
            )

    def forward(self, lr):

        srs = []
        ker_maps = []

        B, C, H, W = lr.shape
        ker_map = self.init_ker_map.repeat([B, 1])

        for i in range(self.loop):

            sr = self.Restorer(lr, ker_map.detach())
            ker_map = self.Estimator(sr.detach(), lr)

            srs.append(sr)
            ker_maps.append(ker_map)
        return [srs, ker_maps]


