# model.py

import torch
import torch.nn as nn

class ResBlock(nn.Module):
    def __init__(self, nc):
        super().__init__()
        self.block = nn.Sequential(
            nn.Conv2d(nc, nc, 3, 1, 1, bias=False),
            nn.BatchNorm2d(nc),
            nn.ReLU(inplace=True),
            nn.Conv2d(nc, nc, 3, 1, 1, bias=False),
            nn.BatchNorm2d(nc)
        )
    def forward(self, x):
        return x + self.block(x)

class ResidualDenseBlock(nn.Module):
    def __init__(self,nf=64,gc=32):
        super(ResidualDenseBlock,self).__init__()
        self.conv1=nn.Conv2d(nf,gc,3,1,1)
        self.conv2=nn.Conv2d(nf+gc,gc,3,1,1)
        self.conv3=nn.Conv2d(nf+2*gc,gc,3,1,1)
        self.conv4=nn.Conv2d(nf+3*gc,gc,3,1,1)
        self.conv5=nn.Conv2d(nf+4*gc,nf,3,1,1)
        self.lrelu=nn.LeakyReLU(0.2,True)
    def forward(self,x):
        x1=self.lrelu(self.conv1(x))
        x2=self.lrelu(self.conv2(torch.cat((x,x1),1)))
        x3=self.lrelu(self.conv3(torch.cat((x,x1,x2),1)))
        x4=self.lrelu(self.conv4(torch.cat((x,x1,x2,x3),1)))
        x5=self.conv5(torch.cat((x,x1,x2,x3,x4),1))
        return x5*0.2+x

class RRDB(nn.Module):
    def __init__(self,nf,gc=32):
        super(RRDB,self).__init__()
        self.RDB1=ResidualDenseBlock(nf,gc)
        self.RDB2=ResidualDenseBlock(nf,gc)
        self.RDB3=ResidualDenseBlock(nf,gc)
    def forward(self,x):
        out=self.RDB1(x)
        out=self.RDB2(out)
        out=self.RDB3(out)
        return out*0.2+x

class TransformerBlock(nn.Module):
    def __init__(self,dim,num_heads,mlp_ratio=4.,qkv_bias=False):
        super().__init__()
        self.norm1=nn.LayerNorm(dim)
        self.attn=nn.MultiheadAttention(dim,num_heads,bias=qkv_bias,batch_first=True)
        self.norm2=nn.LayerNorm(dim)
        self.mlp=nn.Sequential(
            nn.Linear(dim,int(dim*mlp_ratio)),
            nn.GELU(),
            nn.Linear(int(dim*mlp_ratio),dim)
        )
    def forward(self,x):
        x = x + self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0]
        x = x + self.mlp(self.norm2(x))
        return x

class LightInpaintingTransformer(nn.Module):
    def __init__(self,
                 img_size=128, in_nc=3,
                 base_c=16,
                 t_dim=256,
                 t_heads=4, t_layers=2):
        super().__init__()
        self.num_patches = (img_size // 4) * (img_size // 4)

        # Encoder
        self.conv_first   = nn.Conv2d(in_nc, base_c, 3, 1, 1)
        self.enc_res1     = ResBlock(base_c)
        self.down1        = nn.Conv2d(base_c, base_c*2, 4, 2, 1)
        self.enc_res2     = ResBlock(base_c*2)
        self.down2        = nn.Conv2d(base_c*2, base_c*4, 4, 2, 1)

        # Transformer
        self.to_transformer = nn.Conv2d(base_c*4, t_dim, 1)
        self.pos_embed      = nn.Parameter(torch.zeros(1, self.num_patches, t_dim))
        self.transformer_blocks = nn.ModuleList([
            TransformerBlock(t_dim, t_heads) for _ in range(t_layers)
        ])
        self.from_transformer = nn.Conv2d(t_dim, base_c*4, 1)

        # Decoder
        self.up1        = nn.ConvTranspose2d(base_c*4, base_c*2, 2, 2)
        self.dec_res1   = ResBlock(base_c*2)
        self.up2        = nn.ConvTranspose2d(base_c*2, base_c, 2, 2)
        self.dec_res2   = ResBlock(base_c)
        self.conv_last  = nn.Sequential(
            nn.Conv2d(base_c, base_c, 3, 1, 1),
            nn.ReLU(inplace=True),
            nn.Conv2d(base_c, in_nc, 3, 1, 1),
            nn.Sigmoid()
        )

    def forward(self, x):
        B, C, H, W = x.shape
        e1 = self.enc_res1(self.conv_first(x))
        e2 = self.enc_res2(self.down1(e1))
        t  = self.down2(e2)
        t  = self.to_transformer(t).flatten(2).transpose(1, 2) + self.pos_embed
        for blk in self.transformer_blocks:
            t = blk(t)
        t  = self.from_transformer(t.transpose(1, 2).view(B, -1, H//4, W//4))
        d1 = self.dec_res1(self.up1(t))
        d2 = self.dec_res2(self.up2(d1))
        out = self.conv_last(d2)
        return out