import time
import random
import logging
import streamlit as st
import pandas as pd

from dotenv import load_dotenv

import utils
import db
import modeling
import plots

def set_if_not_in_session_state(key, value):
    """Helper function to initialize a session state variable if it doesn't exist."""
    if key not in st.session_state:
        st.session_state[key] = value

def initialize():
    """Initialization function to set up logging, load environment variables, and initialize session state variables."""
    load_dotenv()
    logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', level=logging.INFO)
    
    keys = ['selected_rating', 'collect_data', 'gender_value', 'expert_value', 'show_launch', 'user_id', 'statements', 'current_statement', 'db']
    values = [0, None, None, None, True, random.randint(1, 999_999_999), None, None, None]
    
    for key, value in zip(keys, values):
        set_if_not_in_session_state(key, value)

    connect_to_database()

def connect_to_database():
    """Establishes a connection to the database."""
    if st.session_state.db is None:
        credentials_dict = db.load_credentials()
        connection_attempts = 0

        while st.session_state.db is None and connection_attempts < 3:
            st.session_state.db = db.connect_to_db(credentials_dict)
            if st.session_state.db is None:
                logging.info('Retrying to connect to db...')
                connection_attempts += 1
                time.sleep(1)
    else:
        retrieve_statements()

def retrieve_statements():
    """Retrieves statements from the database."""
    retrieval_attempts = 0

    while st.session_state.statements is None and retrieval_attempts < 3:
        st.session_state.statements = db.get_statements_from_db(st.session_state.db)
        st.session_state.current_statement = db.pick_random(st.session_state.statements)
        if st.session_state.statements is None:
            logging.info('Retrying to retrieve statements from db...')
            retrieval_attempts += 1
            time.sleep(1)

def get_user_consent():
    st.markdown("""
        ### Support Future Research
        Additionally, we kindly ask for your agreement to collect anonymous data from your app usage in order to improve future research. 
        You may choose to agree or decline this data collection.
    """)

    collect_data_options = ['Yes, I agree and want to support and help improve this research', 'No']
    collect_data_input = st.radio(
        label='You may choose to agree or decline this data collection.',
        options=collect_data_options,
        horizontal=True,
        label_visibility='collapsed'
    )
    return collect_data_options.index(collect_data_input) == 0


def get_user_info():
    gender_options = ['[Please select]', 'Female', 'Male', 'Other']
    gender_input = st.selectbox(
        label='Please select your gender',
        options=gender_options,
    )
    gender_value = gender_options.index(gender_input)

    expert_options = [
        '[Please select]', 
        'No, I do not have a background in social or behavioral sciences', 
        'Yes, I have either studied social or behavioral sciences or I am currently a student in this field', 
        'Yes, I have either worked as a researcher in the field of social or behavioral sciences or I have had past experience as a researcher in this area'
    ]
    expert_input = st.selectbox(
        label='Please indicate whether you have any experience or educational background in social or behavioral sciences (e.g., psychology)',
        options=expert_options,
    )
    expert_value = expert_options.index(expert_input)
    
    return expert_value, gender_value

def get_user_rating(placeholder):

    with placeholder:
        with st.container():
            st.markdown(f"""
                ### How desirable is the following statement?
                To support future research, rate the following statement according to whether it is socially desirable or undesirable.
                Is it socially desirable or undesirable to endorse the following statement?
                #### <center>\"{st.session_state.current_statement.capitalize()}\"</center>
            """, unsafe_allow_html=True)
            
            rating_options = ['[Please select]', 'Very undesirable', 'Undesirable', 'Neutral', 'Desirable', 'Very desirable']
            
            selected_rating = st.selectbox(
                label='Rate the statement above according to whether it is socially desirable or undesirable.',
                options=rating_options,
                key='selection'
            )

            suitability_options = ['No, I\'m just playing around', 'Yes, my input can help improve this research']

            research_suitability = st.radio(
                label='Is your input suitable for research purposes?',
                options=suitability_options,
                horizontal=True
            )

            st.session_state.collect_data_optout = st.checkbox(
                label='Don\'t ask me to rate further statements.',
                value=False
            )

            st.session_state.item_rating = rating_options.index(selected_rating)
            st.session_state.suitability_rating = suitability_options.index(research_suitability)

def handle_acceptance(collect_data_value, expert_value, gender_value, message):
    if st.button(label='Accept Disclaimer', type='primary', use_container_width=True):
        if collect_data_value and not (expert_value > 0 and gender_value > 0):
            message.error('Please answer the questions above!')
        else:
            st.session_state.expert_value = expert_value
            st.session_state.gender_value = gender_value
            st.session_state.show_launch = False
            st.session_state.collect_data = collect_data_value
            st.experimental_rerun()

def show_launch(placeholder):
    with placeholder:
        with st.container():
            st.divider()
            st.markdown("""
                ## Before Using the App
                ### Disclaimer
                This application is provided as-is, without any warranty or guarantee of any kind, expressed or implied. It is intended for educational, non-commercial use only. 
                The developers of this app shall not be held liable for any damages or losses incurred from its use. By using this application, you agree to the terms and conditions 
                outlined herein and acknowledge that any commercial use or reliance on its functionality is strictly prohibited.    
            """)

            collect_data_value = False
            if st.session_state.db:
                collect_data_value = get_user_consent()

            expert_value, gender_value = (0, 0)
            if collect_data_value:
                expert_value, gender_value = get_user_info()

            message = st.empty()

            handle_acceptance(collect_data_value, expert_value, gender_value, message)

def show_summary(placeholder):
    with placeholder:
        with st.container():
            st.markdown("""
                ## What is the focus of this research?
                Certain biases can affect how people respond to surveys and psychological questionnaires. 
                For example, survey respondents may attempt to conceal socially undesirable traits (e.g., 
                being ill-tempered) and endorse statements that cast them in a favorable manner (e.g., 
                being cooperative).

                Developers of psychological questionnaires hence sometimes aim to ensure that questions 
                are neutral, or that a subset of questions is equally (un)desirable. In the past, human 
                judges have been tasked with quantifying item desirability. In contrast, the research 
                underlying this web application demonstrates that large language models (LLMs) can 
                achieve this too!
            """)

def handle_demo_input():
    
    if st.session_state.collect_data:
        if st.session_state.item_rating > 0:

            st.session_state.sentiment, st.session_state.desirability = modeling.score_text(st.session_state.input_text)

            payload = {
                'user_id': st.session_state.user_id,
                'gender_value': st.session_state.gender_value,
                'expert_value': st.session_state.expert_value,
                'statement': st.session_state.current_statement,
                'rating': st.session_state.item_rating,
                'suitability': st.session_state.suitability_rating,
                'input_text': st.session_state.input_text,
                'sentiment': st.session_state.sentiment,
                'desirability': st.session_state.desirability,
            }
            write_to_db_success = db.write_to_db(st.session_state.db, payload)

            if st.session_state.collect_data_optout:
                st.session_state.collect_data = False

            if write_to_db_success:
                st.session_state.current_statement = db.pick_random(st.session_state.statements)
                st.session_state.selection = '[Please select]'
        else:
            return None
    else:
        st.session_state.sentiment, st.session_state.desirability = modeling.score_text(st.session_state.input_text)
        

def show_demo(placeholder):
    with placeholder:
        with st.container():
            st.divider()            
            st.markdown("""
                ## Try it yourself!
                Use the text field below to enter a statement that might be part of a psychological 
                questionnaire (e.g., "I love a good fight."). Your input will be processed by 
                language models, returning a machine-based estimate of item sentiment (i.e., valence) 
                and desirability.

            """)
            modeling.load_model()

            if 'sentiment' in st.session_state and 'desirability' in st.session_state:                    
                plots.show_scores(
                    sentiment=st.session_state.sentiment, 
                    desirability=st.session_state.desirability, 
                    input_text=st.session_state.input_text
                )

            st.session_state.input_text = st.text_input(
                label='Item text/statement:',
                value='I love a good fight.',
                placeholder='Enter item text'
            )        

            user_rating_placeholder = st.empty()

            if st.session_state.collect_data:                
                get_user_rating(user_rating_placeholder)

            if st.button(
                label='Evaluate Item Text',
                on_click=handle_demo_input,
                type='primary',
                use_container_width=True
            ):
                if st.session_state.collect_data and st.session_state.item_rating == 0:
                    st.error('Please rate the statement presented above!')

def show_data(placeholder):
    with placeholder:
        with st.container():
            st.divider()   
            st.markdown("""
                ## Explore the data
                Figures show the accuarcy in precitions of human-rated item desirability by 
                the sentiment model (left) and the desirability model (right), using 
                `test`-partition data only. 
            """)

            show_covariates = st.checkbox('Show covariates', value=True)
            if show_covariates:
                option = st.selectbox('Group by', options=list(utils.covariate_columns.values()))
            else:
                option = None
            
            if 'df' not in st.session_state:
                utils.load_data()

            plot = plots.scatter_plot(st.session_state.df, option)
            st.plotly_chart(plot, theme=None, use_container_width=True)

def main():
    st.markdown("""
        # Machine-Based Item Desirability Ratings
        This web application demonstrates how item desirability ratings can be obtained with natural language processing ("AI"), and accompanying the paper "*Expanding the Methodological Toolbox: Machine-Based Item Desirability Ratings as an Alternative to Human-Based Ratings*".
        
        *Hommel, B. E. (2023). Expanding the methodological toolbox: Machine-based item desirability ratings as an alternative to human-based ratings. Personality and Individual Differences, 213, 112307. https://doi.org/10.1016/j.paid.2023.112307*

        <small>https://www.magnolia-psychometrics.com/</small>
    """, unsafe_allow_html=True)

    placeholder_launch = st.empty()
    placeholder_summary = st.empty()
    placeholder_demo = st.empty()
    placeholder_data = st.empty()

    if st.session_state.show_launch is True:
        show_launch(placeholder_launch)
    else:
        placeholder_launch = st.empty()
        show_summary(placeholder_summary)
        show_demo(placeholder_demo)
        show_data(placeholder_data)


if __name__ == '__main__':

    initialize()
    main()