import streamlit as st
from src.caller import OpenAI_Caller
from src.committee import Committee
from datasets import load_dataset

st.image('./img/meeting.psd')

st.header("Committee")

domains = st.multiselect(
    "Let's nominate committee members!",
    ["Age", "Disability Status", "Gender Identity", "Nationality", "Physical Appearance", "Race Ethnicity", "Religion", "Socioeconomic Status", "Sexual Orientation"],
    ["Age", "Nationality", "Religion"])

with st.spinner('Initializing committee members...'):
    committee_dict = {
        "chair": {
            "model_caller": OpenAI_Caller('gpt-4-1106-preview')
        },
        "member": [
            {"bias_type": domain, "model_caller": OpenAI_Caller('gpt-4-1106-preview')} for domain in domains
        ]
    }
with st.spinner('Initializing the committee...'):
    committee = Committee(committee_dict=committee_dict)

with st.spinner('Loding the BBQ datasets...'):
    dataset = load_dataset("Elfsong/BBQ")

category = st.selectbox("Which category you are going to deliberate?", dataset.keys())
# instance_id = st.number_input("Ok, which one?", min_value=0, max_value=len(dataset[category]))
raw_instance = dataset[category][0]

instance = {
    "context": raw_instance['context'],
    "question": raw_instance['question'],
    "ans0": raw_instance['ans0'],
    "ans1": raw_instance['ans1'],
    "ans2": raw_instance['ans2'],
}
instance = st.data_editor(instance)

# Propose
st.header("Propose")
proposals = list()
for member in committee.members:
    with st.spinner(f'Member who concerns {member.bias_type} is proposing...'):
        proposal = member.propose(instance)
        st.subheader(f"Proposal from Member who concerns [{member.bias_type}]:")
        st.markdown(proposal)
        proposals += [proposal]

# Craft Motion
st.header("Motion")
with st.spinner(f'Chair is crafting a motion...'):
    motion = committee.chair.craft_motion(proposals, instance)
    st.subheader(f"Motion from the Chair:")
    st.markdown(motion)

# Vote
st.header("Vote")
for member in committee.members:
    with st.spinner(f'Member who concerns {member.bias_type} is voting...'):
        vote_result = member.vote(motion, instance)
        st.subheader(f"Member who concerns [{member.bias_type}] votes:")
        st.markdown(vote_result)

st.balloons