import skia
import numpy as np
import gradio as gr

def load_image(path):
    """Load an image from the given path using Skia."""
    return skia.Image.open(path)

def save_image(image, path):
    """Save the given image to the specified path."""
    image.save(path)

def image_to_array(image):
    """Convert a Skia image to a numpy array."""
    return np.array(image)

def array_to_image(array):
    """Convert a numpy array to a Skia image."""
    return skia.Image.fromarray(np.uint8(array))

def texture_image_with_normal(A, A_normal, B):
    """Texture image B onto image A based on the normal map A_normal."""
    # Load images
    A_img = load_image(A)
    A_normal_img = load_image(A_normal)
    B_img = load_image(B)

    # Convert images to numpy arrays
    A_array = image_to_array(A_img)
    A_normal_array = image_to_array(A_normal_img)
    B_array = image_to_array(B_img)

    # Ensure all images have the same dimensions
    assert A_array.shape == A_normal_array.shape == B_array.shape, "All images must have the same dimensions."

    # Normalize the normal map (values should be between 0 and 1)
    A_normal_array = A_normal_array / 255.0

    # Create an empty array for the textured image
    textured_array = np.zeros_like(A_array, dtype=np.float32)

    # Texture the image B onto A based on the normal map
    for y in range(A_array.shape[0]):
        for x in range(A_array.shape[1]):
            # Get the normal vector at this pixel
            normal = A_normal_array[y, x]

            # Calculate the offset based on the normal vector
            offset_x = int(normal[0] * B_array.shape[1])
            offset_y = int(normal[1] * B_array.shape[0])

            # Calculate the new coordinates in B
            new_x = (x + offset_x) % B_array.shape[1]
            new_y = (y + offset_y) % B_array.shape[0]

            # Apply the texture from B to A
            textured_array[y, x] = B_array[new_y, new_x]

    # Convert the textured array back to an image
    textured_image = array_to_image(textured_array)

    return textured_image

def gradio_interface(A, A_normal, B):
    """Gradio interface function to texture images."""
    textured_image = texture_image_with_normal(A, A_normal, B)
    return np.array(textured_image)

# Define example inputs
examples = [
    ["examples/image_A_1.png", "examples/image_A_normal_1.png", "examples/image_B_1.png"],
    ["examples/image_A_2.png", "examples/image_A_normal_2.png", "examples/image_B_2.png"],
    ["examples/image_A_3.png", "examples/image_A_normal_3.png", "examples/image_B_3.png"]
]

# Create Gradio interface
iface = gr.Interface(
    fn=gradio_interface,
    inputs=[
        gr.Image(type="filepath", label="Image A"),
        gr.Image(type="filepath", label="Normal Map A_normal"),
        gr.Image(type="filepath", label="Image B")
    ],
    outputs=gr.Image(type="numpy", label="Textured Image"),
    title="Image Texturing with Normal Map",
    description="Upload Image A, its Normal Map A_normal, and Image B to texture B onto A based on A_normal.",
    examples=examples
)

# Launch the interface
iface.launch()