import gradio as gr
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from huggingface_hub import InferenceClient

# WHOQOL-BREF questions (Same as before)
questions = [
    {"id": "Q1", "text": "How would you rate your quality of life?", "domain": "Overall QOL", "scores": ["Very poor", "Poor", "Neither poor nor good", "Good", "Very good"]},
    {"id": "Q2", "text": "How satisfied are you with your health?", "domain": "Overall Health", "scores": ["Very dissatisfied", "Dissatisfied", "Neither satisfied nor dissatisfied", "Satisfied", "Very satisfied"]},
    {"id": "Q3", "text": "To what extent do you feel that physical pain prevents you from doing what you need to do?", "domain": "Physical", "scores": ["Not at all", "A little", "A moderate amount", "Very much", "An extreme amount"], "reverse": True},
    {"id": "Q4", "text": "How much do you need any medical treatment to function in your daily life?", "domain": "Physical", "scores": ["Not at all", "A little", "A moderate amount", "Very much", "An extreme amount"], "reverse": True},
    {"id": "Q5", "text": "How much do you enjoy life?", "domain": "Psychological", "scores": ["Not at all", "A little", "A moderate amount", "Very much", "An extreme amount"]},
    {"id": "Q6", "text": "To what extent do you feel your life to be meaningful?", "domain": "Psychological", "scores": ["Not at all", "A little", "A moderate amount", "Very much", "An extreme amount"]},
    {"id": "Q7", "text": "How well are you able to concentrate?", "domain": "Psychological", "scores": ["Not at all", "A little", "A moderate amount", "Very much", "Extremely"]},
    {"id": "Q8", "text": "How safe do you feel in your daily life?", "domain": "Environment", "scores": ["Not at all", "A little", "A moderate amount", "Very much", "Extremely"]},
    {"id": "Q9", "text": "How healthy is your physical environment?", "domain": "Environment", "scores": ["Not at all", "A little", "A moderate amount", "Very much", "Extremely"]},
    {"id": "Q10", "text": "Do you have enough energy for everyday life?", "domain": "Physical", "scores": ["Not at all", "A little", "Moderately", "Mostly", "Completely"]},
    {"id": "Q11", "text": "Are you able to accept your bodily appearance?", "domain": "Psychological", "scores": ["Not at all", "A little", "Moderately", "Mostly", "Completely"]},
    {"id": "Q12", "text": "Have you enough money to meet your needs?", "domain": "Environment", "scores": ["Not at all", "A little", "Moderately", "Mostly", "Completely"]},
    {"id": "Q13", "text": "How available to you is the information that you need in your day-to-day life?", "domain": "Environment", "scores": ["Not at all", "A little", "Moderately", "Mostly", "Completely"]},
    {"id": "Q14", "text": "To what extent do you have the opportunity for leisure activities?", "domain": "Environment", "scores": ["Not at all", "A little", "Moderately", "Mostly", "Completely"]},
    {"id": "Q15", "text": "How well are you able to get around?", "domain": "Physical", "scores": ["Very poor", "Poor", "Neither poor nor good", "Good", "Very good"]},
    {"id": "Q16", "text": "How satisfied are you with your sleep?", "domain": "Physical", "scores": ["Very dissatisfied", "Dissatisfied", "Neither satisfied nor dissatisfied", "Satisfied", "Very satisfied"]},
    {"id": "Q17", "text": "How satisfied are you with your ability to perform your daily living activities?", "domain": "Physical", "scores": ["Very dissatisfied", "Dissatisfied", "Neither satisfied nor dissatisfied", "Satisfied", "Very satisfied"]},
    {"id": "Q18", "text": "How satisfied are you with your capacity for work?", "domain": "Physical", "scores": ["Very dissatisfied", "Dissatisfied", "Neither satisfied nor dissatisfied", "Satisfied", "Very satisfied"]},
    {"id": "Q19", "text": "How satisfied are you with yourself?", "domain": "Psychological", "scores": ["Very dissatisfied", "Dissatisfied", "Neither satisfied nor dissatisfied", "Satisfied", "Very satisfied"]},
    {"id": "Q20", "text": "How satisfied are you with your personal relationships?", "domain": "Social", "scores": ["Very dissatisfied", "Dissatisfied", "Neither satisfied nor dissatisfied", "Satisfied", "Very satisfied"]},
    {"id": "Q21", "text": "How satisfied are you with your sex life?", "domain": "Social", "scores": ["Very dissatisfied", "Dissatisfied", "Neither satisfied nor dissatisfied", "Satisfied", "Very satisfied"]},
    {"id": "Q22", "text": "How satisfied are you with the support you get from your friends?", "domain": "Social", "scores": ["Very dissatisfied", "Dissatisfied", "Neither satisfied nor dissatisfied", "Satisfied", "Very satisfied"]},
    {"id": "Q23", "text": "How satisfied are you with the conditions of your living place?", "domain": "Environment", "scores": ["Very dissatisfied", "Dissatisfied", "Neither satisfied nor dissatisfied", "Satisfied", "Very satisfied"]},
    {"id": "Q24", "text": "How satisfied are you with your access to health services?", "domain": "Environment", "scores": ["Very dissatisfied", "Dissatisfied", "Neither satisfied nor dissatisfied", "Satisfied", "Very satisfied"]},
    {"id": "Q25", "text": "How satisfied are you with your transport?", "domain": "Environment", "scores": ["Very dissatisfied", "Dissatisfied", "Neither satisfied nor dissatisfied", "Satisfied", "Very satisfied"]},
    {"id": "Q26", "text": "How often do you have negative feelings such as blue mood, despair, anxiety, depression?", "domain": "Psychological", "scores": ["Never", "Seldom", "Quite often", "Very often", "Always"], "reverse": True}
]

# Domain compositions and transformation formula (Same as before)
domain_info = {
    "Physical": {
        "questions": ["Q3", "Q4", "Q10", "Q15", "Q16", "Q17", "Q18"],
        "transform": lambda raw: (raw - 7) * (100/28)
    },
    "Psychological": {
        "questions": ["Q5", "Q6", "Q7", "Q11", "Q19", "Q26"],
        "transform": lambda raw: (raw - 6) * (100/24)
    },
    "Social": {
        "questions": ["Q20", "Q21", "Q22"],
        "transform": lambda raw: (raw - 3) * (100/12)
    },
    "Environment": {
        "questions": ["Q8", "Q9", "Q12", "Q13", "Q14", "Q23", "Q24", "Q25"],
        "transform": lambda raw: (raw - 8) * (100/32)
    },
    "Overall QOL": {
        "questions": ["Q1"],
        "transform": lambda raw: (raw - 1) * (100/4)
    },
    "Overall Health": {
        "questions": ["Q2"],
        "transform": lambda raw: (raw - 1) * (100/4)
    }
}

# Map domain names to full names for display (Same as before)
domain_full_names = {
    "Physical": "Physical Health",
    "Psychological": "Psychological Well-being",
    "Social": "Social Relationships",
    "Environment": "Environment",
    "Overall QOL": "Overall Quality of Life",
    "Overall Health": "Overall Health Satisfaction"
}

# Define interpretation ranges and interventions (Same as before)
interpretation_ranges = {
    "Physical Health": {
        "high_range": 70,
        "medium_range": (50, 69),
        "low_range": 49,
        "high_inference": "Good physical well-being, minimal limitations, good energy and mobility.",
        "medium_inference": "Mild to moderate physical limitations, some pain, or reduced energy.",
        "low_inference": "Significant physical health limitations, pain, fatigue, poor mobility.",
        "high_intervention": "Maintain a healthy lifestyle, regular exercise, balanced nutrition.",
        "medium_intervention": "Increase physical activity, improve sleep, manage minor pain.",
        "low_intervention": "Medical intervention, physical therapy, chronic disease management."
    },
    "Psychological Health": {
        "high_range": 75,
        "medium_range": (50, 74),
        "low_range": 49,
        "high_inference": "Strong emotional resilience, positive self-esteem, low anxiety or stress.",
        "medium_inference": "Moderate psychological health, occasional stress, some mood fluctuations.",
        "low_inference": "High emotional distress, anxiety, depression, low self-worth.",
        "high_intervention": "Continue positive mental health practices, social engagement.",
        "medium_intervention": "Use stress reduction techniques, engage in self-care routines.",
        "low_intervention": "Psychological counseling, cognitive behavioral therapy, medication if needed."
    },
    "Social Relationships": {
        "high_range": 70,
        "medium_range": (50, 69),
        "low_range": 49,
        "high_inference": "Strong social connections, good interpersonal support, high satisfaction.",
        "medium_inference": "Moderate social interactions, may experience loneliness or minor conflicts.",
        "low_inference": "Weak social relationships, lack of support, loneliness, or isolation.",
        "high_intervention": "Sustain social connections, participate in group/community activities.",
        "medium_intervention": "Strengthen personal relationships, seek social support groups.",
        "low_intervention": "Social integration programs, mental health support, relationship therapy."
    },
    "Environment": {
        "high_range": 70,
        "medium_range": (50, 69),
        "low_range": 49,
        "high_inference": "Safe and stable environment, access to resources and healthcare, financial security.",
        "medium_inference": "Some dissatisfaction with living conditions, financial constraints, or safety concerns.",
        "low_inference": "Unstable living conditions, financial stress, poor healthcare access, safety concerns.",
        "high_intervention": "Maintain financial stability, ensure continuous access to healthcare.",
        "medium_intervention": "Identify key environmental issues, seek financial or housing support.",
        "low_intervention": "Government or NGO assistance for financial/housing needs, healthcare access."
    }
}

# Initialize the InferenceClient
client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")  # Replace with a different model if desired


# Function to get interpretation from the LLM
def get_llm_interpretation(domain, score, interpretation_ranges):
    if domain not in interpretation_ranges:
        return "No interpretation available."

    range_info = interpretation_ranges[domain]

    if score >= range_info["high_range"]:
        inference = range_info["high_inference"]
        intervention = range_info["high_intervention"]
        level = "High"
    elif range_info["medium_range"][0] <= score <= range_info["medium_range"][1]:
        inference = range_info["medium_inference"]
        intervention = range_info["medium_intervention"]
        level = "Medium"
    else:
        inference = range_info["low_inference"]
        intervention = range_info["low_intervention"]
        level = "Low"

    prompt = f"""
    Based on the WHOQOL-BREF assessment results for {domain}, the score is {score}, which is categorized as {level}.
    The inference is: {inference}
    Suggested intervention: {intervention}

    Provide a more detailed and personalized interpretation, including potential causes and specific recommendations.

    Response should be concise, empathetic, and helpful.
    """

    try:
        # Call the InferenceClient
        interpretation = client.text_generation(prompt, max_new_tokens=200)  # Adjust max_new_tokens as needed
        return interpretation
    except Exception as e:
        return f"Error getting interpretation from LLM: {e}"


# Function to calculate domain scores (Same as before)
def calculate_scores(responses):
    domain_scores = {}

    for domain, info in domain_info.items():
        domain_questions = info["questions"]
        total_score = 0

        for q_id in domain_questions:
            q_index = next((i for i, q in enumerate(questions) if q["id"] == q_id), None)
            if q_index is not None and responses.get(q_id) is not None:
                if questions[q_index].get("reverse", False):
                    total_score += 6 - responses[q_id]
                else:
                    total_score += responses[q_id]

        transformed_score = info["transform"](total_score)
        domain_scores[domain] = round(transformed_score, 2)

    return domain_scores


# Function to interpret the score range
def interpret_score(domain, score):
    if domain not in interpretation_ranges:
        return "Unknown", "No interpretation available."

    range_info = interpretation_ranges[domain]
    high_range = range_info["high_range"]
    medium_range = range_info["medium_range"]
    low_range = range_info["low_range"]

    if score >= high_range:
        return (
            "High",
            f"{range_info['high_inference']} {range_info['high_intervention']}"
        )
    elif range_info["medium_range"][0] <= score <= range_info["medium_range"][1]:
        return (
            "Medium",
            f"{range_info['medium_inference']} {range_info['medium_intervention']}"
        )
    else:
        return (
            "Low",
            f"{range_info['low_inference']} {range_info['low_intervention']}"
        )

# Define the Gradio interface
def whoqol_assessment(*responses):  # *args to take variable number of responses

    # Map responses to question IDs
    responses_dict = {q["id"]: int(r) for q, r in zip(questions, responses)} # responses from gradio are strings, need convert to int

    # Calculate scores
    scores = calculate_scores(responses_dict)

    if scores is None:
        return "Please answer all questions."  # Error message if not all questions are answered

    # Prepare data for radar chart
    main_domains = ["Physical", "Psychological", "Social", "Environment"]
    main_scores = [scores[domain] for domain in main_domains]

    # Create radar chart (using Matplotlib)
    fig, ax = plt.subplots(figsize=(6, 6), subplot_kw=dict(polar=True))

    # Angle for each domain
    angles = np.linspace(0, 2 * np.pi, len(main_domains), endpoint=False).tolist()
    angles += angles[:1]  # Close the loop

    # Add scores (with loop closure)
    values = main_scores + [main_scores[0]]

    # Plot and fill the radar chart
    ax.plot(angles, values, 'o-', linewidth=2, color='blue')
    ax.fill(angles, values, color='blue', alpha=0.25)

    # Add labels
    domain_labels = [domain_full_names[domain] for domain in main_domains]
    ax.set_xticks(angles[:-1])
    ax.set_xticklabels(domain_labels, fontsize=10)

    # Set y-axis limits
    ax.set_ylim(0, 100)
    ax.set_yticks([0, 25, 50, 75, 100])
    ax.set_yticklabels(['0', '25', '50', '75', '100'], fontsize=8)  # Smaller font

    # Add grid lines
    ax.grid(True, linestyle='-', alpha=0.7)

    # Add title
    ax.set_title('WHOQOL-BREF Domain Scores', size=12, y=1.1) # Smaller Title

    # Detailed Interpretations and store them in text
    interpretations_text = ""
    for domain in ["Physical", "Psychological", "Social", "Environment", "Overall QOL", "Overall Health"]:
        score = scores[domain]
        llm_interpretation = get_llm_interpretation(domain_full_names[domain], score, interpretation_ranges)
        interpretations_text += f"**{domain_full_names[domain]}**: {llm_interpretation}\n\n"

    return fig, interpretations_text


# Create Gradio inputs
inputs = []
for q in questions:
    inputs.append(gr.Radio(choices=q["scores"], label=q["text"]))  # Use Radio instead of Checkbox

# Create Gradio outputs
radar_chart_output = gr.Plot()
interpretations_output = gr.Markdown() # markdown, because I used "**" in text. This format allow to view text in bold format.

# Create Gradio interface
iface = gr.Interface(
    fn=whoqol_assessment,
    inputs=inputs,
    outputs=[radar_chart_output, interpretations_output],
    title="WHOQOL-BREF Quality of Life Assessment",
    description="Complete the questionnaire to receive your quality of life assessment across different domains.",
)

# Launch the Gradio interface
iface.launch()