import os
import gradio as gr
from datasets import load_dataset

val_data_files = {"val": "contextual_val.csv"}
contextual_val = load_dataset("ucla-contextual/contextual_val", data_files=val_data_files)
print(contextual_val)

data = contextual_val['val']
print(data[0])

df = data.to_pandas()

df['image'] = df['image_url'].apply(lambda x: '<a href= "' + str(x) + '" target="_blank"> <img src= "' + str(
    x) + '" width="400"/> </a>')

cols = list(df.columns)
cols.insert(0, cols.pop(cols.index('image')))
cols.insert(4, cols.pop(cols.index('image_url')))

df = df.reindex(columns=cols)

LINES_NUMBER = 20

def display_df():
    df_images = df.head(LINES_NUMBER)
    return df_images

def display_next(dataframe, end):
    start = int(end or len(dataframe))
    end = int(start) + int(LINES_NUMBER)
    global df
    if end >= len(df) - 1:
        start = 0
        end = LINES_NUMBER
        df = df.sample(frac=1)
        print(f"Shuffle")
    df_images = df.iloc[start:end]
    assert len(df_images) == LINES_NUMBER
    return df_images, end

initial_dataframe = display_df()

# Gradio Blocks
with gr.Blocks() as demo:
    gr.Markdown("<h1><center>Contextual Val Dataset Viewer</center></h1>")

    with gr.Row():
        num_end = gr.Number(visible=False)
        b1 = gr.Button("Get Initial dataframe")
        b2 = gr.Button("Next Rows")

    with gr.Row():
        out_dataframe = gr.Dataframe(initial_dataframe, wrap=True, interactive=False, datatype = ['markdown', 'str', 'str', 'str', 'str', 'str'])

    b1.click(fn=display_df, outputs=out_dataframe, api_name="initial_dataframe")
    b2.click(fn=display_next, inputs=[out_dataframe, num_end], outputs=[out_dataframe, num_end],
             api_name="next_rows")

demo.launch(debug=True, show_error=True)