import streamlit as st
import pickle
import os
import time
import json

# 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

# 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()

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

# TAB 1: PLAN GENERATOR (Your existing code)
with tab1:
    # --- SIDEBAR FOR STUDENT PROFILE ---
    with st.sidebar:
        st.header("Student Profile")
        name = st.text_input("Name", "Chaitanya Kharche")
        gpa = st.slider("GPA", 0.0, 4.0, 3.5, 0.1)
        career_goal = st.text_area("Career Goal", "AI Engineer specializing in Large Language Models")
        interests = st.text_input("Interests (comma-separated)", "AI, Machine Learning, LLMs, Agentic AI")
        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)", "CS1800, CS2500")
        
        # 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")
    
    # LOAD DATA
    st.subheader("1. Load Curriculum Data")
    uploaded_file = st.file_uploader("Upload `neu_graph_analyzed_clean.pkl`", 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)
                # Also create visualizer
                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.")
    
    # GENERATE PLAN
    st.subheader("2. 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
        )
        
        # Check if profile changed
        profile_changed = st.session_state.last_profile != profile
        if profile_changed:
            st.session_state.last_profile = profile
        
        col1, col2, col3 = st.columns(3)
        
        if col1.button("🧠 AI-Optimized Plan", use_container_width=True, type="primary"):
            with st.spinner("🚀 Using LLM for intelligent course selection..."):
                start_time = time.time()
                result = optimizer.generate_llm_plan(profile)
                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("Generating personalized rule-based plan..."):
                start_time = time.time()
                result = optimizer.generate_simple_plan(profile)
                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 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:
                    for course_id in summer:
                        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}")
                        else:
                            st.write(f"• {course_id}")
                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:
            if st.button("📥 Export Plan as JSON", use_container_width=True):
                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)
                    }
                }
                plan_json = json.dumps(export_data, indent=2)
                st.download_button(
                    label="Download Complete Plan Data",
                    data=plan_json,
                    file_name=f"curriculum_plan_{name.replace(' ', '_')}.json",
                    mime="application/json"
                )

# 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:
        # Controls
        col1, col2 = st.columns([1, 3])
        
        with col1:
            show_critical = st.checkbox("Show Critical Path", True)
            if st.session_state.display_plan:
                highlight_plan = st.checkbox("Highlight My Courses", False)
        
        # Create visualization
        if st.session_state.visualizer:
            critical_path = []
            if show_critical:
                critical_path = st.session_state.visualizer.find_critical_path()
                if critical_path:
                    st.info(f"Critical Path ({len(critical_path)} courses): {' → '.join(critical_path[:5])}...")
            
            # 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
                
                **Metrics Explained:**
                - **Blocking Factor**: How many courses this prerequisite blocks
                - **Delay Factor**: Length of longest path through this course
                - **Centrality**: How important this course is in the curriculum network
                - **Complexity**: Combined score (research by Prof. Lionelle)
                """)

# 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}")
        
        # Calculate total curriculum complexity
        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']/200, 1.0))
                    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")
        
        # Export to CurricularAnalytics format
        st.divider()
        if st.button("📤 Export to CurricularAnalytics Format"):
            if st.session_state.visualizer:
                ca_format = st.session_state.visualizer.export_to_curricular_analytics_format({})
                st.download_button(
                    "Download CA Format JSON",
                    json.dumps(ca_format, indent=2),
                    "curriculum_analytics.json",
                    "application/json"
                )

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