import torch
import torch.nn as nn
import torch.nn.functional as F


class DenseBlock(nn.Module):
    def __init__(self, feat_dim=64, growth=32):
        super().__init__()
        self.conv1 = nn.Conv2d(feat_dim, growth, 3, padding=1)
        self.conv2 = nn.Conv2d(feat_dim + growth, growth, 3, padding=1)
        self.conv3 = nn.Conv2d(feat_dim + 2 * growth, growth, 3, padding=1)
        self.conv4 = nn.Conv2d(feat_dim + 3 * growth, growth, 3, padding=1)
        self.conv5 = nn.Conv2d(feat_dim + 4 * growth, feat_dim, 3, padding=1)
        self.act = nn.LeakyReLU(0.2, inplace=True)

    def forward(self, x):
        c1 = self.act(self.conv1(x))
        c2 = self.act(self.conv2(torch.cat([x, c1], 1)))
        c3 = self.act(self.conv3(torch.cat([x, c1, c2], 1)))
        c4 = self.act(self.conv4(torch.cat([x, c1, c2, c3], 1)))
        c5 = self.conv5(torch.cat([x, c1, c2, c3, c4], 1))
        return x + 0.2 * c5



class RRDB(nn.Module):
    def __init__(self, feat_dim=64, growth=32):
        super().__init__()
        self.db1 = DenseBlock(feat_dim, growth)
        self.db2 = DenseBlock(feat_dim, growth)
        self.db3 = DenseBlock(feat_dim, growth)

    def forward(self, x):
        out = self.db3(self.db2(self.db1(x)))
        return x + 0.2 * out



class RRDBEncoder(nn.Module):
    def __init__(self, num_channels=3, feat_dim=64, num_blocks=6, growth=32):
        super().__init__()
        self.head = nn.Conv2d(num_channels, feat_dim, 3, padding=1)
        self.body = nn.Sequential(*[RRDB(feat_dim, growth) for _ in range(num_blocks)])
        self.tail = nn.Conv2d(feat_dim, feat_dim, 3, padding=1)

    def forward(self, x):
        x = self.head(x)
        return x + self.tail(self.body(x))



class ImplicitDecoder(nn.Module):
    def __init__(self, in_dim, hidden_dims=(256, 256, 256, 256), out_dim=3):
        super().__init__()
        layers = []
        prev = in_dim
        for h in hidden_dims:
            layers += [nn.Linear(prev, h), nn.ReLU(inplace=True)]
            prev = h
        layers.append(nn.Linear(prev, out_dim))
        self.mlp = nn.Sequential(*layers)

    def forward(self, x):
        return self.mlp(x)



def make_coord(shape, device):
    h, w = shape
    ys = (torch.arange(h, device=device).float() + 0.5) / h * 2 - 1
    xs = (torch.arange(w, device=device).float() + 0.5) / w * 2 - 1
    grid_y, grid_x = torch.meshgrid(ys, xs, indexing="ij")
    return torch.stack([grid_y, grid_x], dim=-1)



class LIIF(nn.Module):
    def __init__(self, feat_dim=64, num_blocks=6):
        super().__init__()
        self.encoder = RRDBEncoder(feat_dim=feat_dim, num_blocks=num_blocks)
        self.decoder = ImplicitDecoder(in_dim=feat_dim * 9 + 4)

    def gen_feat(self, lr_img):
        feat = self.encoder(lr_img)
        feat = F.unfold(feat, kernel_size=3, padding=1)
        B, C9, _ = feat.shape
        H, W = lr_img.shape[-2], lr_img.shape[-1]
        return feat.view(B, C9, H, W)

    def query_rgb(self, feat, coord, cell):
        B, C, H, W = feat.shape
        feat_coord = make_coord((H, W), feat.device)
        feat_coord = feat_coord.permute(2, 0, 1).unsqueeze(0).expand(B, 2, H, W)
        rx, ry = 1.0 / H, 1.0 / W

        preds, areas = [], []
        for vx in (-1, 1):
            for vy in (-1, 1):
                coord_ = coord.clone()
                coord_[:, :, 0] += vx * rx + 1e-6
                coord_[:, :, 1] += vy * ry + 1e-6
                coord_.clamp_(-1 + 1e-6, 1 - 1e-6)
                grid = coord_.flip(-1).unsqueeze(2)

                q_feat = F.grid_sample(feat, grid, mode="nearest",
                                       align_corners=False)[:, :, :, 0].permute(0, 2, 1)
                q_coord = F.grid_sample(feat_coord, grid, mode="nearest",
                                        align_corners=False)[:, :, :, 0].permute(0, 2, 1)
                rel = coord - q_coord
                rel[:, :, 0] *= H
                rel[:, :, 1] *= W

                inp = torch.cat([q_feat, rel, cell], dim=-1)
                Bq, Q, D = inp.shape
                preds.append(self.decoder(inp.view(Bq * Q, D)).view(Bq, Q, -1))
                areas.append(torch.abs(rel[:, :, 0] * rel[:, :, 1]) + 1e-9)

        total = areas[0] + areas[1] + areas[2] + areas[3]
        w0, w1, w2, w3 = areas[3] / total, areas[2] / total, areas[1] / total, areas[0] / total
        out = (preds[0] * w0.unsqueeze(-1) + preds[1] * w1.unsqueeze(-1)
               + preds[2] * w2.unsqueeze(-1) + preds[3] * w3.unsqueeze(-1))
        return out

    def forward(self, lr_img, coord, cell):
        feat = self.gen_feat(lr_img)
        return self.query_rgb(feat, coord, cell)