import torch
from torch import nn
import torch.nn as nn
import torchvision

class DenseLayer(nn.Module):
    def __init__(self, in_channels, out_channels,expansion = 4):
        super(DenseLayer, self).__init__()
        #self.expansion = expansion
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=3 // 2)
        self.conv1 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=3 // 2)
        self.relu = nn.ReLU(inplace=True)
        self.eca = eca_layer(out_channels,k_size=3)
        self.cbam = CBAM(channel=out_channels )
    def forward(self, x):

        f = self.conv(x)
        c = self.relu(f)
        y = self.conv1(c)
        h = self.relu(y)
        #t = self.conv1(h)
        #r = self.relu(t)

        c = self.eca(h)
        #c = self.cbam(h)
        return torch.cat([x, c], 1)


class RDB(nn.Module):
    def __init__(self, in_channels, growth_rate, num_layers):
        super(RDB, self).__init__()
        self.layers = nn.Sequential(*[DenseLayer(in_channels + growth_rate * i, growth_rate) for i in range(num_layers)])

        # local feature fusion
        self.lff = nn.Conv2d(in_channels + growth_rate * num_layers, growth_rate, kernel_size=1)

    def forward(self, x):
        return x + self.lff(self.layers(x))  # local residual learning


class RDN(nn.Module):
    def __init__(self, scale_factor, num_channels, num_features, growth_rate, num_blocks, num_layers,in_channels=64,out_channels=64):
        super(RDN, self).__init__()
        self.G0 = num_features
        self.G = growth_rate
        self.D = num_blocks
        self.C = num_layers

        # shallow feature extraction
        self.sfe1 = nn.Conv2d(num_channels, num_features, kernel_size=3, padding=3 // 2)
        self.sfe2 = nn.Conv2d(num_features, num_features, kernel_size=3, padding=3 // 2)

        # residual dense blocks
        self.rdbs = nn.ModuleList([RDB(self.G0, self.G, self.C)])
        for _ in range(self.D - 1):
            self.rdbs.append(RDB(self.G, self.G, self.C))

        # global feature fusion
        self.gff = nn.Sequential(
            nn.Conv2d(self.G * self.D, self.G0, kernel_size=1),
            nn.Conv2d(self.G0, self.G0, kernel_size=3, padding=3 // 2)
        )
        # 先进行上采样
        self.up = nn.ConvTranspose2d(in_channels=64, out_channels=64, kernel_size=8, stride=4, padding=2)
        self.dfg = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=4, padding=1)
        # 上采样特征进行重建
        self.sr_flow = nn.Sequential(*[
            nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, padding=1),
            nn.PReLU(out_channels),
            nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, padding=1),
            nn.PReLU(out_channels),
            nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, padding=1),
            nn.PReLU(out_channels)
        ])
        # 再进行下采样
        self.down = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=8, stride=4, padding=2)
        # 残差进行重建
        self.bp_flow = nn.Sequential(*[
            nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, padding=1),
            nn.PReLU(in_channels),
            nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, padding=1),
            nn.PReLU(out_channels),
            nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, padding=1),
            nn.PReLU(out_channels)
        ])
        self.eca = eca_layer(out_channels, k_size=3)

        # up-sampling
        assert 2 <= scale_factor <= 4
        if scale_factor == 2 or scale_factor == 4:
            self.upscale = []
            for _ in range(scale_factor // 2):
                self.upscale.extend([nn.Conv2d(self.G0, self.G0 * (2 ** 2), kernel_size=3, padding=3 // 2),
                                     nn.PixelShuffle(2)])
            self.upscale = nn.Sequential(*self.upscale)
        else:
            self.upscale = nn.Sequential(
                nn.Conv2d(self.G0, self.G0 * (scale_factor ** 2), kernel_size=3, padding=3 // 2),
                nn.PixelShuffle(scale_factor)
            )

        self.output = nn.Conv2d(self.G0, num_channels, kernel_size=3, padding=3 // 2)

    def forward(self, x):
        sfe1 = self.sfe1(x)
        sfe2 = self.sfe2(sfe1)

        x = sfe2
        local_features = []
        for i in range(self.D):
            x = self.rdbs[i](x)
            local_features.append(x)
        #bb = self.gff(torch.cat(local_features, 1)) + sfe1
        x11 = self.gff(torch.cat(local_features, 1)) + sfe1  # global residual learning


        up = self.up(sfe1)  # torch.Size([16, 64, 128, 128])
        # ox = self.sr_flow(up)
        # jx = self.eca(ox)  # torch.Size([16, 64, 128, 128])
        # gt = self.dfg(jx)
        down = self.down(up)  # torch.Size([16, 64, 32, 32])
        sub = sfe1 - down
        ix = self.bp_flow(sub)
        ux = self.eca(ix)
        ux += sub
        up2 = self.up(ux)
        ox2 = self.sr_flow(up2)
        x2 = self.eca(ox2)  # torch.Size([16, 64, 128, 128])
        gt2 = self.dfg(x2)  ######
        ######################
        down2 = self.down(up2)
        sub2 = ux - down2
        ix2 = self.bp_flow(sub2)
        ux2 = self.eca(ix2)
        ux2 += sub2  ####
        x = ux2 + gt2 + x11
        xx2 = self.upscale(x)
        x = self.output(xx2)
        return x

import torch
from torch import nn
from torch.nn.parameter import Parameter


class eca_layer(nn.Module):
    """Constructs a ECA module.

    Args:
        channel: Number of channels of the input feature map
        k_size: Adaptive selection of kernel size
    """

    def  __init__(self, out_channels, k_size=3):
        super(eca_layer, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.conv = nn.Conv1d(1, 1, kernel_size=k_size, padding=(k_size - 1) // 2, bias=False)

        self.sigmoid = nn.Sigmoid()

    def forward(self, x):#torch.Size([16, 64, 128, 128])
        # x: input features with shape [b, c, h, w]
        b, c, h, w = x.size()

        # feature descriptor on the global spatial information
        y = self.avg_pool(x)

        # Two different branches of ECA module
        y = self.conv(y.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1)

        # Multi-scale information fusion
        y = self.sigmoid(y)

        return x * y.expand_as(x)
    ###################################################################################
    #cbam
class ChannelAttentionModule(nn.Module):
    def __init__(self, channel, ratio=16):
        super(ChannelAttentionModule, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)

        self.shared_MLP = nn.Sequential(
            nn.Conv2d(channel, channel // ratio, 1, bias=False),
            nn.ReLU(),
            nn.Conv2d(channel // ratio, channel, 1, bias=False)
        )
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        avgout = self.shared_MLP(self.avg_pool(x))
        maxout = self.shared_MLP(self.max_pool(x))
        return self.sigmoid(avgout + maxout)


class SpatialAttentionModule(nn.Module):
    def __init__(self):
        super(SpatialAttentionModule, self).__init__()
        self.conv2d = nn.Conv2d(in_channels=2, out_channels=1, kernel_size=7, stride=1, padding=3)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        avgout = torch.mean(x, dim=1, keepdim=True)
        maxout, _ = torch.max(x, dim=1, keepdim=True)
        out = torch.cat([avgout, maxout], dim=1)
        out = self.sigmoid(self.conv2d(out))
        return out


class CBAM(nn.Module):
    def __init__(self, channel):
        super(CBAM, self).__init__()
        self.channel_attention = ChannelAttentionModule(channel)
        self.spatial_attention = SpatialAttentionModule()

    def forward(self, x):
        out = self.channel_attention(x) * x
        out = self.spatial_attention(out) * out
        return out