import gradio as gr
import tensorflow as tf
import numpy as np
from PIL import Image

# Load VGG for feature extraction
vgg = tf.keras.applications.VGG19(include_top=False, weights="imagenet")
vgg.trainable = False

STYLE_LAYERS = [...]  # your style layers
CONTENT_LAYER = [...] # your content layer(s)


def gram_matrix(A):
    """Compute Gram matrix for style representation.
    Works for VGG layer outputs of shape (batch, H, W, C).
    """
    A = tf.convert_to_tensor(A)

    if len(A.shape) == 4:
        batch, H, W, C = A.shape
        A = tf.reshape(A, (batch, H * W, C))
        gram = tf.matmul(A, A, transpose_a=True) / tf.cast(H * W, tf.float32)
    elif len(A.shape) == 3:
        batch, N, C = A.shape
        gram = tf.matmul(A, A, transpose_a=True) / tf.cast(N, tf.float32)
    elif len(A.shape) == 2:
        N, C = A.shape
        A = tf.expand_dims(A, 0)  
        gram = tf.matmul(A, A, transpose_a=True) / tf.cast(N, tf.float32)
    else:
        raise ValueError(f"Unexpected tensor rank for gram_matrix: {A.shape}")
    return gram


def compute_content_cost(a_C, a_G):
    return tf.reduce_mean(tf.square(a_C - a_G))


def compute_style_cost(a_S, a_G):
    J_style = 0
    for s, g in zip(a_S, a_G):
        J_style += tf.reduce_mean(tf.square(gram_matrix(s) - gram_matrix(g)))
    return J_style / len(a_S)


def total_cost(J_content, J_style, alpha=10, beta=40):
    return alpha * J_content + beta * J_style


def preprocess(img):
    img = Image.fromarray(img).resize((256, 256))
    arr = np.expand_dims(np.array(img) / 255.0, axis=0).astype(np.float32)
    return tf.convert_to_tensor(arr)


def style_transfer(content, style, steps):
    content_tensor = preprocess(content)
    style_tensor = preprocess(style)

    a_C = vgg(content_tensor)
    a_S = vgg(style_tensor)

    generated_image = tf.Variable(content_tensor)
    opt = tf.keras.optimizers.Adam(learning_rate=0.01)

    for i in range(steps):
        with tf.GradientTape() as tape:
            a_G = vgg(generated_image)
            J_style = compute_style_cost(a_S, a_G)
            J_content = compute_content_cost(a_C, a_G)
            J = total_cost(J_content, J_style, alpha=10, beta=40)

        grad = tape.gradient(J, generated_image)
        opt.apply_gradients([(grad, generated_image)])
        generated_image.assign(tf.clip_by_value(generated_image, 0.0, 1.0))

    out_img = (generated_image[0].numpy() * 255).astype("uint8")
    return Image.fromarray(out_img)


demo = gr.Interface(
    fn=style_transfer,
    inputs=[
        gr.Image(type="numpy", label="Content Image"), 
        gr.Image(type="numpy", label="Style Image"),
        gr.Slider(50, 2000, value=1000, step=50, label="Number of Iterations")
    ],
    outputs=gr.Image(type="pil", label="Stylized Image"),
)

demo.launch()