import math
import torch
import torch.nn as nn


class DenseBlock(nn.Module):
    def __init__(self, nf, gc):
        super().__init__()
        self.c1 = nn.Conv2d(nf, gc, 3, 1, 1)
        self.c2 = nn.Conv2d(nf + gc, gc, 3, 1, 1)
        self.c3 = nn.Conv2d(nf + 2 * gc, gc, 3, 1, 1)
        self.c4 = nn.Conv2d(nf + 3 * gc, gc, 3, 1, 1)
        self.c5 = nn.Conv2d(nf + 4 * gc, nf, 3, 1, 1)
        self.act = nn.LeakyReLU(0.2, inplace=True)

    def forward(self, x):
        x1 = self.act(self.c1(x))
        x2 = self.act(self.c2(torch.cat([x, x1], 1)))
        x3 = self.act(self.c3(torch.cat([x, x1, x2], 1)))
        x4 = self.act(self.c4(torch.cat([x, x1, x2, x3], 1)))
        x5 = self.c5(torch.cat([x, x1, x2, x3, x4], 1))
        return x + 0.2 * x5


class RRDB(nn.Module):
    def __init__(self, nf, gc):
        super().__init__()
        self.d1 = DenseBlock(nf, gc)
        self.d2 = DenseBlock(nf, gc)
        self.d3 = DenseBlock(nf, gc)

    def forward(self, x):
        out = self.d1(x)
        out = self.d2(out)
        out = self.d3(out)
        return x + 0.2 * out


class DenoiseSRNet(nn.Module):
    """Noisy LR in -> clean HR out, 4x upsample via pixel-shuffle."""
    def __init__(self, nf=64, gc=32, n_blocks=8, scale=4):
        super().__init__()
        self.scale = scale
        self.head = nn.Conv2d(1, nf, 3, 1, 1)
        self.body = nn.Sequential(*[RRDB(nf, gc) for _ in range(n_blocks)])
        self.body_conv = nn.Conv2d(nf, nf, 3, 1, 1)

        up_layers = []
        n_up = int(math.log2(scale))
        for _ in range(n_up):
            up_layers += [
                nn.Conv2d(nf, nf * 4, 3, 1, 1),
                nn.PixelShuffle(2),
                nn.LeakyReLU(0.2, inplace=True)
            ]
        self.upsample = nn.Sequential(*up_layers)

        self.tail = nn.Sequential(
            nn.Conv2d(nf, nf, 3, 1, 1),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(nf, 1, 3, 1, 1)
        )

    def forward(self, x):
        feat = self.head(x)
        body_out = self.body_conv(self.body(feat))
        feat = feat + body_out
        feat = self.upsample(feat)
        out = self.tail(feat)
        return torch.sigmoid(out)