import streamlit as st
import openai
import json
import os
from dotenv import load_dotenv

# Load the OpenAI API Key
api_key = st.text_input('Enter your OpenAI API Key', type="password")

# Set the OpenAI API key
if api_key:
    openai.api_key = api_key

# English-translated version of the questions (MBTI-related questions)
questions = [
    {"text": "Do you enjoy being spontaneous and keeping your options open?", "trait": "P"},
    {"text": "Do you prefer spending weekends quietly at home rather than going out?", "trait": "I"},
    {"text": "Do you feel more energized when you are around people?", "trait": "E"},
    {"text": "Do you easily set and meet deadlines?", "trait": "J"},
    {"text": "Are your decisions often influenced by how they will affect others emotionally?", "trait": "F"},
    {"text": "Do you like discussing symbolic or metaphorical interpretations of a story?", "trait": "N"},
    {"text": "Do you strive to maintain harmony in group settings, even if it means compromising?", "trait": "F"},
    {"text": "When a friend is upset, is your first instinct to offer emotional support rather than solutions?", "trait": "F"},
    {"text": "In arguments, do you focus more on being rational than on people's feelings?", "trait": "T"},
    {"text": "When you learn something new, do you prefer hands-on experience over theory?", "trait": "S"},
    {"text": "Do you often think about how today's actions will affect the future?", "trait": "N"},
    {"text": "Are you comfortable adapting to new situations as they happen?", "trait": "P"},
    {"text": "Do you prefer exploring different options before making a decision?", "trait": "P"},
    {"text": "At parties, do you start conversations with new people?", "trait": "E"},
    {"text": "When faced with a problem, do you prefer discussing it with others?", "trait": "E"},
    {"text": "When making decisions, do you prioritize logic over personal considerations?", "trait": "T"},
    {"text": "Do you find solitude more refreshing than social gatherings?", "trait": "I"},
    {"text": "Do you prefer having a clear plan and dislike unexpected changes?", "trait": "J"},
    {"text": "Do you find satisfaction in finishing tasks and making final decisions?", "trait": "J"},
    {"text": "Do you tend to process your thoughts internally before speaking?", "trait": "I"},
    {"text": "Are you more interested in exploring abstract theories and future possibilities?", "trait": "N"},
    {"text": "When planning a vacation, do you prefer to have a detailed plan?", "trait": "S"},
    {"text": "Do you often rely on objective criteria to assess situations?", "trait": "T"},
    {"text": "Do you focus more on details and facts in your surroundings?", "trait": "S"}
]

# Function to calculate MBTI scores based on responses
def calculate_weighted_mbti_scores(responses):
    weights = {
        "Strongly Agree": 2,
        "Agree": 1,
        "Neutral": 0,
        "Disagree": -1,
        "Strongly Disagree": -2
    }

    scores = {'E': 0, 'I': 0, 'S': 0, 'N': 0, 'T': 0, 'F': 0, 'J': 0, 'P': 0}

    for i, response in enumerate(responses):
        weight = weights.get(response, 0)
        trait = questions[i]["trait"]
        if trait in scores:
            scores[trait] += weight

    return scores

# Function to determine MBTI type based on weighted scores
def classic_mbti_weighted(responses):
    scores = calculate_weighted_mbti_scores(responses)
    mbti_type = ""
    for trait_pair in ['EI', 'SN', 'TF', 'JP']:
        trait1, trait2 = trait_pair
        if scores[trait1] >= scores[trait2]:
            mbti_type += trait1
        else:
            mbti_type += trait2
    return mbti_type

# Function to save responses to a JSON file
def save_responses_to_json(username, responses):
    user_data = {
        "username": username,
        "responses": [{"text": question["text"], "answer": response} for question, response in zip(questions, responses)]
    }

    # Save to UserChoices.json
    with open("UserChoices.json", "w") as json_file:
        json.dump(user_data, json_file, indent=4)

# Function to save personality results to Output.json
def save_personality_to_output_json(username, mbti_type_classic, mbti_type_llm):
    output_data = {
        "username": username,
        "mbti_type_classic": mbti_type_classic,
        "mbti_type_llm": mbti_type_llm
    }

    # Save to Output.json
    with open("Output.json", "w") as json_file:
        json.dump(output_data, json_file, indent=4)

# Streamlit component to display the quiz and handle responses
def show_mbti_quiz():
    st.title('FlexTemp Personality Test')

    # Step 1: Input name
    participant_name = st.text_input("Enter your name")

    if participant_name:
        responses = []
        st.subheader(f"Hello {participant_name}, let's start the quiz!")

        for i, question in enumerate(questions):
            response = st.radio(
                question["text"],
                ["Strongly Agree", "Agree", "Neutral", "Disagree", "Strongly Disagree"]
            )
            if response:
                responses.append(response)

        if len(responses) == len(questions):
            # Add a button to generate personality information
            if st.button("Generate Personality Trait Information"):
                st.subheader("Your MBTI Personality Type:")
                mbti_type_classic = classic_mbti_weighted(responses)
                st.write(f"Your MBTI type based on weighted answers: {mbti_type_classic}")

                # You can add LLM-based prediction if needed here (example OpenAI-based model)
                if api_key:
                    # Run the LLM (GPT-4, for example) model to generate a personality type.
                    prompt = f"""
                    Determine a person's personality type based on their answers to the following Myers-Briggs Type Indicator (MBTI) questions:
                    The person has answered the following questions:
                    {', '.join([f"{question['text']} {response}" for question, response in zip(questions, responses)])}
                    What is the MBTI personality type based on these answers?
                    """
                    try:
                        response = openai.ChatCompletion.create(
                            model="gpt-4o",
                            messages=[{"role": "system", "content": "You are a helpful assistant."},
                                      {"role": "user", "content": prompt}]
                        )
                        mbti_type_llm = response['choices'][0]['message']['content']
                        st.write(f"Your MBTI type according to AI: {mbti_type_llm}")
                    except Exception as e:
                        st.error(f"Error occurred: {e}")

                # Save responses and personality info to Output.json
                save_responses_to_json(participant_name, responses)
                save_personality_to_output_json(participant_name, mbti_type_classic, mbti_type_llm)

                with open("Output.json", "r") as json_file:
                    json_data = json_file.read()

                st.download_button(
                    label="Download Output.json",
                    data=json_data,
                    file_name="Output.json",
                    mime="application/json"
                )

                with open("UserChoices.json", "r") as json_file:
                    json_data = json_file.read()

                st.download_button(
                    label="Download UserChoices.json",
                    data=json_data,
                    file_name="UserChoices.json",
                    mime="application/json"
                )

        else:
            st.warning("Please answer all the questions!")

# Main function to display the app
def main():
    # Add instructions to the sidebar
    with st.sidebar.expander("How This App Works", expanded=False):
        st.write("""
            ### FlexTemp Personality Test
            This app is designed to help you determine your MBTI personality type based on your answers to a series of questions. The process works as follows:
            1. **Weighted MBTI Scoring**: 
            - Each question corresponds to a trait in the MBTI system. 
            - Your responses are scored on a scale from "Strongly Agree" to "Strongly Disagree", with each level being assigned a weight. 
            - These weights are used to calculate your MBTI type by comparing the scores of trait pairs (E/I, S/N, T/F, J/P).
            2. **LLM-Based Prediction**:
            - Optionally, you can also get your MBTI type based on the answers using a language model (LLM) like GPT-4. This provides an additional prediction that may offer insights into your personality.
            - The LLM is trained on vast amounts of data and can generate responses based on patterns from psychological research and real-world interactions.
            """)
    
    if api_key:
        show_mbti_quiz()
    else:
        st.info("Please enter your OpenAI API Key to begin the quiz.")

if __name__ == "__main__":
    main()