src/ui.py

import streamlit as st
import pickle
import os
import time
import json
import yaml
from datetime import datetime
from typing import Dict, Set, Optional

# Import the optimizer and visualizer
from curriculum_optimizer import HybridOptimizer, StudentProfile
from interactive_visualizer import CurriculumVisualizer

# --- Page Configuration ---
st.set_page_config(page_title="Curriculum Optimizer", layout="wide", initial_sidebar_state="expanded")

# Initialize session state
if "display_plan" not in st.session_state:
    st.session_state.display_plan = None
if "metrics" not in st.session_state:
    st.session_state.metrics = None
if "reasoning" not in st.session_state:
    st.session_state.reasoning = ""
if "graph_data_loaded" not in st.session_state:
    st.session_state.graph_data_loaded = False
if "last_profile" not in st.session_state:
    st.session_state.last_profile = None
if "visualizer" not in st.session_state:
    st.session_state.visualizer = None
if "selected_track" not in st.session_state:
    st.session_state.selected_track = "general" # Default to general

# Title
st.title("🧑‍🎓 Next-Gen Curriculum Optimizer")

# --- Caching and Initialization ---
@st.cache_resource
def get_optimizer():
    """Loads and caches the main optimizer class and its models."""
    try:
        optimizer = HybridOptimizer()
        optimizer.load_models()
        return optimizer
    except Exception as e:
        st.error(f"Fatal error during model loading: {e}")
        st.info("Please ensure you have the required libraries installed.")
        st.stop()
        return None

optimizer = get_optimizer()

# --- DYNAMIC HELPER FUNCTIONS ---

def check_requirements_satisfaction(plan: Dict, track: str) -> Dict:
    """
    Check which requirements are satisfied by the plan.
    This is now dynamic based on the optimizer's config.
    """
    if not optimizer:
        return {}
        
    all_courses = []
    for year_key, year_data in plan.items():
        if year_key.startswith("year_"):
            all_courses.extend(year_data.get("fall", []))
            all_courses.extend(year_data.get("spring", []))
    all_courses_set = set(all_courses)

    # Get the correct requirements dictionary
    if track == "general":
        req_data = {
            "foundations": {"required": ["CS1800", "CS2500", "CS2510", "CS2800"]},
            "core": {"required": ["CS3000", "CS3500", "CS3650"]},
            "math": {"required": ["MATH1341", "MATH1342"], "pick_1_from": ["MATH2331", "MATH3081"]}
        }
    elif track == "game_dev":
         # Use ai_ml as a base for game_dev
         req_data = optimizer.CONCENTRATION_REQUIREMENTS.get("ai_ml", {})
    else:
         req_data = optimizer.CONCENTRATION_REQUIREMENTS.get(track, {})

    satisfaction_report = {}
    for category, reqs in req_data.items():
        report = {}
        if "required" in reqs:
            req_list = reqs["required"]
            report["required"] = req_list
            report["completed"] = list(all_courses_set & set(req_list))
            report["is_satisfied"] = all_courses_set.issuperset(req_list)
        
        for key, courses in reqs.items():
            if key.startswith("pick_"):
                try:
                    num_to_pick = int(key.split("_")[1])
                except Exception:
                    num_to_pick = 1
                
                completed_in_pick = list(all_courses_set & set(courses))
                report[key] = {
                    "options": courses,
                    "completed": completed_in_pick,
                    "count": f"{len(completed_in_pick)} of {num_to_pick}",
                    "is_satisfied": len(completed_in_pick) >= num_to_pick
                }
        satisfaction_report[category] = report
    
    return satisfaction_report

def export_plan_yaml(plan: Dict, profile: StudentProfile, validation: Dict = None, track: str = "general") -> str:
    """Export plan in structured YAML format for verification"""
    
    # Build structured plan data
    structured_plan = {
        "student_profile": {
            "name": profile.name if hasattr(profile, 'name') else "Student",
            "gpa": profile.current_gpa,
            "career_goal": profile.career_goals,
            "interests": profile.interests,
            "completed_courses": profile.completed_courses,
            "time_commitment": profile.time_commitment,
            "preferred_difficulty": profile.preferred_difficulty
        },
        "plan_metadata": {
            "generated": datetime.now().isoformat(),
            "track": track, # --- FIX: Now dynamic ---
            "total_credits": 0,
            "validation_status": "valid" if not validation.get("errors") else "has_errors"
        },
        "validation": validation if validation else {"errors": [], "warnings": []},
        "semesters": [],
        "course_details": {}
    }
    
    # Build semester list with full details
    total_credits = 0
    for year in range(1, 5):
        year_key = f"year_{year}"
        if year_key in plan:
            # Fall
            fall_courses = plan[year_key].get("fall", [])
            if fall_courses:
                semester_data = {"year": year, "term": "fall", "courses": []}
                for course_id in fall_courses:
                    course_info = optimizer.courses.get(course_id, {})
                    course_detail = {
                        "id": course_id,
                        "name": course_info.get("name", "Unknown"),
                        "credits": course_info.get("maxCredits", 4),
                        "complexity": course_info.get("complexity", 0),
                        "prerequisites": list(optimizer.curriculum_graph.predecessors(course_id)) if course_id in optimizer.curriculum_graph else []
                    }
                    semester_data["courses"].append(course_detail)
                    total_credits += course_detail["credits"]
                    structured_plan["course_details"][course_id] = course_detail
                
                semester_data["semester_credits"] = sum(c["credits"] for c in semester_data["courses"])
                semester_data["semester_complexity"] = sum(c["complexity"] for c in semester_data["courses"])
                structured_plan["semesters"].append(semester_data)

            # Spring
            spring_courses = plan[year_key].get("spring", [])
            if spring_courses:
                semester_data = {"year": year, "term": "spring", "courses": []}
                for course_id in spring_courses:
                    course_info = optimizer.courses.get(course_id, {})
                    course_detail = {
                        "id": course_id,
                        "name": course_info.get("name", "Unknown"),
                        "credits": course_info.get("maxCredits", 4),
                        "complexity": course_info.get("complexity", 0),
                        "prerequisites": list(optimizer.curriculum_graph.predecessors(course_id)) if course_id in optimizer.curriculum_graph else []
                    }
                    semester_data["courses"].append(course_detail)
                    total_credits += course_detail["credits"]
                    structured_plan["course_details"][course_id] = course_detail
                
                semester_data["semester_credits"] = sum(c["credits"] for c in semester_data["courses"])
                semester_data["semester_complexity"] = sum(c["complexity"] for c in semester_data["courses"])
                structured_plan["semesters"].append(semester_data)
        
        # Add summer/co-op
        if year in [2, 3]:
            structured_plan["semesters"].append({
                "year": year, "term": "summer", "activity": "co-op", "courses": []
            })
    
    structured_plan["plan_metadata"]["total_credits"] = total_credits
    
    # Calculate requirement satisfaction
    # --- FIX: Pass the dynamic track ---
    requirements_met = check_requirements_satisfaction(plan, track=track)
    structured_plan["requirements_satisfaction"] = requirements_met
    
    return yaml.dump(structured_plan, default_flow_style=False, sort_keys=False)


# --- UI TABS ---
tab1, tab2, tab3 = st.tabs(["📝 Plan Generator", "🗺️ Curriculum Map", "📊 Analytics"])

with tab1:
    # --- SIDEBAR FOR STUDENT PROFILE ---
    with st.sidebar:
        st.header("Student Profile")
        name = st.text_input("Name", "John, son of Jane")
        gpa = st.slider("GPA", 0.0, 4.0, 3.0, 0.1)
        career_goal = st.text_area("Career Goal", " ")
        interests = st.text_input("Interests (comma-separated)", " ")
        learning_style = st.selectbox("Learning Style", ["Visual", "Hands-on", "Auditory"])
        time_commit = st.number_input("Weekly Study Hours", 10, 60, 40, 5)
        difficulty = st.selectbox("Preferred Difficulty", ["easy", "moderate", "challenging"])
        completed_courses_input = st.text_area("Completed Courses (comma-separated)", " ")
        
        # Show profile impact
        st.markdown("---")
        st.markdown("**Profile Impact:**")
        if time_commit < 20:
            st.info("🕒 Part-time load (3 courses/semester)")
        elif time_commit >= 40:
            st.info("🔥 Intensive load (up to 5 courses/semester)")
        else:
            st.info("📚 Standard load (4 courses/semester)")
        
        if difficulty == "easy":
            st.info("😌 Focuses on foundational courses")
        elif difficulty == "challenging":
            st.info("🚀 Includes advanced/specialized courses")
        else:
            st.info("⚖️ Balanced difficulty progression")
    
    # --- MAIN PAGE CONTENT ---
    
    # 1. LOAD DATA
    st.subheader("1. Load Curriculum Data")
    uploaded_file = st.file_uploader("Upload `.pkl` file in the files section of this project", type=["pkl"])
    
    if uploaded_file and not st.session_state.graph_data_loaded:
        with st.spinner("Loading curriculum data and preparing embeddings..."):
            try:
                graph_data = pickle.load(uploaded_file)
                optimizer.load_data(graph_data)
                st.session_state.visualizer = CurriculumVisualizer(graph_data)
                st.session_state.graph_data = graph_data
                st.session_state.graph_data_loaded = True
                st.success(f"Successfully loaded and processed '{uploaded_file.name}'!")
                time.sleep(1)
                st.rerun()
            except Exception as e:
                st.error(f"Error processing .pkl file: {e}")
                st.session_state.graph_data_loaded = False
    elif st.session_state.graph_data_loaded:
        st.success("Curriculum data is loaded and ready.")
    
    # 2. SELECT TRACK (NEW SECTION)
    st.subheader("2. Select a Specialization")
    if not st.session_state.graph_data_loaded:
        st.info("Please load a curriculum file first.")
    else:
        # Map user-friendly names to the internal keys
        track_options = {
            "general": "🤖 General CS (Broadest Focus)",
            "ai_ml": "🧠 Artificial Intelligence & ML",
            "security": "🔒 Cybersecurity",
            "systems": "⚙️ Systems & Networks",
            "game_dev": "🎮 Game Design & Development"
        }
        
        selected_track_key = st.selectbox(
            "Choose your focus area (optional):",
            options=track_options.keys(),
            format_func=lambda key: track_options[key], # Shows the friendly name
            index=0  # Default to "General"
        )
        st.session_state.selected_track = selected_track_key

    # 3. GENERATE PLAN
    st.subheader("3. Generate a Plan")
    if not st.session_state.graph_data_loaded:
        st.info("Please load a curriculum file above to enable plan generation.")
    else:
        # Create student profile
        profile = StudentProfile(
            completed_courses=[c.strip().upper() for c in completed_courses_input.split(',') if c.strip()],
            current_gpa=gpa, 
            interests=[i.strip() for i in interests.split(',') if i.strip()],
            career_goals=career_goal, 
            learning_style=learning_style,
            time_commitment=time_commit, 
            preferred_difficulty=difficulty
        )
        
        # Get the selected track from session state
        selected_track = st.session_state.get("selected_track", "general")
        
        # Check if profile or track changed
        profile_changed = (st.session_state.last_profile != profile) or \
                          (st.session_state.last_track != selected_track)
                          
        if profile_changed:
            st.session_state.last_profile = profile
            st.session_state.last_track = selected_track
        
        col1, col2, col3 = st.columns(3)
        
        if col1.button("🧠 AI-Optimized Plan", use_container_width=True, type="primary"):
            with st.spinner(f"🚀 Performing AI-optimization for '{track_options[selected_track]}' track..."):
                start_time = time.time()
                # --- FIX: Pass selected_track ---
                result = optimizer.generate_llm_plan(profile, selected_track)
                generation_time = time.time() - start_time
                
                plan_raw = result.get('pathway', {})
                st.session_state.reasoning = plan_raw.get("reasoning", "")
                st.session_state.metrics = plan_raw.get("complexity_analysis", {})
                st.session_state.display_plan = plan_raw
                st.session_state.plan_type = "AI-Optimized"
                st.session_state.generation_time = generation_time
                st.success(f"🎉 AI-optimized plan generated in {generation_time:.1f}s!")
        
        if col2.button("⚡ Smart Rule-Based Plan", use_container_width=True):
            with st.spinner(f"Generating rule-based plan for '{track_options[selected_track]}' track..."):
                start_time = time.time()
                # --- FIX: Pass selected_track ---
                result = optimizer.generate_simple_plan(profile, selected_track)
                generation_time = time.time() - start_time
                
                plan_raw = result.get('pathway', {})
                st.session_state.reasoning = plan_raw.get("reasoning", "")
                st.session_state.metrics = plan_raw.get("complexity_analysis", {})
                st.session_state.display_plan = plan_raw
                st.session_state.plan_type = "Smart Rule-Based"
                st.session_state.generation_time = generation_time
                st.success(f"🎉 Smart rule-based plan generated in {generation_time:.1f}s!")
        
        if col3.button("🔄 Clear Plan", use_container_width=True):
            st.session_state.display_plan = None
            st.session_state.metrics = None
            st.session_state.reasoning = ""
            st.rerun()
    
    # Show profile change notification
    if st.session_state.display_plan and profile_changed:
        st.warning("⚠️ Student profile or track changed! Generate a new plan to see updated recommendations.")
    
    # DISPLAY RESULTS
    if st.session_state.display_plan:
        st.subheader(f"📚 {st.session_state.get('plan_type', 'Optimized')} Degree Plan")
        
        # Display generation info
        col_info1, col_info2, col_info3 = st.columns(3)
        with col_info1:
            st.metric("Generation Time", f"{st.session_state.get('generation_time', 0):.1f}s")
        with col_info2:
            st.metric("Plan Type", st.session_state.get('plan_type', 'Unknown'))
        with col_info3:
            if time_commit < 20:
                load_type = "Part-time"
            elif time_commit >= 40:
                load_type = "Intensive"
            else:
                load_type = "Standard"
            st.metric("Course Load", load_type)
        
        # Display reasoning and metrics
        if st.session_state.reasoning or st.session_state.metrics:
            st.markdown("##### 📊 Plan Analysis")
            
            if st.session_state.reasoning:
                st.info(f"**Strategy:** {st.session_state.reasoning}")
            
            if st.session_state.metrics:
                m = st.session_state.metrics
                c1, c2, c3, c4 = st.columns(4)
                
                c1.metric("Avg Complexity", f"{m.get('average_semester_complexity', 0):.1f}")
                c2.metric("Peak Complexity", f"{m.get('peak_semester_complexity', 0):.1f}")
                c3.metric("Total Complexity", f"{m.get('total_complexity', 0):.0f}")
                c4.metric("Balance Score", f"{m.get('balance_score (std_dev)', 0):.2f}")
            
            st.divider()
        
        # Display the actual plan
        plan = st.session_state.display_plan
        total_courses = 0
        
        for year_num in range(1, 5):
            year_key = f"year_{year_num}"
            year_data = plan.get(year_key, {})
            
            st.markdown(f"### Year {year_num}")
            col_fall, col_spring, col_summer = st.columns(3)
            
            # Fall semester
            with col_fall:
                fall_courses = year_data.get("fall", [])
                st.markdown("**🍂 Fall Semester**")
                if fall_courses:
                    for course_id in fall_courses:
                        if course_id in optimizer.courses:
                            course_data = optimizer.courses[course_id]
                            course_name = course_data.get("name", course_id)
                            st.write(f"• **{course_id}**: {course_name}")
                            total_courses += 1
                        else:
                            st.write(f"• {course_id}")
                            total_courses += 1
                else:
                    st.write("*No courses scheduled*")
            
            # Spring semester
            with col_spring:
                spring_courses = year_data.get("spring", [])
                st.markdown("**🌸 Spring Semester**")
                if spring_courses:
                    for course_id in spring_courses:
                        if course_id in optimizer.courses:
                            course_data = optimizer.courses[course_id]
                            course_name = course_data.get("name", course_id)
                            st.write(f"• **{course_id}**: {course_name}")
                            total_courses += 1
                        else:
                            st.write(f"• {course_id}")
                            total_courses += 1
                else:
                    st.write("*No courses scheduled*")
            
            # Summer
            with col_summer:
                summer = year_data.get("summer", [])
                st.markdown("**☀️ Summer**")
                if summer == "co-op":
                    st.write("🏢 *Co-op Experience*")
                elif summer:
                    # This case isn't really used by the optimizer, but good to have
                    st.write("*Summer Classes*")
                else:
                    st.write("*Break*")
        
        # Summary and export
        st.divider()
        col_export1, col_export2 = st.columns(2)
        
        with col_export1:
            st.metric("Total Courses", total_courses)
        
        with col_export2:
            col_yaml, col_json = st.columns(2)
            
            with col_yaml:
                # --- FIX: Get validation from the plan object, DO NOT re-run validate_plan() ---
                validation = st.session_state.display_plan.get("validation", {"errors": [], "warnings": []})
                
                yaml_data = export_plan_yaml(
                    st.session_state.display_plan,
                    profile,
                    validation,
                    st.session_state.get("selected_track", "general") # Pass track
                )
                st.download_button(
                    label="📥 Export as YAML",
                    data=yaml_data,
                    file_name=f"curriculum_plan_{name.replace(' ', '_')}.yaml",
                    mime="text/yaml",
                    use_container_width=True
                )
            
            with col_json:
                export_data = {
                    "student_profile": {
                        "name": name, "gpa": gpa, "career_goals": career_goal,
                        "interests": interests, "learning_style": learning_style,
                        "time_commitment": time_commit, "preferred_difficulty": difficulty,
                        "completed_courses": completed_courses_input
                    },
                    "plan": st.session_state.display_plan,
                    "metrics": st.session_state.metrics,
                    "generation_info": {
                        "plan_type": st.session_state.get('plan_type', 'Unknown'),
                        "generation_time": st.session_state.get('generation_time', 0),
                        "selected_track": st.session_state.get("selected_track", "general")
                    }
                }
                plan_json = json.dumps(export_data, indent=2)
                st.download_button(
                    label="📥 Export as JSON",
                    data=plan_json,
                    file_name=f"curriculum_plan_{name.replace(' ', '_')}.json",
                    mime="application/json",
                    use_container_width=True
                )

# --- TAB 2: CURRICULUM MAP ---
with tab2:
    st.subheader("🗺️ Interactive Curriculum Dependency Graph")
    
    if not st.session_state.graph_data_loaded:
        st.info("Please load curriculum data in the Plan Generator tab first.")
    else:
        # Create visualization
        if st.session_state.visualizer:
            critical_path = st.session_state.visualizer.find_critical_path()
            if critical_path:
                st.info(f"Global Critical Path ({len(critical_path)} courses): {' → '.join(critical_path[:7])}...")
            
            # Create the plot
            fig = st.session_state.visualizer.create_interactive_plot(critical_path)
            st.plotly_chart(fig, use_container_width=True)
            
            # Legend
            with st.expander("📖 How to Read This Graph"):
                st.markdown("""
                **Node (Circle) Size**: Blocking factor - larger circles block more future courses  
                **Node Color**: Complexity score - darker = more complex  
                **Lines**: Prerequisite relationships  
                **Red Path**: Critical path (longest chain)  
                **Hover over nodes**: See detailed metrics for each course
                """)

# --- TAB 3: ANALYTICS ---
with tab3:
    st.subheader("📊 Curriculum Analytics Dashboard")
    
    if not st.session_state.graph_data_loaded:
        st.info("Please load curriculum data in the Plan Generator tab first.")
    else:
        # Overall metrics
        col1, col2, col3, col4 = st.columns(4)
        
        graph = st.session_state.graph_data
        total_courses = graph.number_of_nodes()
        total_prereqs = graph.number_of_edges()
        
        col1.metric("Total Courses", total_courses)
        col2.metric("Total Prerequisites", total_prereqs)
        col3.metric("Avg Prerequisites", f"{total_prereqs/total_courses:.1f}")
        
        if st.session_state.visualizer:
            total_complexity = sum(
                st.session_state.visualizer.calculate_metrics(n)['complexity']
                for n in graph.nodes()
            )
            col4.metric("Curriculum Complexity", f"{total_complexity:,.0f}")
        
        st.divider()
        
        # Most complex courses
        col1, col2 = st.columns(2)
        
        with col1:
            st.subheader("Most Complex Courses")
            if st.session_state.visualizer:
                complexities = []
                for node in graph.nodes():
                    metrics = st.session_state.visualizer.calculate_metrics(node)
                    complexities.append({
                        'course': node,
                        'name': graph.nodes[node].get('name', ''),
                        'complexity': metrics['complexity'],
                        'blocking': metrics['blocking']
                    })
                
                complexities.sort(key=lambda x: x['complexity'], reverse=True)
                
                for item in complexities[:10]:
                    st.write(f"**{item['course']}**: {item['name']}")
                    prog_col1, prog_col2 = st.columns([3, 1])
                    with prog_col1:
                        st.progress(min(item['complexity']/100, 1.0)) # Adjusted scale
                    with prog_col2:
                        st.caption(f"Blocks: {item['blocking']}")
        
        with col2:
            st.subheader("Bottleneck Courses")
            st.caption("(High blocking factor)")
            
            if st.session_state.visualizer:
                bottlenecks = sorted(complexities, key=lambda x: x['blocking'], reverse=True)
                
                for item in bottlenecks[:10]:
                    st.write(f"**{item['course']}**: {item['name']}")
                    st.info(f"Blocks {item['blocking']} future courses")
        
        # Plan vs Global Comparison
        if st.session_state.display_plan:
            st.divider()
            st.subheader("📊 Metric System Comparison")
            st.caption("Comparing metrics for the entire curriculum vs. metrics only within your generated plan.")
            
            plan_courses: Set[str] = set()
            for year_key, year_data in st.session_state.display_plan.items():
                if year_key.startswith("year_"):
                    plan_courses.update(year_data.get("fall", []))
                    plan_courses.update(year_data.get("spring", []))
            
            comparison = st.session_state.visualizer.compare_metric_systems(plan_courses)
            
            col1, col2 = st.columns(2)
            
            with col1:
                st.metric(
                    "Critical Path Match", 
                    "✅ Yes" if comparison['critical_path_match'] else "❌ No"
                )
                st.caption("Global critical path (first 5):")
                st.code(' → '.join(comparison['global_critical']))
            
            with col2:
                st.metric(
                    "Major Metric Differences",
                    len(comparison['major_differences'])
                )
                st.caption("Plan-specific critical path (first 5):")
                st.code(' → '.join(comparison['plan_critical']))
            
            if comparison['major_differences']:
                with st.expander(f"View {len(comparison['major_differences'])} courses with >50% metric difference"):
                    for diff in comparison['major_differences']:
                        st.write(f"**{diff['course']}**: Global blocking={diff['global_blocking']}, Plan blocking={diff['plan_blocking']}")

# Footer
st.divider()
st.caption("🚀 Powered by Students, For Students")