import tensorflow as tf
from tensorflow.keras import layers

# ===========================
# DCGAN Generator: noise (100,) → image (28, 28, 1)
# Uses Conv2DTranspose for learned upsampling
# ===========================
def make_generator(noise_dim=100):
    model = tf.keras.Sequential([
        # Step 1: Dense projection — noise → feature volume
        layers.Dense(7 * 7 * 256, use_bias=False, input_shape=(noise_dim,)),
        layers.BatchNormalization(),
        layers.LeakyReLU(0.2),
        layers.Reshape((7, 7, 256)),    # → (7, 7, 256)

        # Step 2: Conv2DTranspose — upsample 7→7 (strides=1)
        layers.Conv2DTranspose(128, (5, 5), strides=(1, 1),
                               padding='same', use_bias=False),
        layers.BatchNormalization(),
        layers.LeakyReLU(0.2),           # → (7, 7, 128)

        # Step 3: Conv2DTranspose — upsample 7→14
        layers.Conv2DTranspose(64, (5, 5), strides=(2, 2),
                               padding='same', use_bias=False),
        layers.BatchNormalization(),
        layers.LeakyReLU(0.2),           # → (14, 14, 64)

        # Step 4: Conv2DTranspose — upsample 14→28 (output!)
        layers.Conv2DTranspose(1, (5, 5), strides=(2, 2),
                               padding='same', use_bias=False,
                               activation='tanh'),  # → (28, 28, 1)
    ], name="generator")
    return model

G = make_generator()
G.summary()
# Total params: ~1.8M
# Input:  (batch, 100) ← random noise
# Output: (batch, 28, 28, 1) ← generated MNIST-sized image
# Activation tanh → output range [-1, 1] (normalize real data to match!)

# Test
noise = tf.random.normal([1, 100])
fake_image = G(noise, training=False)
print(fake_image.shape)  # (1, 28, 28, 1) ✓

def make_discriminator():
    model = tf.keras.Sequential([
        # Conv Block 1: 28×28 → 14×14
        layers.Conv2D(64, (5, 5), strides=(2, 2), padding='same',
                      input_shape=(28, 28, 1)),
        layers.LeakyReLU(0.2),
        layers.Dropout(0.3),

        # Conv Block 2: 14×14 → 7×7
        layers.Conv2D(128, (5, 5), strides=(2, 2), padding='same'),
        layers.LeakyReLU(0.2),
        layers.Dropout(0.3),

        # Flatten & classify
        layers.Flatten(),                # → (7*7*128) = 6272
        layers.Dense(1),                 # → logit (NO sigmoid!)
    ], name="discriminator")
    return model

D = make_discriminator()
D.summary()
# Total params: ~400k
# Input:  (batch, 28, 28, 1) ← real or fake image
# Output: (batch, 1) ← logit (use from_logits=True in loss!)

# IMPORTANT design choices:
# 1. LeakyReLU (not ReLU) — prevents dead neurons in D
# 2. NO BatchNorm in D (or use SpectralNorm) — BN can destabilize GAN
# 3. NO sigmoid at output — use from_logits=True (numerically stable)
# 4. Dropout in D — prevents D from getting too strong too fast

import tensorflow as tf
from tensorflow import keras
import numpy as np
import time

# ===========================
# 1. Setup
# ===========================
NOISE_DIM = 100
BATCH_SIZE = 256
EPOCHS = 50

G = make_generator()
D = make_discriminator()
g_opt = keras.optimizers.Adam(2e-4, beta_1=0.5)  # beta_1=0.5 for GANs!
d_opt = keras.optimizers.Adam(2e-4, beta_1=0.5)
bce = keras.losses.BinaryCrossentropy(from_logits=True)

# Load MNIST
(train_images, _), (_, _) = keras.datasets.mnist.load_data()
train_images = train_images.reshape(-1, 28, 28, 1).astype('float32')
train_images = (train_images - 127.5) / 127.5  # normalize to [-1, 1]!

train_ds = (tf.data.Dataset.from_tensor_slices(train_images)
    .shuffle(60000).batch(BATCH_SIZE).prefetch(tf.data.AUTOTUNE))

# ===========================
# 2. Training step
# ===========================
@tf.function
def train_step(real_images):
    noise = tf.random.normal([BATCH_SIZE, NOISE_DIM])

    with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
        # Generate fake images
        fake_images = G(noise, training=True)

        # Discriminator predictions
        real_output = D(real_images, training=True)
        fake_output = D(fake_images, training=True)

        # Discriminator loss: real→1, fake→0
        d_loss_real = bce(tf.ones_like(real_output), real_output)
        d_loss_fake = bce(tf.zeros_like(fake_output), fake_output)
        d_loss = d_loss_real + d_loss_fake

        # Generator loss: fool D → make D think fake is real
        g_loss = bce(tf.ones_like(fake_output), fake_output)

    # Update D
    d_grads = disc_tape.gradient(d_loss, D.trainable_variables)
    d_opt.apply_gradients(zip(d_grads, D.trainable_variables))

    # Update G
    g_grads = gen_tape.gradient(g_loss, G.trainable_variables)
    g_opt.apply_gradients(zip(g_grads, G.trainable_variables))

    return d_loss, g_loss

# ===========================
# 3. Training loop
# ===========================
seed_noise = tf.random.normal([16, NOISE_DIM])  # fixed noise for visualization

for epoch in range(EPOCHS):
    start = time.time()
    for batch in train_ds:
        d_loss, g_loss = train_step(batch)

    elapsed = time.time() - start
    print(f"Epoch {epoch+1:3d}/{EPOCHS} ({elapsed:.1f}s) | "
          f"D_loss: {d_loss:.4f} | G_loss: {g_loss:.4f}")

    # Generate sample images every 5 epochs
    if (epoch + 1) % 5 == 0:
        generated = G(seed_noise, training=False)
        # save_grid(generated, f"gen_epoch_{epoch+1}.png")
        print(f"  → Sample images saved")

# Save models
G.save("generator.keras")
D.save("discriminator.keras")
print("🎨 GAN training complete!")

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers

# ===========================
# 1. Sampling layer (reparameterization trick)
# ===========================
class Sampling(layers.Layer):
    """z = mu + sigma * epsilon, where epsilon ~ N(0,1)
    This trick allows gradients to flow through the sampling step!
    """
    def call(self, inputs):
        z_mean, z_log_var = inputs
        epsilon = tf.random.normal(shape=tf.shape(z_mean))
        return z_mean + tf.exp(0.5 * z_log_var) * epsilon

# ===========================
# 2. Encoder: image → z_mean, z_log_var, z
# ===========================
LATENT_DIM = 16

encoder_input = keras.Input(shape=(28, 28, 1))
x = layers.Conv2D(32, 3, strides=2, padding='same', activation='relu')(encoder_input)
x = layers.Conv2D(64, 3, strides=2, padding='same', activation='relu')(x)
x = layers.Flatten()(x)
x = layers.Dense(256, activation='relu')(x)
z_mean = layers.Dense(LATENT_DIM, name='z_mean')(x)
z_log_var = layers.Dense(LATENT_DIM, name='z_log_var')(x)
z = Sampling()([z_mean, z_log_var])

encoder = keras.Model(encoder_input, [z_mean, z_log_var, z], name='encoder')

# ===========================
# 3. Decoder: z → image
# ===========================
decoder_input = keras.Input(shape=(LATENT_DIM,))
x = layers.Dense(7 * 7 * 64, activation='relu')(decoder_input)
x = layers.Reshape((7, 7, 64))(x)
x = layers.Conv2DTranspose(64, 3, strides=2, padding='same', activation='relu')(x)
x = layers.Conv2DTranspose(32, 3, strides=2, padding='same', activation='relu')(x)
decoder_output = layers.Conv2DTranspose(1, 3, padding='same', activation='sigmoid')(x)

decoder = keras.Model(decoder_input, decoder_output, name='decoder')

# ===========================
# 4. VAE Model with custom train_step
# ===========================
class VAE(keras.Model):
    def __init__(self, encoder, decoder, **kwargs):
        super().__init__(**kwargs)
        self.encoder = encoder
        self.decoder = decoder

    def train_step(self, data):
        with tf.GradientTape() as tape:
            z_mean, z_log_var, z = self.encoder(data)
            reconstruction = self.decoder(z)

            # Reconstruction loss (how good is the reconstruction?)
            recon_loss = tf.reduce_mean(
                keras.losses.binary_crossentropy(data, reconstruction)
            ) * 28 * 28

            # KL divergence (how close is latent to N(0,1)?)
            kl_loss = -0.5 * tf.reduce_mean(
                1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var))

            total_loss = recon_loss + kl_loss

        grads = tape.gradient(total_loss, self.trainable_variables)
        self.optimizer.apply_gradients(zip(grads, self.trainable_variables))
        return {"loss": total_loss, "recon": recon_loss, "kl": kl_loss}

# Train
vae = VAE(encoder, decoder)
vae.compile(optimizer=keras.optimizers.Adam(1e-3))
vae.fit(train_images_01, epochs=30, batch_size=128)  # normalized [0,1]

# Generate new images from random latent vectors!
z_sample = tf.random.normal([16, LATENT_DIM])
generated = decoder(z_sample)
# → 16 new MNIST-like images! 🎨

import tensorflow as tf
from tensorflow.keras import layers

# ===========================
# Conditional GAN: noise + label → specific class image
# ===========================

# Generator: (noise, label) → image
def make_cgan_generator(noise_dim=100, num_classes=10):
    # Noise input
    noise_in = layers.Input(shape=(noise_dim,))
    # Label input (one-hot or embedding)
    label_in = layers.Input(shape=(num_classes,))

    # Concatenate noise + label
    combined = layers.Concatenate()([noise_in, label_in])  # (110,)

    x = layers.Dense(7*7*128, activation='relu')(combined)
    x = layers.Reshape((7, 7, 128))(x)
    x = layers.Conv2DTranspose(64, 5, strides=2, padding='same', activation='relu')(x)
    x = layers.BatchNormalization()(x)
    img = layers.Conv2DTranspose(1, 5, strides=2, padding='same', activation='tanh')(x)

    return tf.keras.Model([noise_in, label_in], img, name='cgan_gen')

# Generate specific digit:
# noise = tf.random.normal([1, 100])
# label_7 = tf.one_hot([7], 10)  # "generate a 7"
# fake_7 = generator([noise, label_7])
# → Generates an image that looks like the digit 7! 🎯

import tensorflow as tf

# ===========================
# WGAN: replace BCE with Wasserstein distance
# Critic (not "discriminator") outputs unbounded score
# ===========================

# Critic loss: maximize E[C(real)] - E[C(fake)]
def critic_loss(real_output, fake_output):
    return tf.reduce_mean(fake_output) - tf.reduce_mean(real_output)

# Generator loss: minimize -E[C(fake)]
def generator_loss(fake_output):
    return -tf.reduce_mean(fake_output)

# WGAN-GP: gradient penalty (replaces weight clipping)
def gradient_penalty(critic, real, fake, batch_size):
    alpha = tf.random.uniform([batch_size, 1, 1, 1])
    interpolated = alpha * real + (1 - alpha) * fake

    with tf.GradientTape() as tape:
        tape.watch(interpolated)
        pred = critic(interpolated, training=True)

    grads = tape.gradient(pred, interpolated)
    norm = tf.sqrt(tf.reduce_sum(tf.square(grads), axis=[1, 2, 3]) + 1e-8)
    penalty = tf.reduce_mean((norm - 1.0) ** 2)
    return penalty

# Training: update Critic 5× per Generator update
# for each batch:
#   for _ in range(5):  # train critic more!
#     c_loss = critic_loss + 10 * gradient_penalty
#   g_loss = generator_loss

# WGAN advantages over standard GAN:
# ✅ Meaningful loss curve (correlates with image quality!)
# ✅ No mode collapse
# ✅ More stable training
# ✅ No need to carefully balance G and D

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt

# ===========================
# 1. Linear interpolation between two points
# ===========================
z1 = tf.random.normal([1, 100])  # point A in latent space
z2 = tf.random.normal([1, 100])  # point B in latent space

# Generate images along the path from A to B
steps = 10
images = []
for alpha in np.linspace(0, 1, steps):
    z = z1 * (1 - alpha) + z2 * alpha  # linear interpolation (LERP)
    img = generator(z, training=False)
    images.append(img[0])

# Show: smooth transition from image A to image B!
fig, axes = plt.subplots(1, steps, figsize=(20, 2))
for i, ax in enumerate(axes):
    ax.imshow(images[i][:, :, 0], cmap='gray')
    ax.axis('off')
plt.suptitle('Latent Space Interpolation')
plt.show()

# ===========================
# 2. Spherical interpolation (SLERP — better for high-dim)
# ===========================
def slerp(z1, z2, alpha):
    """Spherical linear interpolation — better than LERP for latent space"""
    z1_norm = z1 / (tf.norm(z1) + 1e-8)
    z2_norm = z2 / (tf.norm(z2) + 1e-8)
    omega = tf.acos(tf.clip_by_value(
        tf.reduce_sum(z1_norm * z2_norm), -1.0, 1.0))
    so = tf.sin(omega)
    return tf.sin((1 - alpha) * omega) / so * z1 + tf.sin(alpha * omega) / so * z2

# ===========================
# 3. Latent space arithmetic (like Word2Vec for images!)
# ===========================
# With a face GAN:
# z_man_glasses = encoder(man_with_glasses)
# z_man = encoder(man_without_glasses)
# z_woman = encoder(woman_without_glasses)
# z_woman_glasses = z_woman + (z_man_glasses - z_man)
# result = decoder(z_woman_glasses)  → woman with glasses! ✨

# ===========================
# 4. 2D latent grid visualization (for VAE with latent_dim=2)
# ===========================
# grid_x = np.linspace(-3, 3, 20)
# grid_y = np.linspace(-3, 3, 20)
# for i, xi in enumerate(grid_x):
#     for j, yi in enumerate(grid_y):
#         z = np.array([[xi, yi]])
#         img = decoder.predict(z)  → shows entire manifold of digits!