<gh_stars>0
import torch
import torch.nn as nn 
import torch.nn.functional as F
from torch.nn import Sequential
from .kernels import convKernel 
from .dwtHaar import dwtHaar_2d, idwtHaar_2d
from .customActivations import semiH_threshold4

""" Class: __Saturation_unit__
    Description:  This class implements the LET saturation unit used in 
            the rLWFSN CNN 
    constructor:
        * nts: Number of thresholds 
    forward inputs:
        * D: Input detail bands generated by a (learned) framelet 
    forward outputs:
        * out: thresholded signal
"""
class __saturation_unit__(nn.Module):
    def __init__(self, nts):
        super(__saturation_unit__, self).__init__()

        amp = 1e-3
        
        # Global thresholds
        self.t = nn.Parameter( torch.zeros((nts,1,1)))
        self.t.data.uniform_(0, 1/nts)
        
        self.a = nn.Parameter( torch.ones((1,nts,1,1)))
        
        self.nts =nts
        
    def forward(self, D):
        ds = D.shape
        Dr = D.reshape([ds[0]*ds[1], 1, ds[2], ds[3]])
        
        # Thresholding
        rep =  Dr.repeat( 1, self.nts, 1, 1)
        t = semiH_threshold4(rep,  self.t)

        
        # Mixing
        b = self.__softmax__(self.a)
        out = F.conv2d(t, b, bias = None, stride=(1,1), padding=0)
        
        return(  D - out.reshape([ds[0], ds[1], ds[2], ds[3]]) )

    def __softmax__(self,x):
        xs = x.shape
        xp = x.reshape(-1)
        xe = torch.exp(x)
        sf = xe/torch.sum(xe)
        return( sf.reshape(xs) )
        
""" Class: lwfsn_2d 
    Description:  This class implements the LWFSN CNN 
    constructor:
        * in_channels: Number of channels of the input image
        * depth: Depth of the CNN
        * wf: Number of channels after first convolution
    forward inputs:
        * x: Input image to be processed
        * bypass_shrinkage: If True, the LET threshold is ignored
    forward outputs:
        * LL: Denoised image
"""
class rlwfsn_2d(nn.Module):
    def __init__( self, in_channels=1, depth=1, wf=4):
        super(rlwfsn_2d, self).__init__()
        self.depth = depth
        
        # Encoding path
        self.shrink = nn.ModuleList()
        self.ew1 = nn.ModuleList()
        self.ew2 = nn.ModuleList()
        self.dw1 = nn.ModuleList()
        self.dw2 = nn.ModuleList()
        
        nts = 5
        prev_channels = in_channels 
        for i in range(depth):
            out_channels = wf*2**i
            # Thresholds
            self.shrink+= [__saturation_unit__( nts ) ]
            # Weight for the encoding
            self.ew1+= [convKernel(out_channels//2, prev_channels, 2, 2) ]
            self.ew2+= [convKernel(out_channels, out_channels//2, 2, 2) ]
            # Weights for the decoding path
            self.dw1+= [convKernel(out_channels//2, prev_channels, 2, 2) ]
            self.dw2+= [convKernel(out_channels, out_channels//2, 2, 2) ]

            # Other relevant constants
            prev_channels = out_channels

        # Defining inverse and forward transforms
        self.dwt = dwtHaar_2d()
        self.idwt = idwtHaar_2d()
        self.__set_requires_grad__(self.dwt, False)
        self.__set_requires_grad__(self.idwt, False)
        

    def forward(self, x, bypass_shrinkage=False):
        xs = x.shape
        LH_list = []; HL_list = []; HH_list = []

        LL = x
        for i in range(self.depth):
            # Loading kernels
            # Convoluting signal with factorized wavelet frame        
            LL, LH, HL, HH = self.__forwardF__(LL, i)

            lls = LL.shape
            # Processing bands, we define also a bypass option for training the
            # Encoding and decoding path
            if not bypass_shrinkage:
                # Shrinklage stage
                cat = torch.cat([LH, HL, HH], axis=1)
                LH, HL, HH = torch.split(self.shrink[i](cat), LH.shape[1], dim=1)
            
            LH_list.append( LH )
            HL_list.append( HL )
            HH_list.append( HH )
            
        if not bypass_shrinkage:
            LL = LL*0
        
        # Inverse transforming
        for i in range(self.depth):
            indx = self.depth - i -1
            LH = LH_list[indx]
            HL = HL_list[indx]
            HH = HH_list[indx]
            
            # Inverse-wavelet transforming
            # Croping extra samples
            LL  = self.__inverseF__(LL, LH, HL, HH, indx)
            
        if not bypass_shrinkage:
            out = x - LL
        else:
            out = LL
        return(out)
    
    
    def __set_requires_grad__(self, net, requires_grad=False):
        for param in net.parameters():
            param.requires_grad = requires_grad
            
    def __forwardF__(self, x, i):
        # Loading the weights
        W1 = self.ew1[i]()
        W2 = self.ew2[i]()

        # Convolutions
        xp = F.pad(x,(1,1,1,1), mode="constant")
        xW1 = F.conv2d(xp, W1, bias = None, stride=(1,1), padding=0)
        xW2 = F.conv2d(xW1, W2, bias = None, stride=(1,1), padding=0)
        return(self.dwt(xW2))
    
    def __inverseF__(self, LL, LH, HL, HH, i):
        # Loading the weights
        WT1 = self.dw1[i]().transpose(1,0)
        WT2 = self.dw2[i]().transpose(1,0)
        
        #Inverse frame transform
        iw= self.idwt(LL, LH, HL, HH)
        xp = F.pad(iw,(1,1,1,1), mode="constant")
        x2 = F.conv2d(xp, WT2, bias = None, stride=(1,1), padding=0)  
        x1 = F.conv2d(x2, WT1, bias = None, stride=(1,1), padding=0)  
        return(x1 )
