import streamlit as st
from ispice import Spice


# Sample example captions
sample_examples = {
    "ID - Example 1": {
        "generated": "P1 sits on the couch. P1 puts her hands on the couch and sits up. P2 sits beside her.",
        "reference": "P1 straightens her posture. P1 sits up. P2 sits beside her on the couch."
    }
}


# Function to compute score
def preprocess_captions(generated_captions, reference_captions):
    hypotheses = {'image'+str(i): [generated_captions[i]] for i in range(len(generated_captions))}
    references = {'image'+str(i): [reference_captions[i]] for i in range(len(reference_captions))}
    return hypotheses, references

# Streamlit app
def main():
    st.title("iSPICE Metric Evaluation")
    
    # Dropdown for comparison option
    mode = st.selectbox("Mode:", ["ID", "Name"])
    
    spice_scorer = Spice(mode=mode)
    
    # Note on how to use the app
    id_label = r'''If you choose the Mode as ID, supported Person IDs are from P1, P2, P3 .. upto P11.'''
    name_label = r'''If you choose the Mode as Name, supported Person Names are - Ray, Sam, Casey, Riley, Morgan, Alex, Quinn, Cameron, Avery, Charlie, Jamie and Mike.'''
        
    st.caption(id_label)
    st.caption(name_label)
    
    # Sample examples dropdown
    st.sidebar.subheader("Choose Sample Example:")
    example_choice = st.sidebar.selectbox("Select Example", list(sample_examples.keys()))


    # Description
    st.write("You can either input single caption or multiple captions separated by new line.")
    generated_caption = sample_examples[example_choice]["generated"]
    reference_caption = sample_examples[example_choice]["reference"]
    generated_caption_input = st.text_area("Generated Caption:", value=generated_caption)
    reference_caption_input = st.text_area("Reference Caption:", value=reference_caption)

    st.caption("We have a default caption here you can try out. You can modify it as required.")
    
    
    # Compute score button
    if st.button("Compute Score"):
        generated_captions = generated_caption_input.split("\n")
        reference_captions = reference_caption_input.split("\n")
        
        print(generated_captions, len(generated_captions))
        print(reference_captions, len(reference_captions))
        
    
        hypotheses, references = preprocess_captions(generated_captions, reference_captions)
        
        if generated_caption.strip() == "" or reference_caption.strip() == "":
            st.error("Please provide both generated and reference captions.")
        else:
            average_spice_score, spice_scores, average_ispice_score, ispice_scores = spice_scorer.compute_score(references, hypotheses)
            st.subheader("Scores :")
            st.write("Average SPICE Score:", average_spice_score)
            st.write("Average iSPICE Score:", average_ispice_score)
            st.write("SPICE Scores:", spice_scores)
            st.write("iSPICE Scores:", ispice_scores)

if __name__ == "__main__":
    main()