import numpy as np
import cv2
import gradio as gr
from PIL import Image

# Load the PCA model
PCA_MODEL_PATH = "pca_texture_model.npy"
pca = np.load(PCA_MODEL_PATH, allow_pickle=True).item()

# PCA attributes
mean_texture = pca.mean_
components = pca.components_
n_components = components.shape[0]
TEXTURE_SIZE = np.sqrt(mean_texture.shape[0] // 3).astype(int)

# Slider ranges based on PCA variance
slider_ranges = [3 * np.sqrt(var) for var in pca.explained_variance_]

# Function to reconstruct texture from PCA components
def reconstruct_texture(component_values):
    # Calculate the texture using PCA components
    new_texture = mean_texture + np.dot(component_values, components)
    new_texture = np.clip(new_texture, 0, 255).astype(np.uint8)  # Ensure valid pixel range
    new_texture = new_texture.reshape((TEXTURE_SIZE, TEXTURE_SIZE, 3))  # Reshape back to image
    new_texture = cv2.cvtColor(new_texture, cv2.COLOR_BGR2RGB)  # Convert to RGB for display
    return Image.fromarray(new_texture)

# Function to handle slider values and generate the texture
def generate_texture(*component_values):
    component_values = np.array(component_values)
    return reconstruct_texture(component_values)

# Function to generate random textures
def random_texture():
    component_values = np.random.normal(0, 1, n_components)
    return reconstruct_texture(component_values), list(component_values)

# Gradio interface
def create_app():
    sliders = [
        gr.Slider(
            minimum=-slider_ranges[i],
            maximum=slider_ranges[i],
            step=0.1,
            label=f"Component {i + 1}",
            value=0,
        )
        for i in range(n_components)
    ]

    # Define interface elements
    with gr.Blocks() as app:
        with gr.Row():
            with gr.Column():
                gr.Markdown("### PCA Texture Generator")
                sliders_container = gr.Group(sliders)
                generate_button = gr.Button("Generate")
                random_button = gr.Button("Random")

            with gr.Column():
                preview = gr.Image(shape=(TEXTURE_SIZE, TEXTURE_SIZE), label="Preview")

        # Link sliders and preview
        generate_button.click(
            generate_texture, inputs=sliders, outputs=preview
        )

        random_button.click(
            random_texture, inputs=None, outputs=[preview] + sliders
        )

    return app


if __name__ == "__main__":
    # Create and launch the app
    app = create_app()
    app.launch()