from datasets import load_dataset
import pandas as pd
import gradio as gr

# Load the Pokémon Cards dataset from Hugging Face
dataset = load_dataset("TheFusion21/PokemonCards")
df = pd.DataFrame(dataset["train"])

display_columns = ["name", "hp", "image_url"] + [col for col in df.columns if col not in ["image_url", "name", "hp"]]
display_df = df[display_columns]

def style_df(df):
    hp_min = df['hp'].min()
    hp_max = df['hp'].max()

    def color_hp(val):
        norm_val = (val - hp_min) / (hp_max - hp_min)
        r = int(255 * (1 - norm_val))
        g = int(255 * norm_val)
        return f'background-color: rgba({r}, {g}, 0, 0.2)'

    styled = df.style.applymap(color_hp, subset=['hp'])
    return styled

# Function to filter data based on user input
def filter_cards(set_name=None):
    filtered_df = df.copy()
    if set_name and set_name != "All":
        filtered_df = filtered_df[filtered_df["set_name"] == set_name]
    return style_df(filtered_df[display_columns])

def update_display(evt: gr.SelectData):
    selected_data = df.iloc[evt.index[0]]
    stats_df = pd.DataFrame({
        'Stat': ['Name', 'HP', 'Set Name', 'Caption'],
        'Value': [selected_data['name'], selected_data['hp'], selected_data['set_name'], selected_data['caption']]
    })
    return selected_data["image_url"], stats_df

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("## Pokémon Cards Explorer")

    with gr.Row():
        set_filter = gr.Dropdown(
            choices=["All"] + df["set_name"].unique().tolist(),
            label="Filter by Set Name"
        )

    filtered_table = gr.DataFrame(
        style_df(display_df),
        show_fullscreen_button=True,
        show_search="search",
        datatype=["str", "number", "image"] + ["str"] * (len(display_columns) - 3)
    )

    with gr.Row():
        card_image = gr.Image(label="Card Image", height=400)
        stats_table = gr.DataFrame(
            headers=["Stat", "Value"],
            label="Pokemon Stats",
            interactive=False,
            wrap=True
        )

    filter_button = gr.Button("Apply Filters")
    filter_button.click(
        filter_cards,
        inputs=[set_filter],
        outputs=filtered_table
    )

    filtered_table.select(
        update_display,
        None,
        [card_image, stats_table]
    )

demo.launch()