import gradio as gr
import pandas as pd

df = pd.read_csv("data.csv")

def filter_records(records, category):
    filtered = records[records["Category"] == category]
    
    if category == "NSFW":
        filtered = filtered.drop("Selective Alignment", axis=1)
    return filtered.iloc[:, 1:]

theme = gr.themes.Base()

with gr.Blocks(theme=theme) as demo:
    with gr.Row():
        button_celebrity = gr.Button(value="Celebrity", variant="primary")
        button_style = gr.Button(value="Style", variant="primary")
        button_ip = gr.Button(value="IP", variant="primary")
        button_nsfw = gr.Button(value="NSFW", variant="primary")
        button_all = gr.Button(value="Overall", variant="primary")
    
    # Initialize dataframe with Celebrity data
    celebrity_data = filter_records(df, "Celebrity")
    dataframe = gr.DataFrame(celebrity_data, column_widths="8%", elem_classes="light-mode")
    
    # Create separate functions for each category that always start with the original data
    def filter_celebrity():
        return filter_records(df, "Celebrity")
    
    def filter_style():
        return filter_records(df, "Style")
    
    def filter_ip():
        return filter_records(df, "IP")
    
    def filter_nsfw():
        return filter_records(df, "NSFW")
    
    def filter_overall():
        return filter_records(df, "Overall")
    
    # Connect each button to its specific filter function
    button_celebrity.click(fn=filter_celebrity, inputs=None, outputs=dataframe)
    button_style.click(fn=filter_style, inputs=None, outputs=dataframe)
    button_ip.click(fn=filter_ip, inputs=None, outputs=dataframe)
    button_nsfw.click(fn=filter_nsfw, inputs=None, outputs=dataframe)
    button_all.click(fn=filter_overall, inputs=None, outputs=dataframe)


demo.launch(share=True)