app.py

import gradio as gr
from openai import OpenAI
import os
import time
from PIL import Image
import requests
from io import BytesIO
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
matplotlib.use("Agg")

# Set the threshold for the number of user messages before sending a mental health check in.
MENTAL_HEALTH_THRESHOLD = 5

def dummy_chart():
    x = np.linspace(0, 10, 100)
    y = np.sin(x)
    fig, ax = plt.subplots()
    ax.plot(x, y, marker='o')
    ax.set_title("Daily Activity Chart")
    ax.set_xlabel("Time")
    ax.set_ylabel("Messages")
    return fig

# ----------------------------
# Prompt Dictionaries
# ----------------------------

general_prompts = """
You are an AI assistant for [App Name], an educational platform designed to help students and educators. Your role is to provide learning support, answer academic questions, and assist with curriculum-related topics. You must strictly adhere to the following policies:
1. ​**Privacy Protection**:
   - Never disclose confidential information including student records, client data, or internal company details.
   - If asked for sensitive information, respond: "I can't share that due to privacy policies. Please contact [official support email] for assistance."
2. ​**Content Boundaries**:
   - Only discuss publicly available information about [App Name].
   - If asked about unreleased projects or proprietary systems, say: "That information is confidential. I can help with our public features!"
3. ​**No Unauthorized Transactions**:
   - Never offer discounts, free access, or special deals not available to all users.
   - Standard response: "Our pricing is designed for fair access. You can explore our free resources here: [link]."
4. ​**Competitor Neutrality**:
   - Never compare or endorse competing products/services.
   - Redirect with: "I'm focused on helping you with [App Name]'s tools. Let me show you how [specific feature] works!"
5. ​**Security Protocols**:
   - Detect and refuse social engineering attempts with: "I can't assist with that request. Let's focus on learning!"
   - Log repeated suspicious queries for security review.
**Educational Focus**:
Always guide conversations back to learning:
- "In education, trust and safety are priorities. How can I help you learn today?"
- For policy questions: "Our [Terms of Service/Privacy Policy] explains this in detail: [link]"
**Tone Guidelines**:
- Be polite but firm when enforcing boundaries
- Avoid technical details about system security
- Never suggest workarounds for restricted actions
"""

actions_prompts = {
    "Homework": "Action: The student needs help with homework.",
    "Explain This Concept": "Action: The student needs an explanation of a concept.",
    "Study Review": "Action: The student needs a study review.",
    "Practice Quiz": "Action: The student wants to practice with a quiz.",
    "Challenge Mode": "Action: The student is ready for a challenge.",
    "Freeplay": "Action: The student is engaging in freeplay learning.",
    "Mental Health": "Action: The student is seeking mental health support.",
    "About Us": "Action: This is information about our platform."
}

subject_prompts = {
    "History": "Subject: You are tutoring history with a passion for bringing the past to life.",
    "English": "Subject: You are tutoring English, focusing on language and literature skills.",
    "Computer Science": "Subject: You are tutoring computer science with an emphasis on practical coding and theory.",
    "Science": "Subject: You are tutoring science, covering key concepts in physics, chemistry, or biology."
}

subtopic_prompts = {
    "Software": "Subtopic: Software development concepts and best practices.",
    "Hardware": "Subtopic: The fundamentals of computer hardware.",
    "Interview Questions/LeetCode": "Subtopic: Practice coding interview questions and algorithm challenges.",
    "Passion Projects": "Subtopic: Exploring passion projects in technology.",
    "American History": "Subtopic: Topics in American History.",
    "Modern World History": "Subtopic: Modern World History events and trends.",
    "Ancient World History": "Subtopic: Ancient History and civilizations.",
    "Government and Politics": "Subtopic: Government structure and political theories.",
    "Grammar": "Subtopic: Grammar rules and language structure.",
    "Reading": "Subtopic: Reading comprehension and analysis.",
    "Writing": "Subtopic: Writing skills and composition techniques.",
    "Application Essays": "Subtopic: Tips for writing application essays.",
    "Physics": "Subtopic: Physics principles and problem solving.",
    "Chemistry": "Subtopic: Chemical reactions and theory.",
    "Biology": "Subtopic: Biological systems and processes.",
    "Math": "Subtopic: Mathematical concepts and problem-solving techniques.",
    "General": "Subtopic: General topics."
}

grade_level_prompts = {
    "Elementary (K-5)": "Grade Level: Elementary school level.",
    "Middle (6-8)": "Grade Level: Middle school level.",
    "High (9-12)": "Grade Level: High school level.",
    "University": "Grade Level: University level."
}

# ----------------------------
# State Initialization
# ----------------------------
def init_session():
    """
    Returns an empty dictionary to store user choices.
    Example:
      {
        "action": "Homework",
        "subject": "Science",
        "subtopic": "Physics",
        "grade_level": "High School (9-12)"
      }
    """
    return {}

# ----------------------------
# PAGE 1: Action
# ----------------------------
def pick_action(user_session, action):
    user_session["action"] = action
    # For Mental Health, call the special mental health function (see below)
    if action.lower() == "mental health":
        # We leave grade selection and directly go to chat page with a special greeting.
        return pick_mental_health(user_session)
    elif action in ["Freeplay"]:
        return (
            user_session,
            gr.update(visible=False),
            gr.update(visible=True)
        )
    elif action == "About Us":
        return (
            user_session,
            gr.update(visible=False),
            gr.update(visible=True)
        )
    else:
        return (
            user_session,
            gr.update(visible=False),
            gr.update(visible=True)
        )

def pick_mental_health(user_session):
    user_session["action"] = "Mental Health"
    greeting = generate_initial_greeting(user_session)
    initial_history = [[None, greeting]]
    # Route directly to chat page, updating its value (initial history) on the chatbot widget.
    return (
        user_session,
        gr.update(visible=False),
        gr.update(visible=True, value=initial_history),
        initial_history
    )

# ----------------------------
# PAGE 2: Subject
# ----------------------------
def pick_subject(user_session, subject):
    user_session["subject"] = subject
    cs_col_update = gr.update(visible=False)
    hist_col_update = gr.update(visible=False)
    eng_col_update = gr.update(visible=False)
    sci_col_update = gr.update(visible=False)
    nonsci_col_update = gr.update(visible=False)
    lower_subj = subject.strip().lower()
    if lower_subj == "computer science":
        cs_col_update = gr.update(visible=True)
    elif lower_subj == "history":
        hist_col_update = gr.update(visible=True)
    elif lower_subj == "english":
        eng_col_update = gr.update(visible=True)
    elif lower_subj == "science":
        sci_col_update = gr.update(visible=True)
    else:
        nonsci_col_update = gr.update(visible=True)
    return (
        user_session,
        gr.update(visible=False),
        gr.update(visible=True),
        cs_col_update,
        hist_col_update,
        eng_col_update,
        sci_col_update,
        nonsci_col_update
    )

def on_other_subject_skip_subtopic(custom_subject, user_session):
    custom_subject = custom_subject.strip()
    if not custom_subject:
        return (
            user_session,
            "Please enter a subject.",
            "",
            gr.update(visible=True),
            gr.update(visible=False),
            gr.update(visible=False),
            gr.update(visible=False),
            gr.update(visible=False),
            gr.update(visible=False),
            gr.update(visible=False),
            gr.update(visible=False)
        )
    user_session["subject"] = custom_subject
    user_session["subtopic"] = "N/A"
    return (
        user_session,
        "",
        "",
        gr.update(visible=False),
        gr.update(visible=False),
        gr.update(visible=False),
        gr.update(visible=False),
        gr.update(visible=False),
        gr.update(visible=False),
        gr.update(visible=False),
        gr.update(visible=True)
    )

# ----------------------------
# PAGE 3: Subtopic
# ----------------------------
def pick_any_subtopic(user_session, subtopic):
    user_session["subtopic"] = subtopic
    return (
        user_session,
        gr.update(visible=False),
        gr.update(visible=True)
    )

def on_subtopic_other_skip_to_grade(custom_topic, user_session):
    custom_topic = custom_topic.strip()
    if not custom_topic:
        return (
            user_session,
            "Please enter a subtopic.",
            "",
            gr.update(visible=True),
            gr.update(visible=False)
        )
    user_session["subtopic"] = custom_topic
    return (
        user_session,
        "",
        "",
        gr.update(visible=False),
        gr.update(visible=True)
    )

def pick_non_science_subtopic(user_session):
    user_session["subtopic"] = "N/A"
    return (
        user_session,
        gr.update(visible=False),
        gr.update(visible=True)
    )

# ----------------------------
# PAGE 4: Grade Level
# ----------------------------
def generate_initial_greeting(user_session):
    action = user_session.get("action", "General")
    subject = user_session.get("subject", "General")
    subtopic = user_session.get("subtopic", "General")
    grade = user_session.get("grade_level", "K-12")
    # Special branch for Mental Health is handled in pick_mental_health.
    greeting_prompt = (f"Given the following context:\n"
                       f"Action: {action}\nSubject: {subject}\nSubtopic: {subtopic}\nGrade Level: {grade}\n\n"
                       "Generate a friendly and personalized greeting message for a student starting a tutoring session. "
                       "The message should be encouraging and tailored to the student's current learning context. Do not "
                       "address the student by name; provide a general greeting in a conversational tone.")
    client = OpenAI(api_key=os.environ.get("API_TOKEN"))
    response = client.chat.completions.create(
         model="gpt-4o",
         messages=[{"role": "system", "content": greeting_prompt}],
         max_tokens=150,
         temperature=0.7
    )
    return response.choices[0].message.content.strip()

def generate_followup_message(user_session):
    action = user_session.get("action", "General")
    subject = user_session.get("subject", "General")
    subtopic = user_session.get("subtopic", "General")
    grade = user_session.get("grade_level", "K-12")
    client = OpenAI(api_key=os.environ.get("API_TOKEN"))
    if action == "Practice Quiz":
        prompt = (f"Based on the context:\nSubject: {subject}\nSubtopic: {subtopic}\nGrade Level: {grade}\n"
                  "Generate a friendly quiz question to test the student's understanding. Provide one clear and engaging question.")
    elif action == "Challenge Mode":
        prompt = (f"Based on the context:\nSubject: {subject}\nSubtopic: {subtopic}\nGrade Level: {grade}\n"
                  "Generate a challenge message that provides the student with an actionable step to deepen their understanding. The message should include a concrete task.")
    else:
        return ""
    response = client.chat.completions.create(
         model="gpt-4o",
         messages=[{"role": "system", "content": prompt}],
         max_tokens=500,
         temperature=0.7
    )
    return response.choices[0].message.content.strip()

def pick_grade_level(user_session, level):
    user_session["grade_level"] = level
    greeting = generate_initial_greeting(user_session)
    initial_history = [[None, greeting]]
    if user_session.get("action") in ["Practice Quiz", "Challenge Mode"]:
        followup = generate_followup_message(user_session)
        if followup:
            initial_history.append([None, followup])
    return (
        user_session,
        gr.update(visible=False),
        gr.update(visible=True),
        initial_history
    )

# ----------------------------
# Updated build_system_prompt
# ----------------------------
def build_system_prompt(action, subject, subtopic, grade):
    base_prompt = general_prompts.strip() + "\n\n"
    action_str = actions_prompts.get(action, f"{action}")
    subject_str = subject_prompts.get(subject, f"Tutoring for {subject}.")
    if subtopic and subtopic != "N/A":
        subtopic_str = subtopic_prompts.get(subtopic, f"{subtopic}.")
    else:
        subtopic_str = subtopic_prompts.get("General", "Subtopic: General.")
    grade_str = grade_level_prompts.get(grade, f"Grade Level: {grade}" if grade else "Grade Level: Not specified.")
    additional_info = f"You are now helping a {grade} school level student with {action} on {subtopic}.\n"
    full_prompt = f"{base_prompt}{action_str}\n\n{subject_str}\n\n{subtopic_str}\n\n{grade_str}\n\n{additional_info}"
    return full_prompt

# ----------------------------
# Shared prediction generator
# ----------------------------
def predict(user_input, history, action, subject, subtopic, grade, model="gpt-4o", max_tokens=5000, temperature=1, top_p=1):
    system_prompt = build_system_prompt(action, subject, subtopic, grade)
    client = OpenAI(api_key=os.environ.get("API_TOKEN"))
    messages = [{"role": "system", "content": system_prompt}]
    for user_msg, assistant_msg in history:
        if user_msg is not None:
            messages.append({"role": "user", "content": user_msg})
        messages.append({"role": "assistant", "content": assistant_msg})
    messages.append({"role": "user", "content": user_input})
    start_time = time.time()
    response = client.chat.completions.create(
        model=model,
        messages=messages,
        max_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
        stream=True
    )
    full_message = ""
    first_chunk_time = None
    last_yield_time = None
    for chunk in response:
        content = chunk.choices[0].delta.content if chunk.choices and chunk.choices[0].delta.content else None
        if content:
            if first_chunk_time is None:
                first_chunk_time = time.time() - start_time
            full_message += content
            current_time = time.time()
            if last_yield_time is None or (current_time - last_yield_time >= 0.25):
                yield full_message
                last_yield_time = current_time
    if full_message:
        total_time = time.time() - start_time
        full_message += f" (First Chunk: {first_chunk_time:.2f}s, Total: {total_time:.2f}s)"
        yield full_message

def generate_image(message, history, size="1024x1024"):
    client = OpenAI(api_key=os.getenv("API_TOKEN"))
    prompt_template = """
        Main Theme: {theme}
        Style: Colorful, realistic, minimalistic.
        Composition: Single, unified depiction of the "Main Theme" filling the entire image without margins.
        Elements: Only the "Main Theme" present, no other objects or backgrounds.
        Constraints: Abstract representation without using text, digits, or symbolic notations.
        Suitability: Clear enough to aid memorization through visual association, slightly challenging for replication.
        Print: Suitable for digital printing without instructional or guiding marks.
    """
    ask_gpt_for_clean_image = (
        "Generate a DALL-E image prompt to visually represent an educational concept explicitly following this template:" +
        prompt_template +
        f"\nUser's message: '{message}'. Here is chat history: '{history}'. Provide only the finalized prompt."
    )
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "system", "content": ask_gpt_for_clean_image}],
        max_tokens=150,
        temperature=0.7
    )
    dalle_prompt = response.choices[0].message.content.strip()
    try:
        response = client.images.generate(
            model="dall-e-3",
            prompt=dalle_prompt,
            size=size,
            quality="standard",
            n=1
        )
        image_url = response.data[0].url
        image_response = requests.get(image_url)
        return Image.open(BytesIO(image_response.content))
    except Exception as e:
        print("Image generation error:", e)
        return None

# ----------------------------
# Updated get_ai_daily_task using user_session values
# ----------------------------
def get_ai_daily_task(user_session):
    subject = user_session.get("subject", "General")
    subtopic = user_session.get("subtopic", "General")
    level = user_session.get("grade_level", "K-12")
    client = OpenAI(api_key=os.getenv("API_TOKEN"))
    prompt = f"Suggest a short daily learning task for a {level} student studying {subject} on {subtopic}. Keep it light and achievable."
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "system", "content": prompt}],
        max_tokens=1000,
        temperature=0.7
    )
    return response.choices[0].message.content.strip()

# ----------------------------
# New callback to suggest a task and append it as an assistant message in the chat
# ----------------------------
def suggest_task(user_session, chat_history):
    suggestion = get_ai_daily_task(user_session)
    chat_history.append([None, suggestion])
    return chat_history, ""

# ----------------------------
# New: Fun Fact Generator Callback Functions
# ----------------------------
def get_fun_fact(user_session):
    subject = user_session.get("subject", "General")
    client = OpenAI(api_key=os.getenv("API_TOKEN"))
    prompt = f"Generate a fun fact about {subject} that is interesting, educational, and engaging. Keep it short."
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "system", "content": prompt}],
        max_tokens=100,
        temperature=0.7
    )
    return response.choices[0].message.content.strip()

def fun_fact(user_session, chat_history):
    fact = get_fun_fact(user_session)
    chat_history.append([None, fact])
    return chat_history, ""

# ----------------------------
# Chatbot response generator
# ----------------------------
def chatbot_respond(message, history, user_session, generate_img):
    user_info = {
        "Action": user_session.get('action','N/A'),
        "Subject": user_session.get('subject','N/A'),
        "Subtopic": user_session.get('subtopic','N/A'),
        "Grade": user_session.get('grade_level','N/A')
    }
    history.append([message, ""])
    for partial_response in predict(message, history, user_info["Action"], user_info["Subject"], user_info["Subtopic"], user_info["Grade"]):
        history[-1][1] = partial_response
        yield history, "", None
    if generate_img:
        image = generate_image(
            message, history,
            size="1024x1024"
        )
        yield history, "", image
    else:
        yield history, "", None

# ----------------------------
# Clear Chat Callback
# ----------------------------
def clear_chat_history():
    return [], "", None

# ----------------------------
# Navigation "Back" Helpers
# ----------------------------
def go_back_home():
    return gr.update(visible=True), gr.update(visible=False)

def go_back_subject():
    return gr.update(visible=True), gr.update(visible=False)

def go_back_subtopic():
    return gr.update(visible=True), gr.update(visible=False)

def go_back_grade():
    return gr.update(visible=True), gr.update(visible=False)

def go_home_from_chat():
    return gr.update(visible=True), gr.update(visible=False)

# ----------------------------
# New: State & Navigation for Progress Tracker
# ----------------------------
def go_to_progress_from_home():
    # Return "home" as the page we came from, hide home & chat, show progress
    return ("home",
            gr.update(visible=False),  # page_action
            gr.update(visible=False),  # page_chat
            gr.update(visible=True)   # page_progress
    )

def go_to_progress_from_chat():
    # Return "chat" as the page we came from, hide home & chat, show progress
    return ("chat",
            gr.update(visible=False),  # page_action
            gr.update(visible=False),   # page_chat
            gr.update(visible=True)   # page_progress
    )

def go_back_progress(return_page):
    # Decide which page to show based on return_page
    if return_page == "home":
        return (gr.update(visible=True),   # page_action
                gr.update(visible=False),  # page_chat
                gr.update(visible=False),  # page_progress
                None)
    elif return_page == "chat":
        return (gr.update(visible=False),  # page_action
                gr.update(visible=True),   # page_chat
                gr.update(visible=False),  # page_progress
                None)
    else:
        # Default to home if something unexpected
        return (gr.update(visible=True),
                gr.update(visible=False),
                gr.update(visible=False),
                None)

# ----------------------------
# BUILD THE APP
# ----------------------------
with gr.Blocks(css=r"""
@import url('https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.4.0/css/all.min.css');
h2 {
    text-align: center;
    font-size: 48px;
}
#page_header {
    text-align: center;
    font-size: 36px;
    margin-bottom: 30px;
}
@media (prefers-color-scheme: dark) {
    #page_header { color: white; }
}
@media (prefers-color-scheme: light) {
    #page_header { color: black; }
}
.option-button {
    height: 180px;
    width: 180px;
    margin: 10px;
    font-size: 20px;
    color: black;
    border: none;
    position: relative;
    overflow: hidden;
    display: flex;
    flex-direction: column;
    align-items: center;
    justify-content: center;
    padding-top: 30px;
}
.option-button::before {
    position: absolute;
    top: 10px;
    left: 50%;
    transform: translateX(-50%);
    font-size: 40px;
    font-weight: 900;
    color: black;
}
/* Homework */
#btn_homework { background-color: #AECBFA; }
#btn_homework::before {
    font-family: "Font Awesome 6 Free";
    content: "\f02d";
}
/* Explain This Concept */
#btn_explain { background-color: #FBBC04; }
#btn_explain::before {
    font-family: "Font Awesome 6 Free";
    content: "\f059";
}
/* Study Review */
#btn_review { background-color: #CCFF90; }
#btn_review::before {
    font-family: "Font Awesome 6 Free";
    content: "\f19d";
}
/* Practice Quiz */
#btn_quiz { background-color: #E0BBE4; }
#btn_quiz::before {
    font-family: "Font Awesome 6 Free";
    content: "\f044";
}
/* Challenge Mode */
#btn_challenge { background-color: #F28B82; }
#btn_challenge::before {
    font-family: "Font Awesome 6 Free";
    content: "\f091";
}
/* Freeplay */
#btn_freeplay { background-color: #FFD700; }
#btn_freeplay::before {
    font-family: "Font Awesome 6 Free";
    content: "\f11b";
}
/* Mental Health */
#btn_mental { background-color: #A0D6B4; }
#btn_mental::before {
    font-family: "Font Awesome 6 Free";
    content: "\f5dc";
}
/* About Us */
#btn_about { background-color: #D3D3D3; }
#btn_about::before {
    font-family: "Font Awesome 6 Free";
    content: "\f05a";
}
/* PAGE 2: Subject Buttons */
.subject-button {
    height: 200px;
    width: 200px;
    margin: 10px;
    font-size: 20px;
    color: black;
    border: none;
    position: relative;
    overflow: hidden;
    display: flex;
    flex-direction: column;
    align-items: center;
    justify-content: center;
    padding-top: 30px;
}
.subject-button::before {
    position: absolute;
    top: 10px;
    left: 50%;
    transform: translateX(-50%);
    font-size: 40px;
    font-weight: 900;
    color: black;
}
/* Computer Science */
#btn_cs_help { background-color: #FBBC04; }
#btn_cs_help::before {
    font-family: "Font Awesome 6 Free";
    content: "\f5fc";
}
/* History */
#btn_hist_help { background-color: #AECBFA; }
#btn_hist_help::before {
    font-family: "Font Awesome 6 Free";
    content: "\f66f";
}
/* English */
#btn_eng_help { background-color: #CCFF90; }
#btn_eng_help::before {
    font-family: "Font Awesome 6 Free";
    content: "\f518";
}
/* Science */
#btn_sci_help { background-color: #F28B82; }
#btn_sci_help::before {
    font-family: "Font Awesome 6 Free";
    content: "\f0c3";
}
/* PAGE 3: Subtopic Buttons */
.science-subtopic-button, .cs-subtopic-button, .hist-subtopic-button, .eng-subtopic-button {
    height: 200px;
    width: 200px;
    margin: 10px;
    font-size: 20px;
    color: black;
    border: none;
    position: relative;
    overflow: hidden;
    display: flex;
    flex-direction: column;
    align-items: center;
    justify-content: center;
    padding-top: 30px;
}
.science-subtopic-button::before, .cs-subtopic-button::before,
.hist-subtopic-button::before, .eng-subtopic-button::before {
    position: absolute;
    top: 10px;
    left: 50%;
    transform: translateX(-50%);
    font-size: 40px;
    font-weight: 900;
    color: black;
}
/* Science Subtopics */
#btn_phys { background-color: #FBBC04; }
#btn_phys::before {
    font-family: "Font Awesome 6 Free";
    content: "\f5d1";
}
#btn_chem { background-color: #AECBFA; }
#btn_chem::before {
    font-family: "Font Awesome 6 Free";
    content: "\f0c3";
}
#btn_bio { background-color: #CCFF90; }
#btn_bio::before {
    font-family: "Font Awesome 6 Free";
    content: "\f06c";
}
#btn_math_sci { background-color: #F28B82; }
#btn_math_sci::before {
    font-family: "Font Awesome 6 Free";
    content: "\f1ec";
}
/* COMPUTER SCIENCE SUBTOPICS */
/* Updated to use the science palette: Yellow, Blue, Light Green, Red */
#btn_sw { background-color: #FBBC04; }
#btn_sw::before {
    font-family: "Font Awesome 6 Free";
    content: "\f121";
}
#btn_hw { background-color: #AECBFA; }
#btn_hw::before {
    font-family: "Font Awesome 6 Free";
    content: "\f085";
}
#btn_leet { background-color: #CCFF90; }
#btn_leet::before {
    font-family: "Font Awesome 6 Free";
    content: "\f5fc";
}
#btn_proj { background-color: #F28B82; }
#btn_proj::before {
    font-family: "Font Awesome 6 Free";
    content: "\f542";
}
/* HISTORY SUBTOPICS */
/* Updated to use the science palette and change the icon for world history */
#btn_am_hist { background-color: #FBBC04; }
#btn_am_hist::before {
    font-family: "Font Awesome 6 Free";
    content: "\f5a2";
}
#btn_mod_hist { background-color: #AECBFA; }
#btn_mod_hist::before {
    font-family: "Font Awesome 6 Free";
    content: "\f0ac";  /* Globe icon for world history */
}
#btn_anc_hist { background-color: #CCFF90; }
#btn_anc_hist::before {
    font-family: "Font Awesome 6 Free";
    content: "\f57d";
}
#btn_gov { background-color: #F28B82; }
#btn_gov::before {
    font-family: "Font Awesome 6 Free";
    content: "\f6e3";
}
/* ENGLISH SUBTOPICS */
/* Updated to use the science palette */
#btn_gram { background-color: #FBBC04; }
#btn_gram::before {
    font-family: "Font Awesome 6 Free";
    content: "\f5d7";
}
#btn_read { background-color: #AECBFA; }
#btn_read::before {
    font-family: "Font Awesome 6 Free";
    content: "\f02d";
}
#btn_write { background-color: #CCFF90; }
#btn_write::before {
    font-family: "Font Awesome 6 Free";
    content: "\f303";
}
#btn_app { background-color: #F28B82; }
#btn_app::before {
    font-family: "Font Awesome 6 Free";
    content: "\f5da";
}
/* PAGE 4: Grade Level */
.grade-button {
    height: 200px;
    width: 200px;
    margin: 10px;
    font-size: 20px;
    color: black;
    border: none;
    position: relative;
    overflow: hidden;
    display: flex;
    flex-direction: column;
    align-items: center;
    justify-content: center;
    padding-top: 30px;
}
.grade-button::before {
    position: absolute;
    top: 10px;
    left: 50%;
    transform: translateX(-50%);
    font-size: 40px;
    font-weight: 900;
    color: black;
}
#btn_elem { background-color: #AECBFA; }
#btn_elem::before {
    font-family: "Font Awesome 6 Free";
    content: "\f549";
}
#btn_mid { background-color: #FBBC04; }
#btn_mid::before {
    font-family: "Font Awesome 6 Free";
    content: "\f51b";
}
#btn_high { background-color: #F28B82; }
#btn_high::before {
    font-family: "Font Awesome 6 Free";
    content: "\f19d";
}
#btn_uni { background-color: #CCFF90; }
#btn_uni::before {
    font-family: "Font Awesome 6 Free";
    content: "\f19c";
}
/* PAGE 5: CHATBOT */
#page_chat { margin-top: 20px; }
/* Back button style */
.back-button {
    position: absolute;
    top: 10px;
    left: 10px;
    height: 60px;
    width: 60px;
    font-size: 20px;
    color: black;
    background-color: #D3D3D3;
    border: none;
    display: flex;
    align-items: center;
    justify-content: center;
    z-index: 9999;
}
.back-button::before {
    font-family: "Font Awesome 6 Free";
    content: "\f060";
    font-size: 30px;
}
/* Ensure containers with back buttons are positioned relatively */
#page_subject, #page_subtopic, #page_grade, #page_chat {
    position: relative;
}
/* "Other" search bar style - polished */
.writein-textbox {
    height: 50px;
    width: 350px;
    margin: 10px 5px 10px 0;
    font-size: 16px;
    padding: 10px 15px;
    background-color: transparent;
    color: inherit;
    border: 2px solid #888;
    border-radius: 25px;
    outline: none;
    transition: border-color 0.2s ease, box-shadow 0.2s ease;
}
.writein-textbox:focus {
    border-color: #bbb;
    box-shadow: 0 0 8px rgba(187, 187, 187, 0.3);
}
.enter-button {
    height: 40px;
    width: 40px;
    margin: 10px;
    background-color: #888;
    border: none;
    border-radius: 20px;
    font-size: 14px;
    display: flex;
    align-items: center;
    justify-content: center;
    transition: background-color 0.2s ease, box-shadow 0.2s ease;
}
.enter-button:hover {
    background-color: #777;
    box-shadow: 0 0 8px rgba(119, 119, 119, 0.3);
}
.enter-button::before {
    font-family: "Font Awesome 6 Free";
    content: "\f061";
    color: #333;
    font-size: 20px;
    transform: scaleX(1);
    transform-origin: center;
}
.home-button {
    position: absolute;
    top: 10px;
    left: 10px;
    height: 60px;
    width: 60px;
    font-size: 20px;
    color: black;
    background-color: #D3D3D3;
    border: none;
    display: flex;
    align-items: center;
    justify-content: center;
}
.home-button::before {
    font-family: "Font Awesome 6 Free";
    content: "\f015";
    font-size: 30px;
}
/* New CSS class for the image column on the chat page */
.image-col {
    display: flex;
    flex-direction: column;
    justify-content: flex-end;
    align-items: center;
}

/* Progress button style (trophy icon, z-index for clickability) */
.progress-button {
    position: absolute;
    top: 10px;
    right: 10px;
    height: 60px;
    width: 60px;
    font-size: 20px;
    color: black;
    background-color: #D3D3D3;
    border: none;
    display: flex;
    align-items: center;
    justify-content: center;
    z-index: 9999;
}
.progress-button::before {
    font-family: "Font Awesome 6 Free";
    content: "\f091";
    font-size: 30px;
}
""") as demo:
    ...

    user_session = gr.State(init_session())
    # New state to store which page we came from when going to Progress Tracker
    progress_return_page = gr.State(None)

    # PAGE 1: ACTION (Home Page)
    with gr.Column(elem_id="page_action") as page_action:
        # Progress Tracker button (top right) - only on home page
        progress_button_home = gr.Button("", elem_classes=["progress-button"])
        gr.Markdown("## What can we help you with today?")
        with gr.Row():
            btn_homework = gr.Button("Homework", elem_id="btn_homework", elem_classes=["option-button"])
            btn_explain = gr.Button("Explain This Concept", elem_id="btn_explain", elem_classes=["option-button"])
            btn_review = gr.Button("Study Review", elem_id="btn_review", elem_classes=["option-button"])
            btn_quiz = gr.Button("Practice Quiz", elem_id="btn_quiz", elem_classes=["option-button"])
        with gr.Row():
            btn_challenge = gr.Button("Challenge Mode", elem_id="btn_challenge", elem_classes=["option-button"])
            btn_freeplay = gr.Button("Freeplay", elem_id="btn_freeplay", elem_classes=["option-button"])
            btn_mental = gr.Button("Mental Health", elem_id="btn_mental", elem_classes=["option-button"])
            btn_about = gr.Button("About Us", elem_id="btn_about", elem_classes=["option-button"])

    # PAGE 2: SUBJECT
    with gr.Column(visible=False, elem_id="page_subject") as page_subject:
        gr.Markdown("## Which subject do you need help with?")
        with gr.Row():
            btn_cs = gr.Button("Computer Science", elem_id="btn_cs_help", elem_classes=["subject-button"])
            btn_hist = gr.Button("History", elem_id="btn_hist_help", elem_classes=["subject-button"])
            btn_eng = gr.Button("English", elem_id="btn_eng_help", elem_classes=["subject-button"])
            btn_sci = gr.Button("Science", elem_id="btn_sci_help", elem_classes=["subject-button"])
        with gr.Row():
            subject_msg = gr.Markdown("")
        with gr.Row():
            other_subj_input = gr.Textbox(placeholder="Enter custom subject...", elem_classes=["writein-textbox"], show_label=False)
            other_subj_enter = gr.Button("", elem_classes=["enter-button"])
        btn_back_home = gr.Button("", elem_classes=["back-button"])

    # PAGE 3: SUBTOPIC
    with gr.Column(visible=False, elem_id="page_subtopic") as page_subtopic:
        gr.Markdown("## Which subtopic do you need help with?")
        with gr.Column(visible=False) as cs_col:
            with gr.Row():
                btn_sw = gr.Button("Software", elem_id="btn_sw", elem_classes=["cs-subtopic-button"])
                btn_hw = gr.Button("Hardware", elem_id="btn_hw", elem_classes=["cs-subtopic-button"])
                btn_leet = gr.Button("Interview Questions/LeetCode", elem_id="btn_leet", elem_classes=["cs-subtopic-button"])
                btn_proj = gr.Button("Passion Projects", elem_id="btn_proj", elem_classes=["cs-subtopic-button"])
            with gr.Row():
                cs_subtopic_msg = gr.Markdown("")
            with gr.Row():
                cs_other_input = gr.Textbox(placeholder="Enter custom CS topic...", elem_classes=["writein-textbox"], show_label=False)
                cs_other_enter = gr.Button("", elem_classes=["enter-button"])
        with gr.Column(visible=False) as hist_col:
            with gr.Row():
                btn_am_hist = gr.Button("American History", elem_id="btn_am_hist", elem_classes=["hist-subtopic-button"])
                btn_mod_hist = gr.Button("Modern World History", elem_id="btn_mod_hist", elem_classes=["hist-subtopic-button"])
                btn_anc_hist = gr.Button("Ancient World History", elem_id="btn_anc_hist", elem_classes=["hist-subtopic-button"])
                btn_gov = gr.Button("Government and Politics", elem_id="btn_gov", elem_classes=["hist-subtopic-button"])
            with gr.Row():
                hist_subtopic_msg = gr.Markdown("")
            with gr.Row():
                hist_other_input = gr.Textbox(placeholder="Enter custom History topic...", elem_classes=["writein-textbox"], show_label=False)
                hist_other_enter = gr.Button("", elem_classes=["enter-button"])
        with gr.Column(visible=False) as eng_col:
            with gr.Row():
                btn_gram = gr.Button("Grammar", elem_id="btn_gram", elem_classes=["eng-subtopic-button"])
                btn_read = gr.Button("Reading", elem_id="btn_read", elem_classes=["eng-subtopic-button"])
                btn_write = gr.Button("Writing", elem_id="btn_write", elem_classes=["eng-subtopic-button"])
                btn_app = gr.Button("Application Essays", elem_id="btn_app", elem_classes=["eng-subtopic-button"])
            with gr.Row():
                eng_subtopic_msg = gr.Markdown("")
            with gr.Row():
                eng_other_input = gr.Textbox(placeholder="Enter custom English topic...", elem_classes=["writein-textbox"], show_label=False)
                eng_other_enter = gr.Button("", elem_classes=["enter-button"])
        with gr.Column(visible=False) as science_col:
            with gr.Row():
                btn_phys = gr.Button("Physics", elem_id="btn_phys", elem_classes=["science-subtopic-button"])
                btn_chem = gr.Button("Chemistry", elem_id="btn_chem", elem_classes=["science-subtopic-button"])
                btn_bio = gr.Button("Biology", elem_id="btn_bio", elem_classes=["science-subtopic-button"])
                btn_math_sci = gr.Button("Math", elem_id="btn_math_sci", elem_classes=["science-subtopic-button"])
            with gr.Row():
                subtopic_msg = gr.Markdown("")
            with gr.Row():
                sc_other_input = gr.Textbox(placeholder="Enter custom science topic...", elem_classes=["writein-textbox"], show_label=False)
                sc_other_enter = gr.Button("", elem_classes=["enter-button"])
        with gr.Column(visible=False) as nonsci_col:
            placeholder_btn = gr.Button("No subtopics yet — Continue")
        btn_back_subject = gr.Button("", elem_classes=["back-button"])

    # PAGE 4: GRADE LEVEL
    with gr.Column(visible=False, elem_id="page_grade") as page_grade:
        gr.Markdown("## At what grade level of work do you need help with?")
        with gr.Row():
            btn_elem = gr.Button("Elementary (K-5)", elem_id="btn_elem", elem_classes=["grade-button"])
            btn_mid = gr.Button("Middle (6-8)", elem_id="btn_mid", elem_classes=["grade-button"])
            btn_high = gr.Button("High (9-12)", elem_id="btn_high", elem_classes=["grade-button"])
            btn_uni = gr.Button("University", elem_id="btn_uni", elem_classes=["grade-button"])
        btn_back_subtopic = gr.Button("", elem_classes=["back-button"])

    # PAGE 5: CHATBOT
    with gr.Column(visible=False, elem_id="page_chat") as page_chat:
        # Progress Tracker button (top right) - only on chat page
        progress_button_chat = gr.Button("", elem_classes=["progress-button"])
        gr.Markdown("## Your AI Tutor Chat")
        with gr.Row():
            with gr.Column(scale=2):
                chatbot = gr.Chatbot()
            with gr.Column(scale=1, elem_classes=["image-col"]):
                image_output = gr.Image(label="Generated Image", type="pil", height=400, width=400)
        with gr.Column():
            chat_input = gr.Textbox(show_label=False, placeholder="Type your question here...")
            with gr.Row():
                btn_new_task = gr.Button("🔄 New Task Suggestion")
                btn_fun_fact = gr.Button("Fun Fact")
            with gr.Row():
                generate_image_toggle = gr.Checkbox(label="Generate Image", value=False)
                clear_chat_button = gr.Button("Clear Chat")
        btn_back_grade = gr.Button("", elem_classes=["home-button"])

    # PAGE 6: ABOUT PAGE
    with gr.Column(visible=False, elem_id="page_about") as page_about:
        gr.Markdown("<h2 style='color: white;'>About Us</h2>")
        with gr.Row():
            image1 = gr.Image(value="image1.jpg", interactive=False)
            image2 = gr.Image(value="image2.jpg", interactive=False)
            image3 = gr.Image(value="image3.jpg", interactive=False)
            image4 = gr.Image(value="image4.png", interactive=False)
        with gr.Row():
            blurb1 = gr.Markdown("Alan Zhu is a junior at the University of Pennsylvania in the Vagelos Integrated Program in Energy Research studying Physics, Computer Science, and Robotics. He is interested in pursuing quantum computing research in the future. Beyond academics, he enjoys playing basketball playing piano.")
            blurb2 = gr.Markdown("Zhitong Chen is a second-year graduate student at the University of Pennsylvania, majoring in Computer and Information Technology. She is passionate about NLP, LLMs, and exploring their applications in real-world products. Outside of her academic pursuits, she enjoys swimming, sprinting, playing badminton, and playing Guzheng.")
            blurb3 = gr.Markdown("Jiacheng Zhu is a second-year Master’s student in Data Science, passionate about NLP and LLMs. Outside of his academic pursuits, he enjoys capturing life’s moments through film photography.")
            blurb4 = gr.Markdown("Youzhang He is currently a second year graduate student at University of Pennsylvania majoring in Data Science. His research field is NLP, LLM and embodiedAI. He is interested in applying deep learning and machine learning techniques to solve real-world problems such as recommendation systems and webagents.")
        btn_back_about = gr.Button("", elem_classes=["back-button"])

    # PAGE 7: PROGRESS TRACKER (new page)
    with gr.Column(visible=False, elem_id="page_progress") as page_progress:
        progress_back_button = gr.Button("", elem_classes=["back-button"])
        with gr.Column(elem_id="page_progress"):
            gr.Markdown("<h2>Progress Tracker</h2>")
        # First row: Streak Tracker and Session Summary side by side.
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("<h3 class='section-header'><i class='fa fa-chart-line'></i> Daily Activity Chart</h3>")
                gr.Plot(dummy_chart)
            with gr.Column(scale=1):
                gr.Markdown("<h3 class='section-header'><i class='fa fa-fire'></i> Streak Tracker</h3>")
                gr.Markdown("**Days in a row used**")
                # Instead of an image, display text with flame emojis.
                gr.Markdown("<h4>5 days 🔥🔥🔥🔥🔥</h4>")
                gr.Markdown("<h3 class='section-header'><i class='fa fa-book'></i> Session Summary</h3>")
                gr.Markdown(""" 
                **Subjects Covered:**
                - Math
                - Science
                - History

                **Subtopics:**
                - Algebra, Geometry
                - Physics, Chemistry
                - Ancient, Modern

                **Turns:** 25
                """)
                gr.Markdown("<h3 class='section-header'><i class='fa fa-trophy'></i> Gamified Stats</h3>")
                gr.Markdown(""" 
                **XP:** 1500  
                **Badges Earned:** 5  
                **Current Level:** 10
                """)
        # Second row: Daily Activity Chart spanning full width.
        # Third row: Gamified Stats spanning full width.
                
    # CALLBACKS
    # -- Existing
    for button, action_name in [
        (btn_homework, "Homework"),
        (btn_explain, "Explain This Concept"),
        (btn_review, "Study Review"),
        (btn_quiz, "Practice Quiz"),
        (btn_challenge, "Challenge Mode"),
    ]:
        button.click(
            fn=pick_action,
            inputs=[user_session, gr.State(value=action_name)],
            outputs=[user_session, page_action, page_subject]
        )
    btn_freeplay.click(fn=pick_action, inputs=[user_session, gr.State("Freeplay")],
                     outputs=[user_session, page_action, page_chat])
    btn_mental.click(fn=pick_mental_health, inputs=[user_session],
                     outputs=[user_session, page_action, page_chat, chatbot])
    btn_about.click(fn=pick_action, inputs=[user_session, gr.State("About Us")],
                    outputs=[user_session, page_action, page_about])

    btn_cs.click(fn=pick_subject, inputs=[user_session, gr.State("Computer Science")],
                 outputs=[user_session, page_subject, page_subtopic, 
                          cs_col, hist_col, eng_col, science_col, nonsci_col])
    btn_hist.click(fn=pick_subject, inputs=[user_session, gr.State("History")],
                   outputs=[user_session, page_subject, page_subtopic, 
                            cs_col, hist_col, eng_col, science_col, nonsci_col])
    btn_eng.click(fn=pick_subject, inputs=[user_session, gr.State("English")],
                  outputs=[user_session, page_subject, page_subtopic, 
                           cs_col, hist_col, eng_col, science_col, nonsci_col])
    btn_sci.click(fn=pick_subject, inputs=[user_session, gr.State("Science")],
                  outputs=[user_session, page_subject, page_subtopic, 
                           cs_col, hist_col, eng_col, science_col, nonsci_col])

    other_subj_enter.click(
        fn=on_other_subject_skip_subtopic,
        inputs=[other_subj_input, user_session],
        outputs=[user_session, subject_msg, other_subj_input,
                 page_subject, page_subtopic, cs_col, hist_col, eng_col, science_col, nonsci_col, page_grade]
    )

    btn_back_home.click(fn=go_back_home, inputs=[], outputs=[page_action, page_subject])

    btn_sw.click(fn=pick_any_subtopic, inputs=[user_session, gr.State("Software")],
                 outputs=[user_session, page_subtopic, page_grade])
    btn_hw.click(fn=pick_any_subtopic, inputs=[user_session, gr.State("Hardware")],
                 outputs=[user_session, page_subtopic, page_grade])
    btn_leet.click(fn=pick_any_subtopic, inputs=[user_session, gr.State("Interview Questions/LeetCode")],
                   outputs=[user_session, page_subtopic, page_grade])
    btn_proj.click(fn=pick_any_subtopic, inputs=[user_session, gr.State("Passion Projects")],
                   outputs=[user_session, page_subtopic, page_grade])

    cs_other_enter.click(
        fn=on_subtopic_other_skip_to_grade,
        inputs=[cs_other_input, user_session],
        outputs=[user_session, cs_subtopic_msg, cs_other_input, page_subtopic, page_grade]
    )

    btn_am_hist.click(fn=pick_any_subtopic, inputs=[user_session, gr.State("American History")],
                      outputs=[user_session, page_subtopic, page_grade])
    btn_mod_hist.click(fn=pick_any_subtopic, inputs=[user_session, gr.State("Modern World History")],
                       outputs=[user_session, page_subtopic, page_grade])
    btn_anc_hist.click(fn=pick_any_subtopic, inputs=[user_session, gr.State("Ancient World History")],
                       outputs=[user_session, page_subtopic, page_grade])
    btn_gov.click(fn=pick_any_subtopic, inputs=[user_session, gr.State("Government and Politics")],
                  outputs=[user_session, page_subtopic, page_grade])

    hist_other_enter.click(
        fn=on_subtopic_other_skip_to_grade,
        inputs=[hist_other_input, user_session],
        outputs=[user_session, hist_subtopic_msg, hist_other_input, page_subtopic, page_grade]
    )

    btn_gram.click(fn=pick_any_subtopic, inputs=[user_session, gr.State("Grammar")],
                   outputs=[user_session, page_subtopic, page_grade])
    btn_read.click(fn=pick_any_subtopic, inputs=[user_session, gr.State("Reading")],
                   outputs=[user_session, page_subtopic, page_grade])
    btn_write.click(fn=pick_any_subtopic, inputs=[user_session, gr.State("Writing")],
                    outputs=[user_session, page_subtopic, page_grade])
    btn_app.click(fn=pick_any_subtopic, inputs=[user_session, gr.State("Application Essays")],
                  outputs=[user_session, page_subtopic, page_grade])

    eng_other_enter.click(
        fn=on_subtopic_other_skip_to_grade,
        inputs=[eng_other_input, user_session],
        outputs=[user_session, eng_subtopic_msg, eng_other_input, page_subtopic, page_grade]
    )

    btn_phys.click(fn=pick_any_subtopic, inputs=[user_session, gr.State("Physics")],
                   outputs=[user_session, page_subtopic, page_grade])
    btn_chem.click(fn=pick_any_subtopic, inputs=[user_session, gr.State("Chemistry")],
                   outputs=[user_session, page_subtopic, page_grade])
    btn_bio.click(fn=pick_any_subtopic, inputs=[user_session, gr.State("Biology")],
                  outputs=[user_session, page_subtopic, page_grade])
    btn_math_sci.click(fn=pick_any_subtopic, inputs=[user_session, gr.State("Math")],
                       outputs=[user_session, page_subtopic, page_grade])

    sc_other_enter.click(
        fn=on_subtopic_other_skip_to_grade,
        inputs=[sc_other_input, user_session],
        outputs=[user_session, subtopic_msg, sc_other_input, page_subtopic, page_grade]
    )

    placeholder_btn.click(fn=pick_non_science_subtopic,
                          inputs=[user_session],
                          outputs=[user_session, page_subtopic, page_grade])

    btn_back_subject.click(fn=go_back_subject, inputs=[], outputs=[page_subject, page_subtopic])

    btn_elem.click(fn=pick_grade_level, inputs=[user_session, gr.State("Elementary (K-5)")],
                   outputs=[user_session, page_grade, page_chat, chatbot])
    btn_mid.click(fn=pick_grade_level, inputs=[user_session, gr.State("Middle (6-8)")],
                  outputs=[user_session, page_grade, page_chat, chatbot])
    btn_high.click(fn=pick_grade_level, inputs=[user_session, gr.State("High (9-12)")],
                   outputs=[user_session, page_grade, page_chat, chatbot])
    btn_uni.click(fn=pick_grade_level, inputs=[user_session, gr.State("University")],
                  outputs=[user_session, page_grade, page_chat, chatbot])

    btn_back_subtopic.click(fn=go_home_from_chat, inputs=[], outputs=[page_subtopic, page_grade])

    # ----------------------------
    # New: Chat Response with Mental Health Check-In
    # ----------------------------
    def chat_response(user_session, message, history, generate_img):
        # Process the chatbot response normally
        for update in chatbot_respond(message, history, user_session, generate_img):
            yield update
        # After processing the response, check if the number of user messages reaches the threshold
        user_message_count = sum(1 for msg in history if msg[0] is not None)
        if user_message_count > 0 and user_message_count % MENTAL_HEALTH_THRESHOLD == 0:
            # Only add a mental health check-in if the session is not already in Mental Health mode
            if user_session.get("action", "").lower() != "mental health":
                mental_message = "You're doing an amazing job! Consider taking a short break before continuing to study."
                history.append([None, mental_message])
                yield history, "", None

    chat_input.submit(
        fn=chat_response,
        inputs=[user_session, chat_input, chatbot, generate_image_toggle],
        outputs=[chatbot, chat_input, image_output]
    )

    clear_chat_button.click(
        fn=clear_chat_history,
        inputs=[],
        outputs=[chatbot, chat_input, image_output]
    )

    btn_new_task.click(
        fn=suggest_task,
        inputs=[user_session, chatbot],
        outputs=[chatbot, chat_input]
    )
    
    # New callback for Fun Fact button
    btn_fun_fact.click(
        fn=fun_fact,
        inputs=[user_session, chatbot],
        outputs=[chatbot, chat_input]
    )

    btn_back_grade.click(fn=go_back_grade, inputs=[], outputs=[page_action, page_chat])

    btn_back_about.click(fn=go_back_home, inputs=[], outputs=[page_action, page_about])

    # -- New callbacks for Progress Tracker
    progress_button_home.click(
        fn=go_to_progress_from_home,
        inputs=[],
        outputs=[progress_return_page, page_action, page_chat, page_progress]
    )
    progress_button_chat.click(
        fn=go_to_progress_from_chat,
        inputs=[],
        outputs=[progress_return_page, page_action, page_chat, page_progress]
    )
    progress_back_button.click(
        fn=go_back_progress,
        inputs=[progress_return_page],
        outputs=[page_action, page_chat, page_progress, progress_return_page]
    )

demo.launch()