import gradio as gr
import matplotlib.pyplot as plt
from io import BytesIO
from PIL import Image as PILImage

# Define the 44 BFI questions
questions = [
    "1. Talks a lot",
    "2. Notices other people’s weak points",
    "3. Does things carefully and completely",
    "4. Is sad, depressed",
    "5. Is original, comes up with new ideas",
    "6. Keeps their thoughts to themselves",
    "7. Is helpful and not selfish with others",
    "8. Can be kind of careless",
    "9. Is relaxed, handles stress well",
    "10. Is curious about lots of different things",
    "11. Has a lot of energy",
    "12. Starts arguments with others",
    "13. Is a good, hard worker",
    "14. Can be tense; not always easy going",
    "15. Clever; thinks a lot",
    "16. Makes things exciting",
    "17. Forgives others easily",
    "18. Isn’t very organized",
    "19. Worries a lot",
    "20. Has a good, active imagination",
    "21. Tends to be quiet",
    "22. Usually trusts people",
    "23. Tends to be lazy",
    "24. Doesn’t get upset easily; steady",
    "25. Is creative and inventive",
    "26. Has a good, strong personality",
    "27. Can be cold and distant with others",
    "28. Keeps working until things are done",
    "29. Can be moody",
    "30. Likes artistic and creative experiences",
    "31. Is kind of shy",
    "32. Kind and considerate to almost everyone",
    "33. Does things quickly and carefully",
    "34. Stays calm in difficult situations",
    "35. Likes work that is the same every time",
    "36. Is outgoing; likes to be with people",
    "37. Is sometimes rude to others",
    "38. Makes plans and sticks to them",
    "39. Gets nervous easily",
    "40. Likes to think and play with ideas",
    "41. Doesn’t like artistic things (plays, music)",
    "42. Likes to cooperate; goes along with others",
    "43. Has trouble paying attention",
    "44. Knows a lot about art, music and books"
]

# Scoring function based on the provided SPSS syntax
def compute_bfi_scores(*args):
    responses = list(args)
    # Convert 'No response' to None, else to int
    processed = []
    for r in responses:
        if r == "No response":
            processed.append(None)
        else:
            processed.append(int(r))
    
    # Define traits with their respective items and scoring parameters
    traits = {
        'Extraversion': {
            'positive': [1, 11, 16, 26, 36],
            'reverse': [6, 21, 31],
            'threshold': 1,
            'formula_pos_mult': 5,
            'formula_reverse_mult': 3,
            'formula_reverse_const': 12
        },
        'Agreeableness': {
            'positive': [7, 17, 22, 32, 42],
            'reverse': [2, 12, 27, 37],
            'threshold': 1,
            'formula_pos_mult': 5,
            'formula_reverse_mult':4,
            'formula_reverse_const':16
        },
        'Conscientiousness': {
            'positive': [3, 13, 28, 33, 38],
            'reverse': [8, 18, 23, 43],
            'threshold': 1,
            'formula_pos_mult':5,
            'formula_reverse_mult':4,
            'formula_reverse_const':16
        },
        'Neuroticism':{
            'positive':[4, 14, 19, 29, 39],
            'reverse':[9, 24, 34],
            'threshold':1,
            'formula_pos_mult':5,
            'formula_reverse_mult':3,
            'formula_reverse_const':12
        },
        'Openness':{
            'positive':[5, 10, 15, 20, 25, 30, 40, 44],
            'reverse':[35, 41],
            'threshold':2,
            'formula_pos_mult':8,
            'formula_reverse_mult':2,
            'formula_reverse_const':8
        }
    }

    scores = {}

    for trait, info in traits.items():
        pos_items = [processed[i-1] for i in info['positive']]
        rev_items = [processed[i-1] for i in info['reverse']]

        missing_pos = pos_items.count(None)
        missing_rev = rev_items.count(None)
        total_missing = missing_pos + missing_rev

        if total_missing > info['threshold']:
            scores[trait] = "Incomplete"
        else:
            # Compute means, ignoring None
            pos_values = [x for x in pos_items if x is not None]
            rev_values = [x for x in rev_items if x is not None]

            mean_pos = sum(pos_values) / len(pos_values) if pos_values else 0
            mean_rev = sum(rev_values) / len(rev_values) if rev_values else 0

            # Apply the scoring formula
            score = (mean_pos * info['formula_pos_mult']) + (info['formula_reverse_const'] - (mean_rev * info['formula_reverse_mult']))
            score = round(score, 2)
            scores[trait] = score

    # Prepare the output in Markdown format with explanations
    explanations = {
        'Extraversion': {
            'high': "You are highly outgoing, energetic, and enjoy being around people. You thrive in social situations and are often perceived as enthusiastic and lively.",
            'low': "You are more reserved and prefer solitary activities. You might find large social gatherings draining and enjoy deep, meaningful interactions over casual conversations."
        },
        'Agreeableness': {
            'high': "You are compassionate, cooperative, and value getting along with others. You tend to be trusting and considerate, often putting others' needs before your own.",
            'low': "You are more competitive and skeptical, prioritizing your own needs and viewpoints. You might be seen as direct or even confrontational in your interactions."
        },
        'Conscientiousness': {
            'high': "You are organized, dependable, and have a strong sense of duty. You strive for achievement and are meticulous in your work, often planning ahead and following through on commitments.",
            'low': "You are more spontaneous and flexible, potentially preferring to adapt as situations arise rather than sticking to a strict plan. You might find rigid structures stifling."
        },
        'Neuroticism':{
            'high': "You tend to experience emotions like anxiety, sadness, and mood swings more frequently. You might be more sensitive to stress and prone to feeling overwhelmed.",
            'low': "You are generally calm, resilient, and emotionally stable. You handle stress well and are less likely to experience negative emotions intensely."
        },
        'Openness':{
            'high': "You are imaginative, curious, and open to new experiences. You appreciate art, creativity, and value intellectual exploration and novelty.",
            'low': "You prefer routine and familiarity, valuing practicality and straightforwardness over abstract ideas. You might be more focused on tangible outcomes rather than theoretical concepts."
        }
    }

    markdown_output = "## Your Big Five Personality Traits Scores\n\n"

    # Prepare data for visualization
    trait_names = []
    trait_scores = []
    for trait, score in scores.items():
        markdown_output += f"### **{trait}**\n"
        if score == "Incomplete":
            markdown_output += "Insufficient responses to compute this trait.\n\n"
        else:
            markdown_output += f"**Score**: {score}\n\n"
            if score >= (max_score(trait)):
                markdown_output += f"{explanations[trait]['high']}\n\n"
            else:
                markdown_output += f"{explanations[trait]['low']}\n\n"
            trait_names.append(trait)
            trait_scores.append(score)
    
    # Function to determine if a score is high or low based on possible range
    def max_score(trait):
        # Define maximum possible scores based on formula
        trait_formula = traits[trait]
        pos_count = len(trait_formula['positive'])
        rev_count = len(trait_formula['reverse'])
        # Maximum score when pos_items are max (5) and rev_items are min (1)
        max_possible = (5 * trait_formula['formula_pos_mult'] * pos_count / pos_count) + (trait_formula['formula_reverse_const'] - (1 * trait_formula['formula_reverse_mult'] * rev_count / rev_count))
        return max_possible - (max_possible / 2)  # Arbitrary threshold at half the max

    # Generate bar chart
    image = None
    if trait_scores:
        fig, ax = plt.subplots(figsize=(10, 6))
        bars = ax.bar(trait_names, trait_scores, color='skyblue')
        ax.set_ylim(0, max(trait_scores) + 10)
        ax.set_ylabel('Score')
        ax.set_title('Big Five Traits Scores')

        # Add score labels on top of bars
        for bar in bars:
            height = bar.get_height()
            ax.annotate(f'{height}',
                        xy=(bar.get_x() + bar.get_width() / 2, height),
                        xytext=(0, 3),  # 3 points vertical offset
                        textcoords="offset points",
                        ha='center', va='bottom')
        
        plt.tight_layout()
        
        # Save the plot to a PNG image in memory
        with BytesIO() as buf:
            plt.savefig(buf, format='png')
            buf.seek(0)
            # Convert bytes to PIL Image and ensure data is loaded
            image = PILImage.open(buf).copy()
            # No need to close the buffer explicitly; 'with' statement handles it
        
        plt.close(fig)  # Close the figure to free memory

        markdown_output += "### **Trait Scores Visualization**\n\n"

    return markdown_output, image

# Create the Gradio interface
def create_interface():
    with gr.Blocks() as demo:
        gr.Markdown("# Big Five Inventory (BFI) Quiz")
        gr.Markdown(
            """
            Please rate the following statements on a scale from **1 (Disagree a lot)** to **5 (Agree a lot)**.
            If you prefer not to respond to a particular statement, select **'No response'**.
            """
        )

        # Organize questions into expandable sections by trait
        trait_question_map = {
            'Extraversion': [1, 11, 16, 26, 36, 6, 21, 31],  # Positive followed by reverse
            'Agreeableness': [7, 17, 22, 32, 42, 2, 12, 27, 37],
            'Conscientiousness': [3, 13, 28, 33, 38, 8, 18, 23, 43],
            'Neuroticism': [4, 14, 19, 29, 39, 9, 24, 34],
            'Openness': [5, 10, 15, 20, 25, 30, 40, 44, 35, 41]
        }

        inputs = []

        with gr.Accordion("Answer the Questions", open=True):
            for trait, q_indices in trait_question_map.items():
                with gr.Accordion(f"{trait}", open=False):
                    for i in q_indices:
                        q = questions[i-1]
                        # Indicate reverse-scored items
                        if i in traits_reverse_map(trait):
                            q_display = f"{q} (Reverse Scored)"
                        else:
                            q_display = q
                        radio = gr.Radio(
                            choices=["No response", 1, 2, 3, 4, 5],
                            label=q_display,
                            value="No response",
                            interactive=True
                        )
                        inputs.append(radio)

        # Submit button
        submit_btn = gr.Button("Submit", variant="primary")

        # Results display
        with gr.Row():
            markdown_result = gr.Markdown(label="Textual Results")
            image_result = gr.Image(label="Trait Scores Visualization")

        # Link the button to the function
        submit_btn.click(
            fn=compute_bfi_scores, 
            inputs=inputs, 
            outputs=[markdown_result, image_result]
        )

    return demo

# Helper function to determine reverse-scored items for display
def traits_reverse_map(trait):
    traits = {
        'Extraversion': [6, 21, 31],
        'Agreeableness': [2, 12, 27, 37],
        'Conscientiousness': [8, 18, 23, 43],
        'Neuroticism': [9, 24, 34],
        'Openness': [35, 41]
    }
    return traits.get(trait, [])

# Launch the interface
demo = create_interface()
demo.launch()