import torch
import torch.nn as nn

class UNetGeneratorImproved(nn.Module):
    """
    UNet simplificado para tareas de colorización.
    Entrada: imagen en escala de grises (1 canal)
    Salida: imagen RGB colorizada (3 canales)
    """
    def __init__(self):
        super().__init__()

        # Codificador: reduce la resolución y captura características
        self.encoder = nn.Sequential(
            nn.Conv2d(1, 64, kernel_size=4, stride=2, padding=1),  # 1 → 64 canales
            nn.ReLU(inplace=True),

            nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1),  # 64 → 128
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),

            nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1),  # 128 → 256
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
        )

        # Decodificador: recupera resolución y predice colores
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1),  # 256 → 128
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),

            nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1),  # 128 → 64
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),

            nn.ConvTranspose2d(64, 3, kernel_size=4, stride=2, padding=1),  # 64 → 3 (RGB)
            nn.Tanh()  # Para que la salida esté en el rango [-1, 1] o lo normalices luego
        )

    def forward(self, x):
        """
        Propagación hacia adelante del modelo
        """
        x = self.encoder(x)  # Pasa por el codificador
        x = self.decoder(x)  # Reconstruye la imagen color
        return x