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"
]

# Define reverse-scored items for each trait
traits_reverse_map = {
    'Extraversion': [6, 21, 31],
    'Agreeableness': [2, 12, 27, 37],
    'Conscientiousness': [8, 18, 23, 43],
    'Neuroticism': [9, 24, 34],
    'Openness': [35, 41]
}

# 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,
        'min_score': 2,   # (1*5) + (12 - (5*3)) = 5 + (12 - 15) = 5 - 3 = 2
        'max_score': 34   # (5*5) + (12 - (1*3)) = 25 + (12 - 3) = 25 + 9 = 34
    },
    '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,
        'min_score': 1,    # (1*5) + (16 - (5*4)) = 5 + (16 - 20) = 5 - 4 = 1
        'max_score': 37    # (5*5) + (16 - (1*4)) = 25 + (16 - 4) = 25 + 12 = 37
    },
    '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,
        'min_score': 1,    # (1*5) + (16 - (5*4)) = 5 + (16 - 20) = 5 - 4 = 1
        'max_score': 37    # (5*5) + (16 - (1*4)) = 25 + (16 - 4) = 25 + 12 = 37
    },
    '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,
        'min_score':2,     # (1*5) + (12 - (5*3)) = 5 + (12 - 15) = 5 - 3 = 2
        'max_score':34     # (5*5) + (12 - (1*3)) = 25 + (12 - 3) = 25 + 9 = 34
    },
    '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,
        'min_score':6,     # (1*8) + (8 - (5*2)) = 8 + (8 - 10) = 8 - 2 = 6
        'max_score':46     # (5*8) + (8 - (1*2)) = 40 + (8 - 2) = 40 + 6 = 46
    }
}

# Define explanations for each trait based on score category
explanations = {
    'Extraversion': {
        'low': "You are more reserved and prefer solitary activities. You might find large social gatherings draining and enjoy deep, meaningful interactions over casual conversations.",
        'medium_low': "You have a balanced approach to social interactions. While you enjoy spending time with others, you also value your alone time and can adapt to different social settings.",
        'medium_high': "You are fairly outgoing and enjoy social interactions, but you also appreciate moments of solitude. You strike a good balance between being sociable and introspective.",
        'high': "You are highly outgoing, energetic, and enjoy being around people. You thrive in social situations and are often perceived as enthusiastic and lively."
    },
    'Agreeableness': {
        '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.",
        'medium_low': "You balance assertiveness with cooperativeness. While you can stand up for yourself, you also recognize the value of collaboration and compromise.",
        'medium_high': "You are generally cooperative and considerate but also maintain your own opinions and boundaries. You strive to get along with others while asserting your needs when necessary.",
        '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."
    },
    'Conscientiousness': {
        '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.",
        'medium_low': "You exhibit a moderate level of organization and dependability. You can be both disciplined and adaptable, adjusting your approach based on the situation.",
        'medium_high': "You are generally organized and dependable but also allow for flexibility when needed. You balance meticulousness with adaptability, ensuring tasks are completed while accommodating changes.",
        '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."
    },
    'Neuroticism':{
        'low': "You are generally calm, resilient, and emotionally stable. You handle stress well and are less likely to experience negative emotions intensely.",
        'medium_low': "You maintain a good level of emotional stability but may occasionally experience stress or negative emotions. Overall, you manage your emotions effectively.",
        'medium_high': "You experience emotional fluctuations more frequently but still retain the ability to manage and cope with stress. You are aware of your emotions and work towards maintaining balance.",
        '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."
    },
    'Openness':{
        '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.",
        'medium_low': "You have a balanced approach to new experiences. While you appreciate some creativity and novelty, you also value tradition and established methods.",
        'medium_high': "You are quite open to new experiences and enjoy exploring creative and intellectual pursuits. You balance your curiosity with practicality, allowing you to adapt and innovate when necessary.",
        'high': "You are imaginative, curious, and open to new experiences. You appreciate art, creativity, and value intellectual exploration and novelty."
    }
}

# 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))
    
    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:
            # Reverse the reverse-scored items
            reversed_rev_values = []
            for r in rev_items:
                if r is not None:
                    reversed_rev_values.append(6 - r)  # Reverse the score: 1↔5, 2↔4, 3↔3
            # Compute means, ignoring None
            pos_values = [x for x in pos_items if x is not None]
            rev_values = [x for x in reversed_rev_values 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

    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} (Range: {traits[trait]['min_score']} - {traits[trait]['max_score']})\n\n"
            # Determine the score category based on the trait's score range
            min_s = traits[trait]['min_score']
            max_s = traits[trait]['max_score']
            range_third = (max_s - min_s) / 3
            mid_low = min_s + range_third
            mid_high = min_s + 2 * range_third

            if score < mid_low:
                category = 'low'
            elif score < mid_high:
                category = 'medium_low'
            elif score < max_s:
                category = 'medium_high'
            else:
                category = 'high'

            # Assign explanations based on the category
            explanation = explanations[trait].get(category, "")
            markdown_output += f"{explanation}\n\n"
            trait_names.append(trait)
            trait_scores.append(score)
    
    # Generate bar chart
    image = None
    if trait_scores:
        fig, ax = plt.subplots(figsize=(12, 7))
        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': traits['Extraversion']['positive'] + traits['Extraversion']['reverse'],
            'Agreeableness': traits['Agreeableness']['positive'] + traits['Agreeableness']['reverse'],
            'Conscientiousness': traits['Conscientiousness']['positive'] + traits['Conscientiousness']['reverse'],
            'Neuroticism': traits['Neuroticism']['positive'] + traits['Neuroticism']['reverse'],
            'Openness': traits['Openness']['positive'] + traits['Openness']['reverse']
        }

        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.get(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

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