Update app.py

app.py

@@ -1,447 +1,581 @@
 import json
-import operator
 import os
-import getpass
+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
+from typing_extensions import TypedDict, Literal
 from langgraph.graph import StateGraph, START, END
-from IPython.display import Image, display
-from langchain_core.messages import HumanMessage, SystemMessage, ToolMessage
-from langgraph.graph import MessagesState
-from typing_extensions import Literal
-from langchain_core.tools import tool
-import json
-from datetime import datetime, timedelta
-import streamlit as st
-import copy
-from openai import OpenAI
-st.set_page_config(layout="wide")
-
-st.title("RoadMap Test")
-
-
-# # setting LLM
-# def _set_env(var: str):
-#     if not os.environ.get(var):
-#         os.environ[var] = getpass.getpass(f"{var}: ")
-
-
-# _set_env("ANTHROPIC_API_KEY")
-
-os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")
-
-llm = ChatOpenAI(model="gpt-4o-mini")
-
 
-with open('fourdayRoadmap.json', 'r') as file:
-    data = json.load(file)  # Load JSON content into a Python dictionary                                                   -->display 1
-
-
-data_old = copy.deepcopy(data)         # Added for stramlit app,not the part of actual code 
-
-#Intermediate Step : Add variables task_completed as false and completion_timestamp as null
-for day in data["schedule"]:
-    for subject in day["subjects"]:
-        for task in subject["tasks"]:
-            task["task_completed"] = False
-            task["completion_timestamp"] = None
-
-# Step 2 : Extract The previous day from the extracted roadmap
-previous_day = data["schedule"][0]
-previous_day_roadmap_str = str(previous_day)
-
-#Step 3 : Extract the incompleted tasks
-# Extract incomplete tasks
-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
+# 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"])
+
+# 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 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)
 
-#Generate a incomplete_tasks_list for the the agent 2
-# Extract incomplete tasks
-incomplete_task_list = []
-incomplete_task_list_json = json.loads(incomplete_tasks_json)
-for subject in incomplete_task_list_json["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"]
+        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": []
             }
-            incomplete_task_list.append(incomplete_task)
-
-#Now we have to make an Orchestrator - worker workflow which generates a report on the roadmap data and defining sections which are most relevant and so giving a report
-
-
-# Schema for structured output to use in planning
-
-# class Section(BaseModel):  # This part defines the structure of a single section of the report, saying that it should return a field name and a description
-#     name: str = Field(
-#         description="Name for this section of the report.",
-#     )
-#     description: str = Field(
-#         description="""Very short overview of the main topics and concepts to be covered in this section.""",
-#     )
-
-
-# class Sections(BaseModel): # This part defines what has to be the output of the llm
-#     sections: List[Section] = Field(
-#         description="Sections of the report.",
-#     )
-
-
-# # Augment the LLM with schema for structured output
-# planner = llm.with_structured_output(Sections)
-
-from langgraph.constants import Send
-
-# This is a state which defines all params that we need to collect in this WF
-# Graph state
-# class State(TypedDict):
-#     previous_day_roadmap: str  # Report topic
-#     sections: list[Section]  # List of report sections
-#     completed_sections: Annotated[
-#         list, operator.add
-#     ]  # All workers write to this key in parallel
-#     final_report: str  # Final report
-
-
-# # Worker state
-# class WorkerState(TypedDict):
-#     section: Section
-#     completed_sections: Annotated[list, operator.add]
-
-
-# # Nodes
-# def orchestrator(state: State):
-#     """Orchestrator that generates a plan for the report"""
-
-#     # Generate queries
-#     report_sections = planner.invoke(
-#         [
-#             SystemMessage(content="""Generate a plan for the report. The report is a review and analysis output, which tells us about
-#             the study pattern, studying hours, non studying hours and future action plan for the user based on
-#             performance of task completion on a day of the user. We only want to understand about how the student studies nothing apart from that
-#             Dont halucinate the data, just keep focus on provided data, no need for data out of the box.
-#             I want a very short report to understand how many tasks student completed and then what is the studying pattern of the user
-#             **Do not generate more than 5 sections**
-#             """),
-#             HumanMessage(content=f"Here is the Previous day roadmap of the user: {state['previous_day_roadmap']}"),
-#         ]
-#     )
-
-#     return {"sections": report_sections.sections}
-
-
-# def llm_call(state: WorkerState):
-#     """Worker writes a section of the report"""
-
-#     # Generate section
-#     section = llm.invoke(
-#         [
-#             SystemMessage(
-#                 content="""Write a report section following the provided name and description. Make this report userful for parents, teachers
-#                  and students himself, also try to motivate the student.This report is for a Joint Entrance Examination aspirant.
-#                  Keep the report very short but will all clear parameter and make sure the report is not so big but contains all the necessary
-#                  details of the student's day. Keep the report very short and avoid texts and focus on numbers
-#                  Include no preamble for each section.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."""
-#             ),
-#             HumanMessage(
-#                 content=f"Here is the section name: {state['section'].name} and description: {state['section'].description}"
-#             ),
-#         ]
-#     )
-
-#     # Write the updated section to completed sections
-#     return {"completed_sections": [section.content]}
-
-
-# def synthesizer(state: State):
-#     """Synthesize full report from sections"""
-
-#     # List of completed sections
-#     completed_sections = state["completed_sections"]
-
-#     # Format completed section to str to use as context for final sections
-#     completed_report_sections = "\n\n---\n\n".join(completed_sections)
-
-#     return {"final_report": completed_report_sections}
-
-
-# # Conditional edge function to create llm_call workers that each write a section of the report
-# def assign_workers(state: State):
-#     """Assign a worker to each section in the plan"""
-
-#     # Kick off section writing in parallel via Send() API
-#     return [Send("llm_call", {"section": s}) for s in state["sections"]]
-
-
-# # Build workflow
-# orchestrator_worker_builder = StateGraph(State)
-
-# # Add the nodes
-# orchestrator_worker_builder.add_node("orchestrator", orchestrator)
-# orchestrator_worker_builder.add_node("llm_call", llm_call)
-# orchestrator_worker_builder.add_node("synthesizer", synthesizer)
-
-# # Add edges to connect nodes
-# orchestrator_worker_builder.add_edge(START, "orchestrator")
-# orchestrator_worker_builder.add_conditional_edges(
-#     "orchestrator", assign_workers, ["llm_call"]
-# )
-# orchestrator_worker_builder.add_edge("llm_call", "synthesizer")
-# orchestrator_worker_builder.add_edge("synthesizer", END)
-
-# # Compile the workflow
-# orchestrator_worker = orchestrator_worker_builder.compile()
-
-# # # Show the workflow
-# # display(Image(orchestrator_worker.get_graph().draw_mermaid_png()))
-
-# # Invoke
-# state = orchestrator_worker.invoke({"previous_day_roadmap": f"{previous_day_roadmap_str}"})
-
-# from IPython.display import Markdown
-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
-final_report = output                                                                                   #-->display 2
-
-#Evaluator-optimizer approach
-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'])
-print(available_dates)
-
-# Graph state
-class State(TypedDict):
-    roadmap: dict
-    available_dates: dict
-    incomplete_task_list: list
-    feedback: str
-    added_or_not: str
-
-
-# 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 not added, give feedback to add those tasks also",
-    )
-
-
-# Augment the LLM with schema for structured output
-evaluator = llm.with_structured_output(Feedback)
-
-
-# Nodes
-def llm_call_generator(state: State):
-    """The function which addes the roadmap and incomplete tasks list to the state"""
-    roadmap_str = """
-    {
-      'roadmap':[
-        'date':'YYYY-MM-DD',
-        'tasks':[
-          <List of tasks>
-        ]
-      ]
-    }
-    """
-    if state.get("feedback"):
-        msg = llm.invoke(
-            f"""Add the following incomplete_tasks {state['incomplete_task_list']} to the roadmap key and take into account the feedback {state['feedback']}
-            and make sure that we dyanamically add the tasks not increasing the load on just one day, also we have to add the tasks on
-            following dates only {state['available_dates']} Make sure you only give roadmap json as output and nothing else, strictly follow the output structure {roadmap_str}
-            """
+            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:
-      ## In this case we have to just add the roadmap and incomplete tasks list to the state
-      msg = llm.invoke(f"""Add the following incomplete_tasks {state['incomplete_task_list']} to the roadmap key
-            and make sure that we dyanamically add the tasks not increasing the load on just one day, also we have to add the tasks on
-            following dates only {state['available_dates']} Make sure you only give roadmap json as output and nothing else, strictly follow the output structure {roadmap_str}
-            """)
-    return {"roadmap": msg.content}
-
-def llm_call_evaluator(state: State):
-    """This the evaluator that checks where all the tasks are added to the roadmap or not"""
-    grade = evaluator.invoke(f"Grade the roadmap {state['roadmap']} by checking the {state['incomplete_task_list']} as all added or not and make sure that we dyanamically add the tasks not increasing the load on just one day, also we need to make sure that we return an roadmap with proper json format as the roadmap was")
-    return {"added_or_not": grade.grade, "feedback": grade.feedback}
-
-
-# Conditional edge function to route back to joke generator or end based upon feedback from the evaluator
-def route_roadmap(state: State):
-    """Route back to the generator if  roadmap does not have all the incomplete tasks"""
-    if state["added_or_not"] == "added":
-        return "Accepted"
-    elif state["added_or_not"] == "not added":
-        return "Rejected + Feedback"
-
-
-# Build workflow
-optimizer_builder = StateGraph(State)
-
-# Add the nodes
-optimizer_builder.add_node("llm_call_generator", llm_call_generator)
-optimizer_builder.add_node("llm_call_evaluator", llm_call_evaluator)
-
-# Add edges to connect nodes
-optimizer_builder.add_edge(START, "llm_call_generator")
-optimizer_builder.add_edge("llm_call_generator", "llm_call_evaluator")
-optimizer_builder.add_conditional_edges(
-    "llm_call_evaluator",
-    route_roadmap,
-    {  # Name returned by route_roadmap : Name of next node to visit
-        "Accepted": END,
-        "Rejected + Feedback": "llm_call_generator",
-    },
-)
-
-# Compile the workflow
-optimizer_workflow = optimizer_builder.compile()
-
-# # Show the workflow
-# display(Image(optimizer_workflow.get_graph().draw_mermaid_png()))
-
-# Invoke
-state = optimizer_workflow.invoke({"available_dates": available_dates, "incomplete_task_list":incomplete_task_list})
-shifted_tasks_roadmap = state["roadmap"]    
-
-
-
-#Now we merge the shifted roadmap
-def add_task(roadmap, task, date_task_to_be_added):
-    subject_name = task["subject"]  # Extract 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:
-            for subject in day['subjects']:
-                if subject['name'] == subject_name:
-                    # Check if task already exists
-                    for existing_task in subject['tasks']:
-                        if existing_task['ChapterName'] == chapter_name and existing_task['type'] == type_name and existing_task['subtopic'] == topic_name:
-                            return roadmap  # Task already exists, return unchanged roadmap
-
-                    # Add task
-                    temp_task = {
-                        "ChapterName": chapter_name,
-                        "type": type_name,
-                        "topic": topic_name,
-                        "time": task["time"],
-                        "task_completed": False,
-                        "completion_timestamp": None
-                    }
-                    subject['tasks'].append(temp_task)
-                    return roadmap  # Return updated roadmap
+        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()
                 
-shifted_tasks_roadmap = shifted_tasks_roadmap.split("```json")[1].split("```")[0]
-for day in json.loads(shifted_tasks_roadmap)["roadmap"]:
-  date = day['date']
-  tasks = day['tasks']
-  for task in tasks:
-    add_task(data,task,date)
-
-
-
-
-#stramlit section 
-
-# Initialize session state for page navigation if it doesn't exist
-if "page" not in st.session_state:
-    st.session_state["page"] = 1
-
-def show_page_1():
-    st.title("Roadmaps")
-
-    col1, col2 = st.columns([1, 1])  
-
-    with col1:
-        st.subheader("Original Roadmap")
-        st.json(data_old)
-
-    with col2:
-        st.subheader("Roadmap with Shifted Task")
-        st.json(data)
-
-
-def show_page_2():
-    st.title("Report")
-    st.markdown(final_report)
-
-# Navigation
-st.sidebar.title("Navigation")
-if st.sidebar.button("Roadmaps Page"):
-    st.session_state["page"] = 1
-if st.sidebar.button("Report Page"):
-    st.session_state["page"] = 2
-
-# Display content based on the current page
-if st.session_state["page"] == 1:
-    show_page_1()
-elif st.session_state["page"] == 2:
-    show_page_2()
+            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)