import torch
import torch.nn as nn
import torch.optim as optim

def train_gan_epoch(generator, discriminator, dataloader,
                    opt_G, opt_D, latent_dim, device):
    """
    One epoch of GAN training.
    Standard practice: update D once, then G once per batch.
    """
    criterion = nn.BCEWithLogitsLoss()  # numerically stable binary cross-entropy

    for batch_idx, (real_imgs, _) in enumerate(dataloader):
        batch_size = real_imgs.size(0)
        real_imgs = real_imgs.to(device)

        # Labels: real=1, fake=0 (smooth labels help stability)
        real_labels = torch.ones(batch_size, 1, device=device)  * 0.9  # label smoothing
        fake_labels = torch.zeros(batch_size, 1, device=device) + 0.1

        # ─────────────────────────────────────
        # STEP 1: Train Discriminator
        # ─────────────────────────────────────
        discriminator.zero_grad()

        # Real images loss: D should output ~1 for real images
        d_real = discriminator(real_imgs)
        loss_real = criterion(d_real, real_labels)

        # Fake images loss: D should output ~0 for generated images
        z = torch.randn(batch_size, latent_dim, device=device)
        fake_imgs = generator(z).detach()  # detach: don't backprop through G yet
        d_fake = discriminator(fake_imgs)
        loss_fake = criterion(d_fake, fake_labels)

        loss_D = (loss_real + loss_fake) / 2
        loss_D.backward()
        opt_D.step()

        # ─────────────────────────────────────
        # STEP 2: Train Generator
        # ─────────────────────────────────────
        generator.zero_grad()

        z = torch.randn(batch_size, latent_dim, device=device)
        fake_imgs = generator(z)
        d_fake_for_g = discriminator(fake_imgs)

        # Generator wants D to output 1 for its fake images (fool D)
        # Use real_labels here: G is trying to make D think fakes are real
        loss_G = criterion(d_fake_for_g, torch.ones(batch_size, 1, device=device))
        loss_G.backward()
        opt_G.step()

        if batch_idx % 100 == 0:
            print(f"  Batch {batch_idx:04d}  |  D loss: {loss_D.item():.4f}  "
                  f"G loss: {loss_G.item():.4f}  "
                  f"D(real): {d_real.mean().item():.3f}  "
                  f"D(fake): {d_fake.mean().item():.3f}")

    return loss_D.item(), loss_G.item()

import torch
import torch.nn as nn

LATENT_DIM  = 100   # noise vector dimension
IMAGE_SIZE  = 64    # 64x64 output images
N_CHANNELS  = 3     # RGB
N_FEATURES  = 64    # base channel count (multiplied per stage)


class DCGANGenerator(nn.Module):
    """
    DCGAN Generator: noise vector (100,) → image (3, 64, 64)
    Uses transposed convolutions to progressively upsample:
    4x4 → 8x8 → 16x16 → 32x32 → 64x64
    """
    def __init__(self, latent_dim=LATENT_DIM, ngf=N_FEATURES):
        super().__init__()
        self.net = nn.Sequential(
            # Input: (latent_dim,) → (ngf*8, 4, 4)
            nn.ConvTranspose2d(latent_dim, ngf * 8, 4, 1, 0, bias=False),
            nn.BatchNorm2d(ngf * 8),
            nn.ReLU(True),
            # (ngf*8, 4, 4) → (ngf*4, 8, 8)
            nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 4),
            nn.ReLU(True),
            # (ngf*4, 8, 8) → (ngf*2, 16, 16)
            nn.ConvTranspose2d(ngf * 4, ngf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 2),
            nn.ReLU(True),
            # (ngf*2, 16, 16) → (ngf, 32, 32)
            nn.ConvTranspose2d(ngf * 2, ngf, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf),
            nn.ReLU(True),
            # (ngf, 32, 32) → (3, 64, 64)
            nn.ConvTranspose2d(ngf, N_CHANNELS, 4, 2, 1, bias=False),
            nn.Tanh(),  # output in [-1, 1] — normalize your real images the same way!
        )

    def forward(self, z):
        # z: (batch, latent_dim, 1, 1)
        if z.dim() == 2:
            z = z.unsqueeze(-1).unsqueeze(-1)
        return self.net(z)


class DCGANDiscriminator(nn.Module):
    """
    DCGAN Discriminator: image (3, 64, 64) → scalar probability of being real
    Uses strided convolutions to progressively downsample:
    64x64 → 32x32 → 16x16 → 8x8 → 4x4 → 1x1
    """
    def __init__(self, ndf=N_FEATURES):
        super().__init__()
        self.net = nn.Sequential(
            # Input: (3, 64, 64) → (ndf, 32, 32)
            # No BatchNorm on first layer (per DCGAN paper)
            nn.Conv2d(N_CHANNELS, ndf, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),
            # (ndf, 32, 32) → (ndf*2, 16, 16)
            nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 2),
            nn.LeakyReLU(0.2, inplace=True),
            # (ndf*2, 16, 16) → (ndf*4, 8, 8)
            nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 4),
            nn.LeakyReLU(0.2, inplace=True),
            # (ndf*4, 8, 8) → (ndf*8, 4, 4)
            nn.Conv2d(ndf * 4, ndf * 8, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 8),
            nn.LeakyReLU(0.2, inplace=True),
            # (ndf*8, 4, 4) → (1, 1, 1)
            # No Sigmoid here — use BCEWithLogitsLoss for numerical stability
            nn.Conv2d(ndf * 8, 1, 4, 1, 0, bias=False),
            nn.Flatten(),
        )

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


# Instantiate and test
G = DCGANGenerator(LATENT_DIM)
D = DCGANDiscriminator()

z = torch.randn(8, LATENT_DIM)
fake_images = G(z)
d_scores    = D(fake_images)

print(f"Noise shape:       {z.shape}")
print(f"Generated images:  {fake_images.shape}")
print(f"D output shape:    {d_scores.shape}")
print(f"G parameters:      {sum(p.numel() for p in G.parameters())/1e6:.2f}M")
print(f"D parameters:      {sum(p.numel() for p in D.parameters())/1e6:.2f}M")

import torch
import torch.nn as nn

def gradient_penalty(critic, real_imgs, fake_imgs, device):
    """WGAN-GP gradient penalty — enforces Lipschitz constraint on the critic."""
    batch_size = real_imgs.size(0)

    # Interpolate between real and fake
    alpha = torch.rand(batch_size, 1, 1, 1, device=device)
    interpolated = (alpha * real_imgs + (1 - alpha) * fake_imgs).requires_grad_(True)

    # Critic score on interpolated samples
    interp_score = critic(interpolated)

    # Gradient of critic output with respect to interpolated input
    grad = torch.autograd.grad(
        outputs=interp_score,
        inputs=interpolated,
        grad_outputs=torch.ones_like(interp_score),
        create_graph=True,
        retain_graph=True,
    )[0]

    # Gradient norm — should be 1 for Lipschitz constraint
    grad_norm = grad.view(batch_size, -1).norm(2, dim=1)
    penalty   = ((grad_norm - 1) ** 2).mean()
    return penalty


def train_wgan_gp_step(generator, critic, real_imgs, opt_G, opt_C,
                        latent_dim, lambda_gp=10, n_critic=5, device='cpu'):
    """
    WGAN-GP update step.
    n_critic: number of critic updates per generator update (typically 5)
    lambda_gp: weight of gradient penalty (typically 10)
    """
    batch_size = real_imgs.size(0)
    real_imgs = real_imgs.to(device)

    # ── Train Critic (n_critic times) ──
    for _ in range(n_critic):
        z = torch.randn(batch_size, latent_dim, device=device)
        fake_imgs = generator(z).detach()

        critic.zero_grad()

        # WGAN loss: maximize E[C(real)] - E[C(fake)]
        # (minimize its negative)
        c_real = critic(real_imgs).mean()
        c_fake = critic(fake_imgs).mean()
        gp     = gradient_penalty(critic, real_imgs, fake_imgs, device)

        loss_C = -(c_real - c_fake) + lambda_gp * gp
        loss_C.backward()
        opt_C.step()

    # ── Train Generator (once) ──
    generator.zero_grad()
    z         = torch.randn(batch_size, latent_dim, device=device)
    fake_imgs = generator(z)
    loss_G    = -critic(fake_imgs).mean()   # maximize C(fake) → fool critic
    loss_G.backward()
    opt_G.step()

    return loss_C.item(), loss_G.item()

print("WGAN-GP: more stable training than standard GAN")
print("Key change: critic outputs unconstrained real numbers, not probabilities")
print("Gradient penalty: ensures the critic function is 1-Lipschitz")

# Trick 1: Label Smoothing
# Instead of real=1.0, use real=0.9 (one-sided label smoothing)
# Prevents D from becoming over-confident, maintains gradient flow to G
real_labels = torch.full((batch_size, 1), 0.9, device=device)  # not 1.0

# Trick 2: Add Noise to D Inputs
# Add small Gaussian noise to both real and fake images fed to D
# Blurs the decision boundary slightly, preventing D from memorizing
noise_std = 0.05
noisy_real = real_imgs + noise_std * torch.randn_like(real_imgs)
noisy_fake = fake_imgs + noise_std * torch.randn_like(fake_imgs)

# Trick 3: Spectral Normalization on D
# Normalize D's weight matrices by their spectral norm → bounds the Lipschitz constant
# from torch.nn.utils import spectral_norm
# self.conv1 = spectral_norm(nn.Conv2d(in_ch, out_ch, 3, padding=1))

# Trick 4: Different Learning Rates
# G typically: lr=2e-4, betas=(0.5, 0.999)
# D typically: lr=2e-4, betas=(0.5, 0.999)  or slightly lower than G
opt_G = torch.optim.Adam(None, lr=2e-4, betas=(0.5, 0.999))  # β1=0.5, not 0.9
opt_D = torch.optim.Adam(None, lr=2e-4, betas=(0.5, 0.999))

# Why beta1=0.5? Standard 0.9 makes Adam accumulate gradient history too aggressively
# for GAN training — 0.5 is more responsive to recent gradient changes

import torch
import torch.nn as nn


class ConditionalGenerator(nn.Module):
    """
    cGAN Generator: (noise, label) → image
    Conditions on class label via embedding concatenation.
    """
    def __init__(self, latent_dim=100, num_classes=10, img_size=28, channels=1):
        super().__init__()
        self.img_size  = img_size
        self.channels  = channels

        # Learnable label embedding (one vector per class)
        self.label_emb = nn.Embedding(num_classes, num_classes)

        self.net = nn.Sequential(
            nn.Linear(latent_dim + num_classes, 256),
            nn.LeakyReLU(0.2, inplace=True),
            nn.BatchNorm1d(256),
            nn.Linear(256, 512),
            nn.LeakyReLU(0.2, inplace=True),
            nn.BatchNorm1d(512),
            nn.Linear(512, 1024),
            nn.LeakyReLU(0.2, inplace=True),
            nn.BatchNorm1d(1024),
            nn.Linear(1024, channels * img_size * img_size),
            nn.Tanh(),
        )

    def forward(self, z, labels):
        # Embed the label and concatenate with noise
        label_emb = self.label_emb(labels)   # (batch, num_classes)
        gen_input = torch.cat([z, label_emb], dim=1)
        img = self.net(gen_input)
        return img.view(-1, self.channels, self.img_size, self.img_size)


class ConditionalDiscriminator(nn.Module):
    """
    cGAN Discriminator: (image, label) → real/fake probability
    """
    def __init__(self, num_classes=10, img_size=28, channels=1):
        super().__init__()
        self.img_size  = img_size

        self.label_emb = nn.Embedding(num_classes, num_classes)

        self.net = nn.Sequential(
            nn.Linear(channels * img_size * img_size + num_classes, 512),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(512, 256),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(256, 128),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(128, 1),
        )

    def forward(self, img, labels):
        flat_img  = img.view(img.size(0), -1)
        label_emb = self.label_emb(labels)
        d_input   = torch.cat([flat_img, label_emb], dim=1)
        return self.net(d_input)


# Test conditional generation
cG = ConditionalGenerator(latent_dim=100, num_classes=10)
cD = ConditionalDiscriminator(num_classes=10)

z      = torch.randn(16, 100)
labels = torch.randint(0, 10, (16,))   # which digit to generate

fake_imgs = cG(z, labels)
d_scores  = cD(fake_imgs, labels)

print(f"Generated images shape:  {fake_imgs.shape}")    # (16, 1, 28, 28)
print(f"Discriminator scores:    {d_scores.shape}")     # (16, 1)

# Generate a specific digit on demand:
z_digit7 = torch.randn(1, 100)
label_7  = torch.tensor([7])
generated_7 = cG(z_digit7, label_7)
print(f"Generated digit 7 image shape: {generated_7.shape}")  # (1, 1, 28, 28)

# StyleGAN-style AdaIN (Adaptive Instance Normalization) block
# The key mechanism: style vector controls per-channel mean and variance
import torch
import torch.nn as nn

class AdaIN(nn.Module):
    """
    Adaptive Instance Normalization: applies style to feature maps.
    Normalises feature map, then affine-transforms with style-derived params.
    """
    def __init__(self, n_channels, style_dim):
        super().__init__()
        self.norm  = nn.InstanceNorm2d(n_channels)
        # Map style vector to per-channel scale (gamma) and shift (beta)
        self.style_scale = nn.Linear(style_dim, n_channels)
        self.style_bias  = nn.Linear(style_dim, n_channels)

    def forward(self, x, style):
        # x: (B, C, H, W); style: (B, style_dim)
        x_norm = self.norm(x)

        gamma = self.style_scale(style).view(-1, x.size(1), 1, 1)
        beta  = self.style_bias(style).view(-1, x.size(1), 1, 1)

        return gamma * x_norm + beta   # modulate the feature map with style


# Demo
adain = AdaIN(n_channels=64, style_dim=512)
features = torch.randn(2, 64, 32, 32)
style    = torch.randn(2, 512)
out      = adain(features, style)
print(f"AdaIN input:  {features.shape}")  # (2, 64, 32, 32)
print(f"AdaIN output: {out.shape}")       # (2, 64, 32, 32) — same shape, different style

# FID computation using torch-fidelity (pip install torch-fidelity)
from torchvision.utils import save_image
import torch_fidelity

# Step 1: Save real images to a directory
# Step 2: Generate and save fake images to another directory
# Step 3: Compute FID

# Example with generated images saved to 'fake_images/' and real to 'real_images/'
# metrics = torch_fidelity.calculate_metrics(
#     input1='real_images/',
#     input2='fake_images/',
#     cuda=True,
#     isc=True,    # also compute Inception Score
#     fid=True,    # compute FID
#     verbose=True
# )
# print(f"FID: {metrics['frechet_inception_distance']:.2f}")
# print(f"IS:  {metrics['inception_score_mean']:.2f}")

# Reference FID values for comparison:
reference_fids = {
    "Random noise":          10000,
    "Vanilla GAN (MNIST)":   25,
    "DCGAN (CelebA 64x64)":  12,
    "Progressive GAN (1024)": 8.0,
    "StyleGAN2 (FFHQ 1024)":  2.8,
    "StyleGAN3 (FFHQ)":       2.1,
    "Real images (lower bound)": 0.0,
}

print(f"{'Model':<35} {'FID':>8}")
print("-" * 45)
for name, fid in reference_fids.items():
    bar = '█' * max(1, int(50 * min(fid, 100) / 100))
    print(f"{name:<35} {fid:>8.1f}  {bar}")