import pandas as pd
import streamlit as st
from util.evaluator import evaluator, write_evaluation_commentary
import os

# Predefined examples
examples = {
    'good': {
        'question': "What causes rainbows to appear in the sky?",
        'explanation': "Rainbows appear when sunlight is refracted, dispersed, and reflected inside water droplets in the atmosphere, resulting in a spectrum of light appearing in the sky."
    },
    'bad': {
        'question': "What causes rainbows to appear in the sky?",
        'explanation': "Rainbows happen because light in the sky gets mixed up and sometimes shows colors when it's raining or when there is water around."
    }
}

# Function to check password
def check_password():
    with st.sidebar:
        password_input = st.text_input("Enter Password:", type="password")
        submit_button = st.button("Submit")
        if submit_button:
            if password_input == os.getenv('PASSWORD'):
                st.session_state['password_verified'] = True
                st.experimental_rerun()
            else:
                st.error("Incorrect Password, please try again.")

# Function to evaluate batch data
def evaluate_batch(uploaded_file):
    df = pd.read_csv(uploaded_file)
    eval_instance = evaluator(model_name=st.session_state['model_name'])
    results = []

    for _, row in df.iterrows():
        scores = eval_instance(row['question'], row['explanation'])
        commentary = write_evaluation_commentary(scores)
        result = {
            'Question': row['question'],
            'Explanation': row['explanation'],
            **{c['Principle']: c['Score'] for c in commentary}
        }
        results.append(result)

    return pd.DataFrame(results)

# Main app logic
def main():
    st.title('Natural Language Explanation Demo')

    model_name = st.selectbox('Select a model:', ['gpt4-1106', 'gpt35-1106'])
    st.session_state['model_name'] = model_name  # Save model name to session state for use in batch processing

    input_type = st.radio("Choose input type:", ('Use predefined example', 'Enter your own', 'Upload CSV for batch evaluation'))

    if input_type == 'Use predefined example':
        example_type = st.radio("Select an example type:", ('good', 'bad'))
        question = examples[example_type]['question']
        explanation = examples[example_type]['explanation']
    elif input_type == 'Enter your own':
        question = st.text_input('Enter your question:', '')
        explanation = st.text_input('Enter your explanation:', '')
    else:
        uploaded_file = st.file_uploader("Upload a CSV file", type='csv')
        if uploaded_file and st.button('Evaluate Batch'):
            result_df = evaluate_batch(uploaded_file)
            st.write('### Evaluated Results')
            st.dataframe(result_df)
            csv = result_df.to_csv(index=False)
            st.download_button(
                label="Download evaluated results as CSV",
                data=csv,
                file_name='batch_evaluation_results.csv',
                mime='text/csv'
            )
            return

    if st.button('Evaluate Explanation'):
        if question and explanation:
            eval_instance = evaluator(model_name)
            scores = eval_instance(question, explanation)
            st.write('### Scores')
            details = write_evaluation_commentary(scores)
            df = pd.DataFrame(details)
            st.write(df)
            data = {
                'Question': question,
                'Explanation': explanation,
                **{detail['Principle']: detail['Score'] for detail in details}
            }
            df = pd.DataFrame([data])

            # Convert DataFrame to CSV for download
            csv = df.to_csv(index=False)
            st.download_button(
                label="Download evaluation as CSV",
                data=csv,
                file_name='evaluation.csv',
                mime='text/csv',
            )
        else:
            st.error('Please enter both a question and an explanation to evaluate.')

if __name__ == '__main__':
    if 'password_verified' not in st.session_state or not st.session_state['password_verified']:
        check_password()
    else:
        main()