import tensorflow as tf
import numpy as np
import cv2
import os
from tqdm import tqdm

def load_dataset(dataset_path, image_size=(512, 512)):
    images = []
    for file in tqdm(os.listdir(dataset_path)):
        img_path = os.path.join(dataset_path, file)
        img = cv2.imread(img_path)
        img = cv2.resize(img, image_size)
        img = (img / 127.5) - 1.0  # Normalize
        images.append(img)
    return np.array(images)

def build_generator():
    inputs = tf.keras.layers.Input(shape=(512, 512, 3))
    x = tf.keras.layers.Conv2D(64, (7, 7), padding="same", activation="relu")(inputs)
    x = tf.keras.layers.Conv2D(128, (3, 3), strides=2, padding="same")(x)
    x = tf.keras.layers.LeakyReLU(alpha=0.2)(x)
    x = tf.keras.layers.Conv2DTranspose(64, (3, 3), strides=2, padding="same")(x)
    x = tf.keras.layers.LeakyReLU(alpha=0.2)(x)
    x = tf.keras.layers.Conv2D(3, (7, 7), activation="tanh", padding="same")(x)
    return tf.keras.models.Model(inputs, x)

def train_animegan(dataset_path, epochs=100, batch_size=8):
    dataset = load_dataset(dataset_path)
    generator = build_generator()
    generator.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.0002, beta_1=0.5), loss="mse")
    for epoch in range(epochs):
        for i in range(0, len(dataset), batch_size):
            batch_images = dataset[i:i+batch_size]
            noise = np.random.normal(0, 1, (batch_size, 512, 512, 3))
            generator.train_on_batch(noise, batch_images)
        print(f"Epoch {epoch+1}/{epochs} completed")
    generator.save("AnimeGANv2_Hayao.h5")

if __name__ == "__main__":
    train_animegan("path/to/dataset")