import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms


class Generator(nn.Module):
    def __init__(self, latent_dim=100):
        super().__init__()
        self.net = nn.Sequential(
            nn.ConvTranspose2d(latent_dim, 128, 7, 1, 0),  # 1->7
            nn.BatchNorm2d(128), nn.ReLU(True),
            nn.ConvTranspose2d(128, 64, 4, 2, 1),           # 7->14
            nn.BatchNorm2d(64), nn.ReLU(True),
            nn.ConvTranspose2d(64, 1, 4, 2, 1),             # 14->28
            nn.Tanh(),
        )

    def forward(self, z):
        z = z.view(z.size(0), -1, 1, 1)
        return self.net(z)


class Discriminator(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(1, 64, 4, 2, 1),                      # 28->14
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(64, 128, 4, 2, 1),                    # 14->7
            nn.BatchNorm2d(128), nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(128, 1, 7, 1, 0),                     # 7->1
        )

    def forward(self, x):
        return self.net(x).view(-1)


def train_dcgan(epochs=5, batch_size=128, latent_dim=100, lr=2e-4):
    tfm = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5,), (0.5,)),  # scale to [-1, 1]
    ])
    # Fashion-MNIST: open, public-domain-licensed dataset.
    train_set = datasets.FashionMNIST(
        root="./data", train=True, download=True, transform=tfm
    )
    loader = DataLoader(train_set, batch_size=batch_size, shuffle=True)

    G = Generator(latent_dim)
    D = Discriminator()
    opt_G = torch.optim.Adam(G.parameters(), lr=lr, betas=(0.5, 0.999))
    opt_D = torch.optim.Adam(D.parameters(), lr=lr, betas=(0.5, 0.999))
    bce = nn.BCEWithLogitsLoss()

    for epoch in range(epochs):
        for real_images, _ in loader:
            batch = real_images.size(0)
            real_labels = torch.ones(batch)
            fake_labels = torch.zeros(batch)

            # --- Train Discriminator ---
            z = torch.randn(batch, latent_dim)
            fake_images = G(z).detach()
            d_real = D(real_images)
            d_fake = D(fake_images)
            d_loss = bce(d_real, real_labels) + bce(d_fake, fake_labels)

            opt_D.zero_grad()
            d_loss.backward()
            opt_D.step()

            # --- Train Generator (non-saturating loss) ---
            z = torch.randn(batch, latent_dim)
            fake_images = G(z)
            d_fake_for_g = D(fake_images)
            # maximize log D(G(z)) == minimize BCE against real_labels
            g_loss = bce(d_fake_for_g, real_labels)

            opt_G.zero_grad()
            g_loss.backward()
            opt_G.step()

        print(f"epoch {epoch}  d_loss {d_loss.item():.3f}  "
              f"g_loss {g_loss.item():.3f}")

    return G, D
