Update app.py

app.py

@@ -4,12 +4,14 @@ import copy
 import sqlite3
 import operator
 import streamlit as st
-from openai import OpenAI
+from math import ceil
+from datetime import datetime, timedelta
+from collections import defaultdict
 from langchain_openai import ChatOpenAI
 from langchain_core.messages import HumanMessage, SystemMessage
 from typing import Annotated, List
 from pydantic import BaseModel, Field
-from typing_extensions import TypedDict, Literal
+from typing_extensions import TypedDict
 from langgraph.graph import StateGraph, START, END
 from langgraph.constants import Send
 
@@ -19,8 +21,8 @@ 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 "full_roadmap" not in st.session_state:
+    st.session_state.full_roadmap = None
 if "report_data" not in st.session_state:
     st.session_state.report_data = None
 if "incomplete_tasks" not in st.session_state:
@@ -35,180 +37,23 @@ 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
-def get_chapters_and_subtopics():
-    with open("full_roadmap.json", "r") as f:
-        data = json.load(f)
-
-    ch_subt = {
-        "Physics": {},
-        "Chemistry": {},
-        "Maths": {}
-    }
-
-    for day in data["schedule"]:
-        for subject in day['subjects']:
-            sub = ch_subt[subject['name']]
-            for task in subject['tasks']:
-                sub[task['ChapterName']] = []
-
-    for day in data["schedule"]:
-        for subject in day['subjects']:
-            sub = ch_subt[subject['name']]
-            for task in subject['tasks']:
-                if task['subtopic'] not in sub[task['ChapterName']]:
-                  sub[task['ChapterName']].append(task['subtopic'])
-
-    return ch_subt
-
-# 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, -- [yyyy-mm-dd]
-                        subject TEXT, -- [Physics, Chemistry or Maths]
-                        chapter_name TEXT,
-                        task_type TEXT, -- (Concept Understanding, Question Practice, Revision, Test)
-                        time TEXT, -- formatted like '0.5 hour', '1 hour', '2 Hours', and so on
-                        subtopic TEXT,
-                    )"""
-    
-    ch_subt = get_chapters_and_subtopics()
-
-    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 to create queries that compare just the numbers within the text.
-             - If the student mention about any chapters or subtopics, browse through this json file {ch_subt},
-               find the one with the closest match to the users query and use only those exact names of Chapers 
-               and Subtopics present in this file to create SQL the query.
-             
-             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, query, 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. If you get empty data
-             as an input, deeply analyze the user's prompt and the sql query and give a suitable reply."""},
-            {"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()
-    query = generate_sql_from_nl(prompt)
-    st.write(query)
-    rows = fetch_data_from_sql(query)
-    st.write(rows)
-    return generate_nl_from_sql_output(prompt, query, 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
+# AGENT 1
 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
+            with open("full_roadmap.json", 'r') as file:
+                data = json.load(file)
+                st.session_state.full_roadmap = data
+            st.success("Data loaded successfully!")
+            return True
         except Exception as e:
             st.error(f"Error loading data: {e}")
             return False
@@ -225,60 +70,204 @@ def process_task_completion_data():
         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": []
-        }
+def check_tot_time(day, max_hours_per_day):
+  tot_time = 0
+  for subject in day:
+    for task in subject["tasks"]:
+      tot_time += float(task['time'].split(" ")[0])
+
+  if tot_time > max_hours_per_day:
+    return tot_time, True
+  else:
+    return tot_time, False
+
+def shift_roadmap(roadmap, max_hours_per_day):
+    roadmap = copy.deepcopy(roadmap)
+    incomplete_tasks_by_subject = defaultdict(list)
+
+    prev_day = roadmap[0]
+    for subject in prev_day["subjects"]:
+        subject_name = subject["name"]
+        tasks = subject["tasks"]
+
+        # Separate completed and incomplete tasks
+        incomplete_tasks = [task for task in tasks if task['task_completed'] == False]
+        completed_tasks = [task for task in tasks if task['task_completed'] == True]
+
+        # Store incomplete tasks per subject
+        if incomplete_tasks:
+            incomplete_tasks_by_subject[subject_name].extend(incomplete_tasks)
+
+        # Keep only completed tasks in the previous day
+        subject["tasks"] = completed_tasks
+
+
+    # Step 2: Redistribute tasks across the next 3 days
+    for i, next_day in enumerate(roadmap[1:]):  # Next 3 days (Day 2, Day 3, Day 4)
+        for subject in next_day["subjects"]:
+            subject_name = subject["name"]
+            if subject_name in incomplete_tasks_by_subject and incomplete_tasks_by_subject[subject_name]:
+                total_tasks = len(incomplete_tasks_by_subject[subject_name])
+
+                # Task distribution based on 1/6, 2/6, and remaining
+                if i == 0:  # First day gets 1/6 of total
+                    tasks_to_add = ceil(total_tasks * (1 / 6))
+                elif i == 1:  # Second day gets 2/6 of total
+                    tasks_to_add = ceil(total_tasks * (2 / 6))
+                else:  # Remaining tasks on the last day
+                    tasks_to_add = len(incomplete_tasks_by_subject[subject_name])
+
+                # Append tasks to the current day's subject
+                subject["tasks"].extend(incomplete_tasks_by_subject[subject_name][:tasks_to_add])
+
+                # Remove assigned tasks from backlog
+                incomplete_tasks_by_subject[subject_name] = incomplete_tasks_by_subject[subject_name][tasks_to_add:]
+
+        # Make sure the time limit doesn't exceed for any day
+        _, check_time = check_tot_time(next_day["subjects"], max_hours_per_day)
+        while check_time:
+            for subject in next_day["subjects"]:
+                subject_name = subject["name"]
+                if subject["tasks"]:
+                    task_to_add = subject["tasks"].pop()
+                    incomplete_tasks_by_subject[subject_name].append(task_to_add)
+                    _, check_time = check_tot_time(next_day["subjects"], max_hours_per_day)
+                    if not check_time:
+                        break
 
-        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!")
+    return roadmap, incomplete_tasks_by_subject
+
+def get_shifted_roadmap(roadmap, dayNumber, max_hours_per_day):
+  day_index = dayNumber-1
+  if day_index+4 <= len(roadmap['schedule']):
+    shifted_roadmap, incomplete_tasks_by_subject = shift_roadmap(roadmap['schedule'][day_index:day_index+4], max_hours_per_day)
+  else:
+    shifted_roadmap, incomplete_tasks_by_subject = shift_roadmap(roadmap['schedule'][day_index:], max_hours_per_day)
+  for day in shifted_roadmap:
+      new_date = day["date"]
+
+      for idx, existing_day in enumerate(roadmap['schedule']):
+          if existing_day['date'] == new_date:
+              roadmap['schedule'][idx] = day
+              break
+
+  if any(len(v) != 0 for v in incomplete_tasks_by_subject.values()):
+    next_date = (datetime.strptime(roadmap['schedule'][-1]['date'], "%Y-%m-%d") + timedelta(days=1)).strftime("%Y-%m-%d")
+    next_day = roadmap['schedule'][-1]['dayNumber'] + 1
+    subjects = [{"name": subject_name, "tasks": tasks} for subject_name, tasks in incomplete_tasks_by_subject.items()]
+    roadmap['schedule'].append({
+        "dayNumber": next_day,
+        "date": next_date,
+        "subjects": subjects
+        })
+  return roadmap
+
+# Step 0: Get Subjectwise Chapter and Topic order ready
+def get_subjectwise_tasks(roadmap):
+    sub_tasks = {
+        "Physics": defaultdict(list),
+        "Chemistry": defaultdict(list),
+        "Maths": defaultdict(list)
+    }
+
+    for day in roadmap["schedule"]:
+        for subject in day['subjects']:
+            sub = sub_tasks[subject['name']]
+            for task in subject['tasks']:
+                if task['subtopic'] not in sub[task['ChapterName']]:
+                  sub[task['ChapterName']].append(task['subtopic'])
+
+    return sub_tasks
+
+# Step 1: Extract all tasks per subject and track time allocation per day
+def extract_tasks(roadmap):
+    subjectwise_tasks = defaultdict(list)
+    daily_subjectwise_time_allocation = defaultdict(lambda: defaultdict(float))
+
+    for day_index, day in enumerate(roadmap['schedule']):
+        for subject in day["subjects"]:
+            subject_name = subject["name"]
+            total_time = sum(float(task['time'].split(" ")[0]) for task in subject["tasks"])
+
+            daily_subjectwise_time_allocation[day_index][subject_name] = total_time
+            subjectwise_tasks[subject_name].extend(subject["tasks"])
+
+    return subjectwise_tasks, daily_subjectwise_time_allocation
+
+# Step 2: Sort all tasks for each subject
+def sort_tasks(tasks, reference):
+    task_type_priority = {
+      "Concept Understanding": 0,
+      "Question Practice": 1,
+      "Revision": 2,
+      "Test": 3
+    }
+    chapter_order = list(reference.keys())
+
+    def task_sort_key(task):
+        chapter = task["ChapterName"]
+        subtopic = task["subtopic"]
+        type_priority = task_type_priority.get(task["type"], 99)
+
+        chapter_idx = chapter_order.index(chapter) if chapter in reference else float('inf')
+        subtopic_idx = reference[chapter].index(subtopic) if subtopic in reference.get(chapter, []) else float('inf')
+
+        return (chapter_idx, subtopic_idx, type_priority)
+
+    return sorted(tasks, key=task_sort_key)
+
+# Helper function to get task time in hours
+def get_task_time(task):
+    return float(task['time'].split(" ")[0])
+
+# Step 3: Sort the roadmap by arranging the sorted tasks, preserving original time allocation
+def shift_and_sort_the_roadmap(full_roadmap, roadmap, dayNumber, max_hours_per_day):
+    roadmap = copy.deepcopy(roadmap)
+    roadmap = get_shifted_roadmap(roadmap, dayNumber, max_hours_per_day)
+    subject_refs = get_subjectwise_tasks(full_roadmap) # Load the full roadmap to obtain the correct chapter orders
+    subject_all_tasks, subject_day_time_allocation = extract_tasks(roadmap)
+
+    # Sort all tasks for each subject
+    for subject in subject_all_tasks:
+        subject_all_tasks[subject] = sort_tasks(subject_all_tasks[subject], subject_refs[subject])
+
+    # Redistribute tasks based on time allocation, strictly maintaining sequence
+    for day_index, day in enumerate(roadmap['schedule']):
+        day_time = 0
+        for subject in day["subjects"]:
+            subject_name = subject["name"]
+            target_time = subject_day_time_allocation[day_index][subject_name]
+
+            selected_tasks = []
+            current_time = 0
+            tasks = subject_all_tasks[subject_name]
+            while tasks and current_time < target_time:
+                next_task = tasks[0]  # Take the next task from the sorted sequence
+                task_time = get_task_time(next_task)
+
+                if day_time + task_time <= max_hours_per_day:  # Allow if its under the max limit
+                    selected_tasks.append(tasks.pop(0))
+                    current_time += task_time
+                    day_time += task_time
+                else:
+                    # If Task doesn't fit, save for next day
+                    break
+
+            if day_index == len(roadmap['schedule']) - 1:
+                if tasks:
+                    for task in tasks:
+                        selected_tasks.append(task)
+
+            # Update the subject's tasks
+            subject["tasks"] = selected_tasks
+            subject_all_tasks[subject_name] = tasks
 
+    with open("current_roadmap.json", "w") as f:
+        json.dump(roadmap, f, indent=4)
+
+    st.session_state.updated_roadmap = roadmap
+
+# AGENT 2
 def generate_sql_for_report(llm, prompt):
     table_struct = """
       CREATE TABLE IF NOT EXISTS roadmap (
@@ -301,12 +290,10 @@ def generate_sql_for_report(llm, prompt):
         which helps students for their JEE Preparation. Now your job is to analyze the user's 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, it tells the amount of time required to complete that specific task. So make sure
         to create queries that compare just the numbers within the text. For the task_type column,
         the data is either of these (Concept Understanding, Question Practice, Revision or Test)
-
         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 anything else like new line statements, ```sql or any other text. Your output
@@ -622,279 +609,171 @@ def generate_report(full_roadmap):
 
         st.session_state.final_report = state["final_report"]
 
-# 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..."):
+# AGENT 3
+def initialize_roadmap_db():
+    if not os.path.exists("jee_roadmap.db"):
         try:
-            # Initialize needed components
-            llm = ChatOpenAI(model="gpt-4o-mini", api_key = os.getenv("OPENAI_API_KEY"))
-            
-            # 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}")
+            with open("full_roadmap.json") as f:
+                roadmap_data = json.load(f)
 
-# 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!")
-            
+            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:
-            st.error(f"Error merging tasks: {e}")
+            print(f"⚠️ Error initializing database: {e}")
+
+def get_chapters_and_subtopics():
+    with open("full_roadmap.json", "r") as f:
+        data = json.load(f)
+
+    ch_subt = {
+        "Physics": {},
+        "Chemistry": {},
+        "Maths": {}
+    }
+
+    for day in data["schedule"]:
+        for subject in day['subjects']:
+            sub = ch_subt[subject['name']]
+            for task in subject['tasks']:
+                sub[task['ChapterName']] = []
+
+    for day in data["schedule"]:
+        for subject in day['subjects']:
+            sub = ch_subt[subject['name']]
+            for task in subject['tasks']:
+                if task['subtopic'] not in sub[task['ChapterName']]:
+                  sub[task['ChapterName']].append(task['subtopic'])
+
+    return ch_subt
+
+# Function to convert NL query to SQL
+def generate_sql_from_nl(prompt):
+    table_struct = """CREATE TABLE IF NOT EXISTS roadmap (
+                        id INTEGER PRIMARY KEY AUTOINCREMENT,
+                        day_num INTEGER,
+                        date TEXT, -- [yyyy-mm-dd]
+                        subject TEXT, -- [Physics, Chemistry or Maths]
+                        chapter_name TEXT,
+                        task_type TEXT, -- (Concept Understanding, Question Practice, Revision, Test)
+                        time TEXT, -- formatted like '0.5 hour', '1 hour', '2 Hours', and so on
+                        subtopic TEXT,
+                    )"""
+
+    ch_subt = get_chapters_and_subtopics()
+    response = llm.invoke(
+        [
+            SystemMessage(
+                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 to create queries that compare just the numbers within the text.
+             - If the student mention about any chapters or subtopics, browse through this json file {ch_subt},
+               find the one with the closest match to the users query and use only those exact names of Chapers
+               and Subtopics present in this file to create SQL the query.
+             - For date related queries, refer today's date {datetime.now().date()}
+
+             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.
+             """
+            ),
+            HumanMessage(
+                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 completed section
+    return response.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)
+    columns = [desc[0] for desc in cursor.description]
+    rows = cursor.fetchall()
+    data = {
+        "query": sql_query,
+        "columns": columns,
+        "rows": rows
+        }
+    conn.close()
+    return data
+
+# Function to convert SQL output to natural language
+def generate_nl_from_sql_output(prompt, data):
+    response = llm.invoke(
+        [
+            SystemMessage(
+                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. If you get empty data
+             as an input, deeply analyze the user's prompt and the sql query and give a suitable reply."""
+            ),
+            HumanMessage(
+                content=f"""Answer to this users query using the data given to you, while keeping
+             your role in mind: {prompt}"""
+            ),
+        ]
+    )
+
+    # Return completed section
+    return response.content.strip()
+
+# Main function for chatbot
+def answer_user_query(prompt):
+    initialize_roadmap_db()
+    query = generate_sql_from_nl(prompt)
+    data = fetch_data_from_sql(query)
+    return generate_nl_from_sql_output(prompt, data)
+
 
 # ---- HOME PAGE ----
 if page == "Home":
@@ -912,16 +791,6 @@ if page == "Home":
     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
@@ -931,77 +800,44 @@ if page == "Home":
             st.session_state.first_load = True
 
 # ---- ROADMAP MANAGER PAGE ----
-elif page == "Roadmap Manager":
+elif page == "Roadmap Manager": # AGENT 2
     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 1: Process Tasks")
+        if st.button("1️⃣ Mark Tasks as Incomplete"):
+            process_task_completion_data()
         
-        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()
+        st.subheader("Step 2: Reschedule Tasks")
+        if st.button("2️⃣ Optimize Task Distribution"):
+            shift_and_sort_the_roadmap(st.session_state.full_roadmap, 
+                                        st.session_state.data, 
+                                        dayNumber = 1, 
+                                        max_hours_per_day = 8)
                 
         # Display original and updated roadmaps side by side
-        if st.session_state.data_old and st.session_state.updated_roadmap:
+        if st.session_state.data 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)
+                    st.json(st.session_state.data)
             
             with col2:
                 st.markdown("#### Updated Roadmap")
                 with st.expander("View Updated Roadmap"):
                     st.json(st.session_state.updated_roadmap)
+                    for day in st.session_state.updated_roadmap['schedule']:
+                        st.write(f"Day: {day['dayNumber']} -> Total Time: {check_tot_time(day['subjects'], 8)[0]} Hours")
 
 # ---- TASK ANALYSIS PAGE ----
-elif page == "Task Analysis":
+elif page == "Task Analysis": # AGENT 1
     st.title("📊 Task Analysis")
     
     choice = st.selectbox("Choose the roadmap to use for building report", ["Four Day Roadmap", "Full Roadmap"])
@@ -1061,7 +897,7 @@ elif page == "Task Analysis":
             with col2:
                 st.subheader("Task Type Distribution")
                 st.bar_chart(type_counts)
-# ---- ROADMAP CHATBOT PAGE ----
+# ---- ROADMAP CHATBOT PAGE ---- # AGENT 3
 elif page == "Roadmap Chatbot":
     st.title("🤖 Roadmap Chatbot Assistant")