import tensorflow as tf
from tensorflow.keras import layers
import numpy as np
import matplotlib.pyplot as plt
import os
import time

# Suppress TensorFlow warnings
import logging
logging.getLogger("tensorflow").setLevel(logging.ERROR)

# Define constants
DATASET_PATH = '/home/leo/Documents/dataset_images/train'
BATCH_SIZE = 128
IMG_HEIGHT = 64
IMG_WIDTH = 64
BUFFER_SIZE = 60000
EPOCHS = 100
NOISE_DIM = 100
NUM_EXAMPLES_TO_GENERATE = 16
AUTOTUNE = tf.data.AUTOTUNE

# Prepare the dataset
def load_and_preprocess_image(file_path):
    """
    Loads an image from the given file path and preprocesses it:
    - Reads the image file
    - Decodes the image
    - Converts it to RGB
    - Resizes to the target dimensions
    - Normalizes pixel values to [-1, 1]
    """
    image = tf.io.read_file(file_path)
    image = tf.image.decode_image(image, channels=3, expand_animations=False)
    image = tf.image.convert_image_dtype(image, tf.float32)  # Converts to [0,1]
    image = tf.image.resize(image, [IMG_HEIGHT, IMG_WIDTH])
    image = (image * 2.0) - 1.0  # Normalize to [-1, 1]
    
    # Debugging: Print shape
    tf.print("Image shape:", tf.shape(image))
    
    return image

def load_dataset(dataset_path):
    """
    Loads and preprocesses the dataset from the specified directory.
    Assumes that the dataset directory contains subdirectories for each class,
    but for GANs, class labels are not utilized.
    """
    # Collect all image file paths
    all_image_paths = []
    for root, dirs, files in os.walk(dataset_path):
        for file in files:
            if file.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif')):
                all_image_paths.append(os.path.join(root, file))
    
    print(f"Total images found: {len(all_image_paths)}")
    
    if len(all_image_paths) == 0:
        raise ValueError(f"No images found in the directory {dataset_path}. Please check the path.")
    
    # Create TensorFlow dataset from file paths
    path_ds = tf.data.Dataset.from_tensor_slices(all_image_paths)
    
    # Shuffle and preprocess
    image_ds = path_ds.shuffle(buffer_size=BUFFER_SIZE, reshuffle_each_iteration=True)
    image_ds = image_ds.map(load_and_preprocess_image, num_parallel_calls=AUTOTUNE)
    image_ds = image_ds.batch(BATCH_SIZE, drop_remainder=True).prefetch(buffer_size=AUTOTUNE)
    
    # Inspect the shape of a single batch
    for image_batch in image_ds.take(1):
        tf.print("Batch shape:", tf.shape(image_batch))
    
    return image_ds

# Build the Generator
def build_generator():
    model = tf.keras.Sequential()
    model.add(layers.Dense(8*8*256, use_bias=False, input_shape=(NOISE_DIM,)))
    model.add(layers.BatchNormalization())
    model.add(layers.LeakyReLU())
    
    model.add(layers.Reshape((8, 8, 256)))  # 8x8x256
    
    model.add(layers.Conv2DTranspose(128, (5,5), strides=(1,1), padding='same', use_bias=False))
    model.add(layers.BatchNormalization())
    model.add(layers.LeakyReLU())
    
    model.add(layers.Conv2DTranspose(64, (5,5), strides=(2,2), padding='same', use_bias=False))
    model.add(layers.BatchNormalization())
    model.add(layers.LeakyReLU())
    
    model.add(layers.Conv2DTranspose(3, (5,5), strides=(2,2), padding='same', use_bias=False, activation='tanh'))
    # Output shape: (64,64,3)
    
    return model

# Build the Discriminator
def build_discriminator():
    model = tf.keras.Sequential()
    model.add(layers.Conv2D(64, (5,5), strides=(2,2), padding='same',
                                     input_shape=[IMG_HEIGHT, IMG_WIDTH, 3]))
    model.add(layers.LeakyReLU())
    model.add(layers.Dropout(0.3))
    
    model.add(layers.Conv2D(128, (5,5), strides=(2,2), padding='same'))
    model.add(layers.LeakyReLU())
    model.add(layers.Dropout(0.3))
    
    model.add(layers.Flatten())
    model.add(layers.Dense(1))
    
    return model

# Loss and optimizers
cross_entropy = tf.keras.losses.BinaryCrossentropy(from_logits=True)

def discriminator_loss(real_output, fake_output):
    real_loss = cross_entropy(tf.ones_like(real_output), real_output)  # Real images labeled as 1
    fake_loss = cross_entropy(tf.zeros_like(fake_output), fake_output)  # Fake images labeled as 0
    total_loss = real_loss + fake_loss
    return total_loss

def generator_loss(fake_output):
    # Try to fool the discriminator, labels are 1
    return cross_entropy(tf.ones_like(fake_output), fake_output)

# Generator and Discriminator optimizers
generator_optimizer = tf.keras.optimizers.Adam(1e-4)
discriminator_optimizer = tf.keras.optimizers.Adam(1e-4)

# Prepare the dataset
train_dataset = load_dataset(DATASET_PATH)

# Build models
generator = build_generator()

# Verify generator output shape
sample_noise = tf.random.normal([1, NOISE_DIM])
generated_sample = generator(sample_noise, training=False)
tf.print("Generated sample shape:", tf.shape(generated_sample))

discriminator = build_discriminator()

# Checkpointing to save the model
checkpoint_dir = './training_checkpoints_dcgan'
checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt")
checkpoint = tf.train.Checkpoint(generator_optimizer=generator_optimizer,
                                 discriminator_optimizer=discriminator_optimizer,
                                 generator=generator,
                                 discriminator=discriminator)

# Define training step
@tf.function
def train_step(images):
    noise = tf.random.normal([BATCH_SIZE, NOISE_DIM])
    
    with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
        generated_images = generator(noise, training=True)
    
        real_output = discriminator(images, training=True)
        fake_output = discriminator(generated_images, training=True)
    
        gen_loss = generator_loss(fake_output)
        disc_loss = discriminator_loss(real_output, fake_output)
    
    gradients_of_generator = gen_tape.gradient(gen_loss, generator.trainable_variables)
    gradients_of_discriminator = disc_tape.gradient(disc_loss, discriminator.trainable_variables)
    
    generator_optimizer.apply_gradients(zip(gradients_of_generator, generator.trainable_variables))
    discriminator_optimizer.apply_gradients(zip(gradients_of_discriminator, discriminator.trainable_variables))
    
    return gen_loss, disc_loss

# Generate and save images
seed = tf.random.normal([NUM_EXAMPLES_TO_GENERATE, NOISE_DIM])

def generate_and_save_images(model, epoch, test_input):
    """
    Generates and saves images using the generator model.
    """
    # Generate images
    predictions = model(test_input, training=False)
    
    # Rescale images from [-1,1] to [0,1]
    predictions = (predictions + 1.0) / 2.0
    
    # Create a figure to display the images
    fig = plt.figure(figsize=(4,4))
    
    for i in range(predictions.shape[0]):
        plt.subplot(4,4,i+1)
        plt.imshow(predictions[i])
        plt.axis('off')
    
    plt.tight_layout()
    plt.savefig(f'generated_images_epoch_{epoch:04d}.png')
    plt.close()

# Training loop
def train(dataset, epochs):
    for epoch in range(1, epochs+1):
        start = time.time()
        
        gen_loss_list = []
        disc_loss_list = []
        
        for image_batch in dataset:
            g_loss, d_loss = train_step(image_batch)
            gen_loss_list.append(g_loss)
            disc_loss_list.append(d_loss)
        
        # Produce images for the current epoch
        generate_and_save_images(generator, epoch, seed)
        
        # Save the model every 10 epochs
        if (epoch) % 10 == 0:
            checkpoint.save(file_prefix = checkpoint_prefix)
        
        # Calculate mean losses
        mean_gen_loss = tf.reduce_mean(gen_loss_list)
        mean_disc_loss = tf.reduce_mean(disc_loss_list)
        
        print (f'Epoch {epoch}, Time: {time.time()-start:.2f} sec | '
               f'Gen Loss: {mean_gen_loss:.4f} | Disc Loss: {mean_disc_loss:.4f}')
    
    # Generate after the final epoch
    generate_and_save_images(generator, epochs, seed)

# Start training
train(train_dataset, EPOCHS)

# To restore the latest checkpoint (if needed)
# checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))

# To generate new images after training, you can use the following code:
def generate_images_after_training(model, num_images=16):
    """
    Generates and displays new images using the trained generator model.
    """
    noise = tf.random.normal([num_images, NOISE_DIM])
    predictions = model(noise, training=False)
    predictions = (predictions + 1.0) / 2.0  # Rescale to [0,1]
    
    fig = plt.figure(figsize=(4,4))
    for i in range(predictions.shape[0]):
        plt.subplot(4,4,i+1)
        plt.imshow(predictions[i])
        plt.axis('off')
    plt.tight_layout()
    plt.show()

# Example usage:
# generate_images_after_training(generator, 16)

results in 

Traceback (most recent call last):
  File "/home/leo/Documents/vsc/just a script.py", line 231, in <module>
    train(train_dataset, EPOCHS)
  File "/home/leo/Documents/vsc/just a script.py", line 209, in train
    g_loss, d_loss = train_step(image_batch)
  File "/home/leo/.local/lib/python3.10/site-packages/tensorflow/python/util/traceback_utils.py", line 153, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/tmp/__autograph_generated_fileuc8zuc1b.py", line 14, in tf__train_step
    fake_output = ag__.converted_call(ag__.ld(discriminator), (ag__.ld(generated_images),), dict(training=True), fscope)
  File "/home/leo/.local/lib/python3.10/site-packages/keras/src/utils/traceback_utils.py", line 122, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/home/leo/.local/lib/python3.10/site-packages/keras/src/layers/input_spec.py", line 227, in assert_input_compatibility
    raise ValueError(
ValueError: in user code:

    File "/home/leo/Documents/vsc/just a script.py", line 162, in train_step  *
        fake_output = discriminator(generated_images, training=True)
    File "/home/leo/.local/lib/python3.10/site-packages/keras/src/utils/traceback_utils.py", line 122, in error_handler  **
        raise e.with_traceback(filtered_tb) from None
    File "/home/leo/.local/lib/python3.10/site-packages/keras/src/layers/input_spec.py", line 227, in assert_input_compatibility
        raise ValueError(

    ValueError: Exception encountered when calling Sequential.call().
    
    Input 0 of layer "dense_1" is incompatible with the layer: expected axis -1 of input shape to have value 32768, but received input with shape (128, 8192)
    
    Arguments received by Sequential.call():
      • inputs=tf.Tensor(shape=(128, 32, 32, 3), dtype=float32)
      • training=True
      • mask=None