import gradio as gr
from datasets import load_dataset
from transformers import pipeline

# Load the WikiArt dataset in streaming mode
dataset = load_dataset("huggan/wikiart", streaming=True)

# Function to display artwork details
def display_artwork(index):
    for i, record in enumerate(dataset["train"]):
        if i == index:
            return record["image"], f"Welcome to the Gallery\n\nTitle: {record['title']}\nArtist: {record['artist']}\nStyle: {record['style']}\nGenre: {record['genre']}\n\nStep into the world of art and explore its details."

# Function to filter artworks based on metadata
def filter_artworks(artist=None, genre=None, style=None):
    results = []
    for record in dataset["train"]:
        if (artist is None or record["artist"] == artist) and \
           (genre is None or record["genre"] == genre) and \
           (style is None or record["style"] == style):
            results.append(record)
    return results

# Function to display filtered artworks
def display_filtered_artworks(artist, genre, style):
    filtered_results = filter_artworks(artist, genre, style)
    return [(r["image"], f"Title: {r['title']}\nArtist: {r['artist']}\nStyle: {r['style']}\nGenre: {r['genre']}") for r in filtered_results]

# Chatbot functionality for museum guide persona
chatbot = pipeline("text-generation", model="gpt-4")

def museum_guide_query(prompt):
    return chatbot(prompt, max_length=100, num_return_sequences=1)[0]["generated_text"]

# Gradio interfaces
artwork_interface = gr.Interface(
    fn=display_artwork,
    inputs=gr.Number(label="Artwork Index"),
    outputs=[gr.Image(label="Artwork"), gr.Text(label="Details")],
    title="Exhibit AI - Virtual Art Gallery"
)

filter_interface = gr.Interface(
    fn=display_filtered_artworks,
    inputs=[gr.Text(label="Artist"), gr.Text(label="Genre"), gr.Text(label="Style")],
    outputs=gr.Gallery(label="Filtered Artworks", caption="Explore artworks based on your preferences."),
    title="Filter Artworks"
)

chatbot_interface = gr.Interface(
    fn=museum_guide_query,
    inputs=gr.Textbox(label="Ask the Museum Guide"),
    outputs=gr.Text(label="Guide Response"),
    title="Museum Guide Chatbot"
)

# Launch Gradio Blocks to combine all interfaces
def launch_app():
    with gr.Blocks() as demo:
        gr.Markdown("# Exhibit AI - Virtual Art Gallery")
        gr.TabbedInterface(
            [artwork_interface, filter_interface, chatbot_interface],
            ["View Artwork", "Filter Artworks", "Ask the Museum Guide"]
        )
    demo.launch()

if __name__ == "__main__":
    launch_app()