import pandas as pd
import gradio as gr
df = pd.read_json("data.json")


def filter_data(x, language):
    language_filtered_df = df[df['language'] == language]
    lower_bound = x
    upper_bound = x + 0.1
    mask = (language_filtered_df['prev_nn_sim'] >= lower_bound) & (language_filtered_df['prev_nn_sim'] < upper_bound)
    filtered_df = language_filtered_df.loc[mask, ["prompt", "prev_nn_prompt", "prev_nn_sim"]]
    filtered_df = filtered_df.sort_values(by="prev_nn_sim", ascending=True)
    return filtered_df

custom_css = """
#my_table table {
    table-layout: fixed;
    width: 100%;               /* or a fixed width like 700px */
}
#my_table table th,
#my_table table td {
    /* Force wrapping within cells: */
    white-space: normal;       
    word-wrap: break-word;     
    overflow-wrap: break-word;
    /* Example fixed width for all columns (or use nth-child to target individually): */
    width: 200px;              
}
"""

with gr.Blocks(css=custom_css) as demo:
    gr.Markdown("## Prompt Freshness Nearest Neighbor")
    gr.Markdown("### Select a similarity threshold (x) and a language to see the prompts that are within 0.1 of x. ")
    gr.Markdown("The nearest neighbor prompt is the prompt that is most similar to the original prompt that appears at a previous time step.")

    dropdown = gr.Dropdown(
        choices=[round((i + 3) * 0.1, 1) for i in range(7)], 
        value=0.7,  
        label="Select a similarity threshold (x)"
    )
    initial_data = filter_data(0.7, "English")
    
    language_dropdown = gr.Dropdown(
        choices=df['language'].unique().tolist(), 
        value="English", 
        label="Select a language"
    )
    
    output = gr.DataFrame(
        value=initial_data,
        label="Filtered Data",
        headers=["Prompt", "Nearest Neighbor Prompt", "Nearest Neighbor Similarity"],
        elem_id="my_table"  
    )
    
    dropdown.change(fn=filter_data, inputs=[dropdown, language_dropdown], outputs=output)
    language_dropdown.change(fn=filter_data, inputs=[dropdown, language_dropdown], outputs=output)

demo.launch()