import streamlit as st
import pandas as pd
import logging
import json
import yaml
import pandas as pd
import numpy as np
from dotenv import load_dotenv

import modeling

def show_launch(placeholder):
    with placeholder.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.

            Furthermore, by using this application, you consent to the collection of anonymous usage data. This data will be used for research purposes and to improve the
            application's functionality. No personal information will be recorded or stored.
        """, unsafe_allow_html=True)

        button_placeholder = st.empty()

        if button_placeholder.button(label='Accept Disclaimer', type='primary', use_container_width=True):
            st.session_state.show_launch = False
            placeholder.empty()
            button_placeholder.empty()

def show_demo(placeholder):

    with placeholder:
        with st.container():
            st.divider()
            st.markdown("""
                ## Try it yourself!
                Our recent research shows that sentence transformer ("AI" models)
                can predict respondent patterns in survey data! The model accurately 
                infers item-correlation with *r* = **.71** 🧨, and shows even higher 
                precision for scale correlations (*r* = **.89** 💥) and reliability 
                coefficients (*r* = **.86** 💣)!
                        
                Try it yourself by defining a scale structure using the input field 
                below and let the **SurveyBot3000** predict the expected response 
                pattern. Use the [YAML](https://yaml.org/) format or follow the structure
                outlined by the preset example.
            """)

            with st.form("my_form"):
                
                input_yaml = st.text_area(
                    label="Questionnaire Structure (YAML-Formatted)",
                    value=st.session_state['input_yaml'],
                    height=250
                )

                st.session_state.results_as_matrix = st.checkbox(
                    label="Result as matrix",
                    help="Results will be list-formated (long) by default. Enable to get (wide-format) matrices."
                )
                
                submitted = st.form_submit_button(
                    label="Get Synthetic Estimates",
                    type="primary",
                    use_container_width=True
                )
                if submitted:
                    
                    try:
                        yaml_dict = yaml.safe_load(input_yaml)
                    except yaml.YAMLError as e:
                        st.error(f"Yikes, you better get your YAML straight! Check https://yaml.org/ for help!")
                        return(None)

                    try:
                        modeling.load_model()
                    except Exception as error:
                        st.error(f"Error while loading model: {error}")
                        st.json(yaml_dict)
                        return(None)

                    try:
                        st.session_state.input_data = modeling.process_yaml_input(yaml_dict)
                    except Exception as error:                        
                        st.error(error)
                        st.json(yaml_dict)
                        return(None)
                
                    try:
                        st.session_state.input_data = modeling.encode_input_data()
                    except Exception as error:
                        st.error(error)
                        st.json(yaml_dict)
                        return(None)

                if 'input_data' in st.session_state:

                    tab1, tab2, tab3 = st.tabs(["Item Correlations", "Scale Correlations", "Scale Reliabilities"])
                    
                    with tab1:                        
                        st.markdown("Θ = Synthetic Item Correlation")
                        df = modeling.synthetic_item_correlations()
                        st.dataframe(df, use_container_width=True)

                    with tab2:
                        st.markdown("Θ = Synthetic Scale Correlation")
                        df = modeling.synthetic_scale_correlations()
                        st.dataframe(df, use_container_width=True)

                    with tab3:
                        st.markdown("alpha (Θ) = Synthetic Reliability Estimate")
                        if np.min(modeling.get_items_per_scale()) < 3:
                            st.error("Please make sure that each scale consits of at least 3 items!")
                        else:                        
                            df = modeling.synthetic_reliabilities()
                            st.dataframe(df, use_container_width=True)

                    if 'yaml_dict' in locals():
                        st.markdown("### Input Structure:")
                        st.json(yaml_dict)

def handle_checkbox_change():
    # Update session state
    st.session_state.checkbox_state = not st.session_state.checkbox_state
    # You can also add additional actions to be triggered by the checkbox here
def initialize():
    load_dotenv()
    logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', level=logging.INFO)

    if 'state_loaded' not in st.session_state:
        
        st.session_state['state_loaded'] = True
        with open('init.json') as json_data:
            st.session_state.update(json.load(json_data))

def main():
    st.set_page_config(page_title='Synthetic Correlations')

    col1, col2 = st.columns([2, 5])
    with col1:
        st.image('logo-130x130.svg')

    with col2:
        st.markdown("# Synthetic Correlations")
        st.markdown("#### Estimate Item and Scale Correlations, as well as Reliability Coefficients based on nothing but Text!")

    st.markdown("""

        📖 **Preprint (Open Access)**: https://osf.io/preprints/psyarxiv/kjuce

        🖊️ **Cite**: *Hommel, B. E., & Arslan, R. C. (2024). Language models accurately infer correlations between psychological items and scales from text alone. https://doi.org/10.31234/osf.io/kjuce*

        🌐 **Project website**: https://synth-science.github.io/surveybot3000/

        💾 **Data**: https://osf.io/z47qs/

        #️⃣ **Social Media**:
            - [Björn Hommel on X/Twitter](https://twitter.com/BjoernHommel)
            - [Ruben Arslan on X/Twitter](https://twitter.com/rubenarslan/)

        The web application is maintained by [magnolia psychometrics](https://www.magnolia-psychometrics.com/).
    """, unsafe_allow_html=True)

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

    if 'input_yaml' not in st.session_state:

        with open('sample_input.yaml', 'r') as file:
            try:
                st.session_state['input_yaml'] = file.read()
            except Exception as error:
                print(error)

    if 'disclaimer' not in st.session_state:
        show_launch(placeholder_launch)
        st.session_state['disclaimer'] = True
    else:
        show_demo(placeholder_demo)

if __name__ == '__main__':
    initialize()
    main()