import streamlit as st
import requests
from datasets import load_dataset

NAME='bigcode/bigcodebench'

# Function to load data
@st.cache_data
def load_data(file):
    return load_dataset('bigcode/bigcodebench', split='v0.1.0_hf', use_auth_token=True)

# Function to filter data based on search keyword
def filter_data(data, keyword):
    if not keyword:
        return data
    filtered_data = [item for item in data if keyword.lower() in item['complete_prompt'].lower()]
    return filtered_data

data = load_data(NAME)

search_keyword = st.sidebar.text_input('Search by keyword', '')
filtered_data = filter_data(data, search_keyword)
if filtered_data:
    index = st.sidebar.number_input('Select Index', min_value=0, max_value=len(filtered_data)-1, value=0, step=1)    
    show_solution = st.sidebar.checkbox('Show Solution')
    show_test = st.sidebar.checkbox('Show Test Cases')
    snippet1 = filtered_data[index]['complete_prompt']
    snippet2 = filtered_data[index]['instruct_prompt']
    task_id = filtered_data[index]['task_id']
    st.title(task_id)
    
    st.subheader("Code Completion")
    st.code(snippet1, language='python')
    
    st.subheader("Natural Language Instruction to Code")
    st.code(snippet2, language='python')
    
    if show_solution:
        snippet3 = filtered_data[index]['canonical_solution']
        st.subheader("Solution")
        st.code(snippet3, language='python')
    
    if show_test:
        snippet4 = filtered_data[index]['test']
        st.subheader("Test Cases")
        st.code(snippet4, language='python')
else:
    st.error("No data available. Check the URL or search criteria.")