# Modelo Attention Res-UNet 2D para Audio Super Resolution

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

class AttentionGate(nn.Module):
    """Módulo de atención para ponderar skip connections."""
    def __init__(self, F_g, F_l, F_int):
        """
        Inicializa el módulo de atención.
        Args:
            F_g (int): Número de canales del gating signal.
            F_l (int): Número de canales de la skip connection.
            F_int (int): Número de canales intermedios.
        """
        super().__init__()

        # Uso de Attention Gates basado en el paper "https://arxiv.org/abs/1804.03999"
        self.W_g = nn.Conv2d(F_g, F_int, kernel_size=1)
        self.W_x = nn.Conv2d(F_l, F_int, kernel_size=1)
        self.psi = nn.Conv2d(F_int, 1, kernel_size=1)
        self.relu = nn.ReLU(inplace=True)
        self.sigmoid = nn.Sigmoid()

    def forward(self, g, x):
        # Interpolar g para que tenga el mismo tamaño que x
        if g.shape[-2:] != x.shape[-2:]:
            g = F.interpolate(g, size=x.shape[-2:], mode="bilinear", align_corners=False)
        # Calcular atención
        att = self.sigmoid(self.psi(self.relu(self.W_g(g) + self.W_x(x))))
        return x * att


class DilatedBlock(nn.Module):
    """Bloque con capas convolucionales dilatadas para capturar contexto de largo alcance."""
    def __init__(self, channels):
        """
        Inicializa el bloque dilatado.
        Args:
            channels (int): Número de canales de entrada y salida.
        """
        super().__init__()

        self.net = nn.Sequential(
            nn.Conv2d(channels, channels, kernel_size=(7,3), padding=(3,1), dilation=(1,1)),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(channels, channels, kernel_size=(7,3), padding=(6,2), dilation=(2,2)),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(channels, channels, kernel_size=(7,3), padding=(12,4), dilation=(4,4)),
            nn.LeakyReLU(0.2, inplace=True),
        )

    def forward(self, x):
        return self.net(x) + x


class ResBlock(nn.Module):
    """Bloque Residual para Attention Res-UNet."""
    def __init__(self, in_ch, out_ch):
        """
        Inicializa el bloque residual.
        Args:
            in_ch (int): Número de canales de entrada.
            out_ch (int): Número de canales de salida.
        """
        super().__init__()
        
        self.conv1 = nn.Conv2d(in_ch, out_ch, kernel_size=(7,3), padding=(3,1))
        self.norm1 = nn.GroupNorm(out_ch//4, out_ch) # GroupNorm para batchs pequeños
        self.relu1 = nn.LeakyReLU(0.2, inplace=True)
        
        self.conv2 = nn.Conv2d(out_ch, out_ch, kernel_size=(7,3), padding=(3,1))
        self.norm2 = nn.GroupNorm(out_ch//4, out_ch)
        self.relu2 = nn.LeakyReLU(0.2, inplace=True)
        
        # Skip connection
        if in_ch != out_ch:
            self.skip = nn.Sequential(
                nn.Conv2d(in_ch, out_ch, kernel_size=1, bias=False),
                nn.GroupNorm(out_ch//4, out_ch)
            )
        else:
            self.skip = nn.Identity()

    def forward(self, x):
        identity = self.skip(x)
        
        out = self.conv1(x)
        out = self.norm1(out)
        out = self.relu1(out)
        
        out = self.conv2(out)
        out = self.norm2(out)
        
        out += identity
        out = self.relu2(out)
        
        return out


class UNetAudio2D(nn.Module):
    """
    Modelo Attention Res-UNet 2D para super resolución de audio.
    Entrada y salida con shape (B, 2, F, T).
    """
    def __init__(self):
        """Inicializa la arquitectura UNet con encoder, bottleneck y decoder."""
        super().__init__()

        # Encoder
        # Entrada: (B, 2, F, T)
        self.enc1 = ResBlock(2, 32)
        self.down1 = nn.Conv2d(32, 32, kernel_size=(4,4), stride=(2,2), padding=(1,1))  # Strided conv

        self.enc2 = ResBlock(32, 64)
        self.down2 = nn.Conv2d(64, 64, kernel_size=(4,4), stride=(2,2), padding=(1,1))

        self.enc3 = ResBlock(64, 128)
        self.down3 = nn.Conv2d(128, 128, kernel_size=(4,4), stride=(2,2), padding=(1,1))

        self.enc4 = ResBlock(128, 256)
        self.down4 = nn.Conv2d(256, 256, kernel_size=(4,4), stride=(2,2), padding=(1,1))

        # Bottleneck
        self.bottleneck_conv = ResBlock(256, 512)
        self.bottleneck_dilated = DilatedBlock(512)

        # Decoder
        self.up4 = self.up_block(512,256)
        self.dec4 = ResBlock(512,256)

        self.up3 = self.up_block(256,128)
        self.dec3 = ResBlock(256,128)

        self.up2 = self.up_block(128,64)
        self.dec2 = ResBlock(128,64)

        self.up1 = self.up_block(64,32)
        self.dec1 = ResBlock(64,32)

        # Attention gates
        self.att4 = AttentionGate(256,256,128)
        self.att3 = AttentionGate(128,128,64)
        self.att2 = AttentionGate(64,64,32)
        self.att1 = AttentionGate(32,32,16)

        # Output
        self.final = nn.Conv2d(32,2,kernel_size=1)

    def up_block(self, in_ch, out_ch):
        """
        Crea un bloque de upsampling con ConvTranspose2d.
        Args:
            in_ch (int): Número de canales de entrada.
            out_ch (int): Número de canales de salida.
        Returns:
            nn.Sequential: Bloque de upsampling en frecuencia y tiempo.
        """
        return nn.Sequential(
            nn.ConvTranspose2d(in_ch, out_ch, kernel_size=(4,4), stride=(2,2), padding=(1,1)),
            nn.LeakyReLU(0.2, inplace=True)
        )

    def forward(self, x):
        # Encoder
        e1 = self.enc1(x)
        e2 = self.enc2(self.down1(e1))
        e3 = self.enc3(self.down2(e2))
        e4 = self.enc4(self.down3(e3))

        # Bottleneck
        b = self.bottleneck_conv(self.down4(e4))
        b = self.bottleneck_dilated(b)

        # Decoder con skip connections y attention gates
        up4 = self.up4(b)
        d4 = self.dec4(torch.cat([up4, self.att4(up4, e4)], dim=1))

        up3 = self.up3(d4)
        d3 = self.dec3(torch.cat([up3, self.att3(up3, e3)], dim=1))

        up2 = self.up2(d3)
        d2 = self.dec2(torch.cat([up2, self.att2(up2, e2)], dim=1))

        up1 = self.up1(d2)
        d1 = self.dec1(torch.cat([up1, self.att1(up1, e1)], dim=1))

        return self.final(d1) + x