app.py

import json
import os
import copy
import streamlit as st
from openai import OpenAI
from datetime import datetime
from langchain_openai import ChatOpenAI
from typing import Annotated, List
from pydantic import BaseModel, Field
from typing_extensions import TypedDict, Literal
from langgraph.graph import StateGraph, START, END

# Page configuration
st.set_page_config(layout="wide", page_title="JEE Roadmap Planner")

# Initialize session state variables
if "data" not in st.session_state:
    st.session_state.data = None
if "data_old" not in st.session_state:
    st.session_state.data_old = None
if "incomplete_tasks" not in st.session_state:
    st.session_state.incomplete_tasks = None
if "incomplete_task_list" not in st.session_state:
    st.session_state.incomplete_task_list = None
if "final_report" not in st.session_state:
    st.session_state.final_report = None
if "shifted_roadmap" not in st.session_state:
    st.session_state.shifted_roadmap = None
if "available_dates" not in st.session_state:
    st.session_state.available_dates = []
if "updated_roadmap" not in st.session_state:
    st.session_state.updated_roadmap = None
if "max_optimizer_iterations" not in st.session_state:
    st.session_state.max_optimizer_iterations = 3  # Limit optimizer to 3 iterations

# Navigation sidebar setup
st.sidebar.title("JEE Roadmap Planner")
page = st.sidebar.radio("Navigation", ["Home", "Roadmap Manager", "Task Analysis","Roadmap Chatbot"])


# For roadmap chatbot
import sqlite3

# Function to convert NL query to SQL
def generate_sql_from_nl(prompt):
    client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

    table_struct = """
        CREATE TABLE IF NOT EXISTS roadmap (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            day_num INTEGER,
            date TEXT,
            subject TEXT,
            chapter_name TEXT,
            task_type TEXT,
            time TEXT,
            subtopic TEXT
        )
    """

    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": f""""You are an helper who runs in the background of an AI agent, 
             which helps students for their JEE Preparation. Now your Job is to analyze the users prompt and 
             create an SQL query to extract the related Information from an sqlite3 database with the table 
             structure: {table_struct}. 

             Note: For the time column, the data is formatted like '0.5 hour', '1 hour', '2 hours' and 
             so on. So make sure create queries that compare just the numbers within the text.

             You will also make sure multiple times that you give an SQL
             Query that adheres to the given table structure, and you Output just the SQL query.
             Do not include anyting else like new line statements, ```sql or any other text. Your output
             is going to be directly fed into a Python script to extract the required information. So,
             please follow all the given Instructions."""},
            {"role": "user", "content": f"""Keeping the table structure in mind: {table_struct}, 
             Convert this prompt to an SQL query for the given table: {prompt}. Make sure your
             output is just the SQL query, which can directly be used to extract required content"""}
        ]
    )
    return response.choices[0].message.content.strip()

# Function to convert SQL output to natural language
def generate_nl_from_sql_output(prompt, data):
    client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": f"""You are an helpful AI chatbot working under the roadmap 
             section of an AI Agent, whose role is to aid students in their preparation for the JEE examination. 
             You are going to play a very crucial role of a Roadmap Assistant, who helps the student out with whatever query
             they have related to their roadmap, the data required to answer the users query is already extracted 
             from the Roadmap table of a SQLite3 database and given to you here {data}. Analyse the users query deeply and 
             reply to it with the relevant information from the given data in a supportive manner."""},
            {"role": "user", "content": f"""Answer to this users query using the data given to you, while keeping 
             your role in mind: {prompt}"""}
        ]
    )
    return response.choices[0].message.content.strip()

# Function to fetch data from SQLite
def fetch_data_from_sql(sql_query):
    conn = sqlite3.connect("jee_roadmap.db")
    cursor = conn.cursor()
    cursor.execute(sql_query)
    rows = cursor.fetchall()
    conn.close()
    return rows

# Main function for chatbot
def answer_user_query(prompt):
    initialize_roadmap_db()
    sql = generate_sql_from_nl(prompt)
    st.write(sql)
    rows = fetch_data_from_sql(sql)
    st.write(rows)
    return generate_nl_from_sql_output(prompt, rows)
    
def initialize_roadmap_db():
    if not os.path.exists("jee_roadmap.db"):
        try:
            with open("full_roadmap.json") as f:
                roadmap_data = json.load(f)

            conn = sqlite3.connect("jee_roadmap.db")
            cursor = conn.cursor()

            cursor.execute("""
            CREATE TABLE IF NOT EXISTS roadmap (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                day_num INTEGER,
                date TEXT,
                subject TEXT,
                chapter_name TEXT,
                task_type TEXT,
                time TEXT,
                subtopic TEXT
            )
            """)

            for day in roadmap_data["schedule"]:
                date = day["date"]
                day_num = day["dayNumber"]
                for subj in day["subjects"]:
                    subject = subj["name"]
                    for task in subj["tasks"]:
                        cursor.execute("""
                            INSERT INTO roadmap (day_num, date, subject, chapter_name, task_type, time, subtopic)
                            VALUES (?, ?, ?, ?, ?, ?, ?)
                        """, (
                            day_num,
                            date,
                            subject,
                            task["ChapterName"],
                            task["type"],
                            task["time"],
                            task["subtopic"]
                        ))

            conn.commit()
            conn.close()
            print("✅ Database created and data inserted successfully.")
        except Exception as e:
            print(f"⚠️ Error initializing database: {e}")
# Function to load initial data
def load_initial_data():
    with st.spinner("Loading roadmap data..."):
        try:
            with open('fourdayRoadmap.json', 'r') as file:
                data = json.load(file)
                st.session_state.data = data
                st.session_state.data_old = copy.deepcopy(data)
                st.success("Data loaded successfully!")
                return True
        except Exception as e:
            st.error(f"Error loading data: {e}")
            return False

# Function to mark tasks as incomplete
def process_task_completion_data():
    with st.spinner("Processing task completion data..."):
        data = st.session_state.data
        for day in data["schedule"]:
            for subject in day["subjects"]:
                for task in subject["tasks"]:
                    task["task_completed"] = False
                    task["completion_timestamp"] = None
        st.session_state.data = data
        st.success("Task completion data processed!")

# Function to extract incomplete tasks
def extract_incomplete_tasks():
    with st.spinner("Extracting incomplete tasks..."):
        data = st.session_state.data
        previous_day = data["schedule"][0]
        
        incomplete_tasks = {
            "dayNumber": previous_day["dayNumber"],
            "date": previous_day["date"],
            "subjects": []
        }

        for subject in previous_day["subjects"]:
            incomplete_subject_tasks = [
                {
                    "ChapterName": task["ChapterName"],
                    "type": task["type"],
                    "subtopic": task["subtopic"],
                    "time": task["time"],
                    "task_completed": False,
                    "completion_timestamp": None
                }
                for task in subject["tasks"] if not task["task_completed"]
            ]

            if incomplete_subject_tasks:
                incomplete_tasks["subjects"].append({
                    "name": subject["name"],
                    "tasks": incomplete_subject_tasks
                })

        # Convert to JSON format
        incomplete_tasks_json = json.dumps(incomplete_tasks, indent=4)
        st.session_state.incomplete_tasks = incomplete_tasks
        
        # Generate a list of incomplete tasks for the agent
        incomplete_task_list = []
        for subject in incomplete_tasks["subjects"]:
            for task in subject["tasks"]:
                if not task["task_completed"]:
                    incomplete_task = {
                        "subject": subject["name"],
                        "ChapterName": task["ChapterName"],
                        "type": task["type"],
                        "subtopic": task["subtopic"],
                        "time": task["time"],
                        "task_completed": task["task_completed"],
                        "completion_timestamp": task["completion_timestamp"]
                    }
                    incomplete_task_list.append(incomplete_task)
        
        st.session_state.incomplete_task_list = incomplete_task_list
        st.success("Incomplete tasks extracted!")

# Function to generate report
def generate_report():
    with st.spinner("Generating performance report using AI..."):
        previous_day = st.session_state.data["schedule"][0]
        previous_day_roadmap_str = str(previous_day)
        
        try:
            client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
            response = client.chat.completions.create(
                model="gpt-4o-mini",
                messages=[
                    {"role": "system", "content": """You will be given a JEE student's previous_day_roadmap and then you have to create
                    a completely interactive and useful report for the user. 
                    The report should include a table for task completion rates and data
                    The student's study pattern
                    The student's weaknesses and tips to improve.
                    Make sure that a task is completed only when the "task_completed" key is true and the "time" key tells about how much
                    tentative that task can take time
                    Use markdown formatting.
                    """},
                    {"role": "user", "content": f"""Here is the user's previous day roadmap in json : {previous_day_roadmap_str}"""}
                ]
            )
            output = response.choices[0].message.content
            st.session_state.final_report = output
            st.success("Report generated successfully!")
        except Exception as e:
            st.error(f"Error generating report: {e}")

# Function to extract available dates
def extract_available_dates():
    with st.spinner("Extracting available dates for rescheduling..."):
        data = st.session_state.data
        
        def remove_the_first_day(roadmap):
            new_roadmap = {
                "schedule": []
            }
            for day in roadmap['schedule']:
                if day['dayNumber'] != 1:
                    new_roadmap['schedule'].append(day)
            return new_roadmap
        
        roadmap = remove_the_first_day(data)
        
        available_dates = []
        for day in roadmap['schedule']:
            available_dates.append(day['date'])
        
        st.session_state.available_dates = available_dates
        st.success(f"Found {len(available_dates)} available dates for rescheduling!")

# Function to shift incomplete tasks using the evaluator-optimizer approach
def shift_incomplete_tasks():
    with st.spinner("Optimizing task distribution using evaluator-optimizer approach..."):
        try:
            # Initialize needed components
            llm = ChatOpenAI(model="gpt-4o-mini")
            
            # Schema for structured output to use in evaluation
            class Feedback(BaseModel):
                grade: Literal["added", "not added"] = Field(
                    description="Check if all the incomplete tasks are added to the roadmap or not",
                )
                feedback: str = Field(
                    description="If some tasks are not added, give feedback to add those tasks also",
                )
            
            # Augment the LLM with schema for structured output
            evaluator = llm.with_structured_output(Feedback)
            
            # Graph state
            class State(TypedDict):
                roadmap: dict
                available_dates: list
                incomplete_task_list: list
                feedback: str
                added_or_not: str
                iteration_count: int
            
            # Initialize state
            current_state = {
                "roadmap": {},
                "available_dates": st.session_state.available_dates,
                "incomplete_task_list": st.session_state.incomplete_task_list,
                "feedback": "",
                "added_or_not": "",
                "iteration_count": 0
            }
            
            # Progress bar for iterations
            progress_bar = st.progress(0)
            iteration_status = st.empty()
            
            # First call to generator
            iteration_status.write("Iteration 1: Generating initial task distribution...")
            if current_state.get("feedback"):
                msg = llm.invoke(
                    f"""Add the following incomplete_tasks {current_state['incomplete_task_list']} to the roadmap key and take into account the feedback {current_state['feedback']}
                    and make sure that we dynamically add the tasks not increasing the load on just one day, also we have to add the tasks on
                    following dates only {current_state['available_dates']} Make sure you only give roadmap json as output and nothing else, strictly follow the output structure:
                    ```json
                    {{
                      "roadmap": [
                        {{
                          "date": "YYYY-MM-DD",
                          "tasks": [
                            {{
                              "subject": "Subject Name",
                              "ChapterName": "Chapter Name",
                              "type": "Type of Task",
                              "subtopic": "Subtopic Name",
                              "time": "estimated time"
                            }}
                          ]
                        }}
                      ]
                    }}
                    ```
                    """
                )
            else:
                msg = llm.invoke(
                    f"""Add the following incomplete_tasks {current_state['incomplete_task_list']} to the roadmap key
                    and make sure that we dynamically add the tasks not increasing the load on just one day, also we have to add the tasks on
                    following dates only {current_state['available_dates']} Make sure you only give roadmap json as output and nothing else, strictly follow the output structure:
                    ```json
                    {{
                      "roadmap": [
                        {{
                          "date": "YYYY-MM-DD",
                          "tasks": [
                            {{
                              "subject": "Subject Name",
                              "ChapterName": "Chapter Name",
                              "type": "Type of Task",
                              "subtopic": "Subtopic Name",
                              "time": "estimated time"
                            }}
                          ]
                        }}
                      ]
                    }}
                    ```
                    """
                )
            
            current_state["roadmap"] = msg.content
            progress_bar.progress(1/6)
            
            # Enter optimization loop
            max_iterations = st.session_state.max_optimizer_iterations
            current_iteration = 1
            
            while current_iteration <= max_iterations:
                # Evaluate the current roadmap
                iteration_status.write(f"Iteration {current_iteration}: Evaluating task distribution...")
                
                grade = evaluator.invoke(
                    f"Grade the roadmap {current_state['roadmap']} by checking if {current_state['incomplete_task_list']} are all added or not and make sure that we dynamically add the tasks not increasing the load on just one day"
                )
                
                current_state["added_or_not"] = grade.grade
                current_state["feedback"] = grade.feedback
                current_state["iteration_count"] += 1
                
                progress_bar.progress((current_iteration * 2 - 1)/6)
                
                # Check if we're done or need another iteration
                if current_state["added_or_not"] == "added":
                    iteration_status.write(f"✅ Success! All tasks distributed after {current_iteration} iterations.")
                    break
                
                if current_iteration == max_iterations:
                    iteration_status.write(f"⚠️ Reached maximum iterations ({max_iterations}). Using best result so far.")
                    break
                
                # Generate an improved roadmap based on feedback
                iteration_status.write(f"Iteration {current_iteration + 1}: Improving task distribution based on feedback...")
                
                msg = llm.invoke(
                    f"""Add the following incomplete_tasks {current_state['incomplete_task_list']} to the roadmap key 
                    and take into account the feedback: {current_state['feedback']}
                    Make sure that we dynamically add the tasks not increasing the load on just one day.
                    Only add tasks on these available dates: {current_state['available_dates']}
                    Make sure you only give roadmap json as output and nothing else, strictly follow the output structure:
                    ```json
                    {{
                      "roadmap": [
                        {{
                          "date": "YYYY-MM-DD",
                          "tasks": [
                            {{
                              "subject": "Subject Name",
                              "ChapterName": "Chapter Name",
                              "type": "Type of Task",
                              "subtopic": "Subtopic Name",
                              "time": "estimated time"
                            }}
                          ]
                        }}
                      ]
                    }}
                    ```
                    """
                )
                
                current_state["roadmap"] = msg.content
                current_iteration += 1
                progress_bar.progress((current_iteration * 2 - 2)/6)
            
            # Process the final roadmap content
            shifted_tasks_roadmap = current_state["roadmap"]
            
            # Extract JSON part from the response
            if "```json" in shifted_tasks_roadmap:
                shifted_tasks_roadmap = shifted_tasks_roadmap.split("```json")[1].split("```")[0]
            elif "```" in shifted_tasks_roadmap:
                shifted_tasks_roadmap = shifted_tasks_roadmap.split("```")[1].split("```")[0]
            
            st.session_state.shifted_roadmap = shifted_tasks_roadmap
            st.success(f"Tasks rescheduled successfully after {current_iteration} iterations!")
            
            progress_bar.progress(1.0)
            
        except Exception as e:
            st.error(f"Error in optimization process: {e}")

# Function to merge shifted tasks into main roadmap
def merge_shifted_tasks():
    with st.spinner("Merging rescheduled tasks into main roadmap..."):
        try:
            data = st.session_state.data
            shifted_roadmap = json.loads(st.session_state.shifted_roadmap)
            
            def add_task(roadmap, task, date_task_to_be_added):
                subject_name = task["subject"]
                chapter_name = task["ChapterName"]
                topic_name = task["subtopic"]
                type_name = task["type"]

                # Check if the date exists
                for day in roadmap['schedule']:
                    if day['date'] == date_task_to_be_added:
                        # Find or create subject
                        subject_exists = False
                        for subject in day['subjects']:
                            if subject['name'] == subject_name:
                                subject_exists = True
                                # Check if task already exists
                                task_exists = False
                                for existing_task in subject['tasks']:
                                    if (existing_task.get('ChapterName') == chapter_name and 
                                        existing_task.get('type') == type_name and 
                                        (existing_task.get('subtopic', '') == topic_name or 
                                         existing_task.get('topic', '') == topic_name)):
                                        task_exists = True
                                        break
                                
                                if not task_exists:
                                    # Add task
                                    temp_task = {
                                        "ChapterName": chapter_name,
                                        "type": type_name,
                                        "subtopic": topic_name,
                                        "time": task["time"],
                                        "task_completed": False,
                                        "completion_timestamp": None
                                    }
                                    subject['tasks'].append(temp_task)
                                break
                        
                        # If subject doesn't exist, create it
                        if not subject_exists:
                            new_subject = {
                                "name": subject_name,
                                "tasks": [{
                                    "ChapterName": chapter_name,
                                    "type": type_name,
                                    "subtopic": topic_name,
                                    "time": task["time"],
                                    "task_completed": False,
                                    "completion_timestamp": None
                                }]
                            }
                            day['subjects'].append(new_subject)
                        break
                        
                return roadmap

            # Process each task in the shifted roadmap
            for day in shifted_roadmap["roadmap"]:
                date = day['date']
                tasks = day['tasks']
                for task in tasks:
                    data = add_task(data, task, date)
            
            st.session_state.updated_roadmap = data
            st.success("Tasks merged into roadmap successfully!")
            
        except Exception as e:
            st.error(f"Error merging tasks: {e}")

# ---- HOME PAGE ----
if page == "Home":
    st.title("📚 JEE Roadmap Planner")
    
    st.markdown("""
    ### Welcome to your JEE Study Roadmap Planner!
    
    This tool helps you manage your JEE preparation schedule by:
    
    1. 📊 **Analyzing your study performance**
    2. 🔄 **Redistributing incomplete tasks**
    3. 📝 **Providing personalized feedback**
    
    Get started by loading your roadmap data and following the step-by-step process.
    """)
    
    # Settings section
    with st.expander("⚙️ Advanced Settings"):
        st.session_state.max_optimizer_iterations = st.slider(
            "Maximum Optimizer Iterations", 
            min_value=1, 
            max_value=5, 
            value=st.session_state.max_optimizer_iterations,
            help="Limit how many times the optimizer will try to improve task distribution"
        )
    
    st.info("Navigate using the sidebar to access different features of the app.")
    
    # Initial data loading
    if st.button("📂 Load Roadmap Data"):
        success = load_initial_data()
        if success:
            st.session_state.first_load = True

# ---- ROADMAP MANAGER PAGE ----
elif page == "Roadmap Manager":
    st.title("🗓️ Roadmap Manager")
    
    if st.session_state.data is None:
        st.warning("Please load roadmap data first from the Home page.")
    else:
        st.markdown("### Roadmap Management Steps")
        
        col1, col2 = st.columns(2)
        
        with col1:
            st.subheader("Step 1: Process Tasks")
            if st.button("1️⃣ Mark Tasks as Incomplete"):
                process_task_completion_data()
            
            st.subheader("Step 2: Extract Tasks")
            if st.button("2️⃣ Extract Incomplete Tasks"):
                extract_incomplete_tasks()
                
            if st.session_state.incomplete_task_list:
                st.write(f"Found {len(st.session_state.incomplete_task_list)} incomplete tasks")
                with st.expander("View Incomplete Tasks"):
                    st.json(st.session_state.incomplete_task_list)
        
        with col2:
            st.subheader("Step 3: Prepare for Rescheduling")
            if st.button("3️⃣ Extract Available Dates"):
                extract_available_dates()
                
            if st.session_state.available_dates:
                with st.expander("View Available Dates"):
                    st.write(st.session_state.available_dates)
                    
            st.subheader("Step 4: Reschedule Tasks")
            if st.button("4️⃣ Optimize Task Distribution"):
                if not st.session_state.incomplete_task_list or not st.session_state.available_dates:
                    st.error("Please complete steps 2 and 3 first!")
                else:
                    shift_incomplete_tasks()
            
            if st.session_state.shifted_roadmap:
                with st.expander("View Task Distribution Plan"):
                    try:
                        st.json(json.loads(st.session_state.shifted_roadmap))
                    except:
                        st.text(st.session_state.shifted_roadmap)
        
        st.subheader("Step 5: Update Roadmap")
        if st.button("5️⃣ Merge Tasks into Main Roadmap"):
            if not st.session_state.shifted_roadmap:
                st.error("Please complete step 4 first!")
            else:
                merge_shifted_tasks()
                
        # Display original and updated roadmaps side by side
        if st.session_state.data_old and st.session_state.updated_roadmap:
            st.subheader("Roadmap Comparison")
            col1, col2 = st.columns(2)
            
            with col1:
                st.markdown("#### Original Roadmap")
                with st.expander("View Original Roadmap"):
                    st.json(st.session_state.data_old)
            
            with col2:
                st.markdown("#### Updated Roadmap")
                with st.expander("View Updated Roadmap"):
                    st.json(st.session_state.updated_roadmap)

# ---- TASK ANALYSIS PAGE ----
elif page == "Task Analysis":
    st.title("📊 Task Analysis")
    
    if st.session_state.data is None:
        st.warning("Please load roadmap data first from the Home page.")
    else:
        st.markdown("### Performance Report")
        
        if st.button("🔍 Generate Performance Report"):
            generate_report()
        
        if st.session_state.final_report:
            st.markdown(st.session_state.final_report)
        else:
            st.info("Click the button above to generate your performance report.")
        
        # Add visualization options
        if st.session_state.data:
            st.subheader("Task Breakdown")
            
            # Simple task statistics
            if st.checkbox("Show Task Statistics"):
                task_count = 0
                subject_counts = {}
                type_counts = {}
                
                for day in st.session_state.data["schedule"]:
                    for subject in day["subjects"]:
                        subject_name = subject["name"]
                        if subject_name not in subject_counts:
                            subject_counts[subject_name] = 0
                        
                        for task in subject["tasks"]:
                            subject_counts[subject_name] += 1
                            task_count += 1
                            
                            # Count by task type
                            task_type = task.get("type", "Unknown")
                            if task_type not in type_counts:
                                type_counts[task_type] = 0
                            type_counts[task_type] += 1
                
                st.write(f"Total tasks: {task_count}")
                
                # Create charts for data visualization
                col1, col2 = st.columns(2)
                
                with col1:
                    st.subheader("Subject Distribution")
                    st.bar_chart(subject_counts)
                
                with col2:
                    st.subheader("Task Type Distribution")
                    st.bar_chart(type_counts)
# ---- ROADMAP CHATBOT PAGE ----
elif page == "Roadmap Chatbot":
    st.title("🤖 Roadmap Chatbot Assistant")

    user_query = st.text_input("Ask a question about your roadmap:", placeholder="e.g., What are my tasks on 14 Feb 2025?")

    if st.button("Ask") and user_query:
        with st.spinner("Thinking..."):
            try:
                response = answer_user_query(user_query)
                st.markdown(response)
            except Exception as e:
                st.error(f"Error: {e}")