import os
import sys

current_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(current_dir)
sys.path.append(parent_dir)

import streamlit as st
import numpy as np
import pandas as pd
import tensorflow as tf
import pickle
from utils.data_helper import *
from environment import raw_world as Env
from utils import irl_helper
from utils.trajectory_comparison import analyze_chapter_engagement
from typing import List, Dict, Union, Tuple
import plotly.graph_objects as go
st.set_page_config(layout="wide")
from utils.evaluation import *
from utils.rnn_models import *
from PIL import Image
import re

def performance_prediction(model, x_test):
    y_pred = model.predict(x_test)
    y_pred = np.array([1 if y[0] >= 0.5 else 0 for y in y_pred])
    num_fail = sum(y_pred) 
    num_pass = len(y_pred) - num_fail
    return num_fail, num_pass, y_pred


def extract_signal_name_regex(input_string):
    match = re.match(r"^\d+\.\d\s*-\s*(.*?)\s*\[\d{1,2}:\d{2}\]$", input_string)
    if match:
        return match.group(1).strip()
    return input_string # Return original if no match

def plot_performance(df=None, show_image=False, show_performance=False, col_image=None, col_performance=None):
    if show_image and col_image is not None:
        with col_image:
            fig = Image.open('streamlit-assets/whatif_results_week_6.png') 
            st.image(fig, caption="Predicted effectiveness of interventions in week 6. \
                        For example, adding content from topic 7 might improve students' performance, \
                            while adding content from week 2 might harm", use_container_width=True)
    elif show_image and col_image is None:
        fig = Image.open('streamlit-assets/whatif_results_week_6.png') 
        st.image(fig, caption="Predicted effectiveness of interventions in week 6. \
                        For example, adding content from topic 7 might improve students' performance, \
                            while adding content from week 2 might harm", use_container_width=False)

    if show_performance and df is not None and col_performance is not None:
        with col_performance:
            fig = go.Figure()
            bar_height_px = 250  
            total_height = bar_height_px * len(df)
            fig.add_trace(go.Bar(
                y=df['class'],
                x=df['fail_percent'],
                name='Fail',
                orientation='h',
                text=[f'{x:.1f}%' for x in df['fail_percent']],
                textfont=dict(size=25),
                textposition='auto',
                customdata=df[['fail_num']], 
                hovertemplate='%{y}<br> No. Students: %{customdata[0]}<extra></extra>',
                marker=dict(color='#c0392b')
            ))

            fig.add_trace(go.Bar(
                y=df['class'],
                x=df['pass_percent'],
                name='Pass',
                orientation='h',
                text=[f'{x:.1f}%' for x in df['pass_percent']],
                textfont=dict(size=25),
                textposition='auto',
                customdata=df[['pass_num']],  
                hovertemplate='%{y}<br> No. Students: %{customdata[0]}<extra></extra>',
                marker=dict(color='#27ae60')
            ))

            fig.update_layout(
                barmode='stack',
                xaxis=dict(title='Percentage', range=[0, 100]),
                legend=dict(orientation='h', yanchor='bottom', y=1.02, xanchor='right', x=1),
                height=total_height,
                margin=dict(l=200, r=40, t=120, b=60),  
            )
            st.plotly_chart(fig, use_container_width=True)

    
def compare_chapter_engagement(syn_trajectories: Union[List, np.ndarray, pd.Series],
                             real_trajectories: Union[List, np.ndarray, pd.Series],
                             world,
                             save_dir=None) -> Tuple[Dict, Dict]:
    """
    Compare engagement metrics between real and synthetic students per chapter
    
    Parameters:
    -----------
    real_trajectories : Union[List, np.ndarray, pd.Series]
        Real student trajectories
    syn_trajectories : Union[List, np.ndarray, pd.Series]
        Synthetic student trajectories
    world : ClickstreamWorld
        World object containing state and action information
    save_dir : str, optional
        Directory to save visualizations
        
    Returns:
    --------
    Tuple[Dict, Dict]: Real and synthetic engagement metrics
    """
    st.header("Engagement Analysis: Before vs. After Intervention")
    
    real_metrics = analyze_chapter_engagement(real_trajectories, world, "Real")
    # print('finished real metrics, starting synthetic metrics')
    syn_metrics = analyze_chapter_engagement(syn_trajectories, world, "Synthetic")
    
    chapter_id_to_name = st.session_state.metadata.set_index('Topic')['Skills'].to_dict()
    metrics_to_plot = [
    ('visit_count', 'Total Visits'),
    ('completion_rate', 'Completion Rate (%)'),
    ('problem_attempts', 'Quiz Attempts'),
    ('video_views', 'Video Views')
    ]
    increase_color = '#27ae60'  # Darker green
    decrease_color = '#c0392b' 
    neutral_color = '#bdc3c7'
    plotly_figures = []

    for idx, (metric, ylabel) in enumerate(metrics_to_plot):
        chapters = sorted(set(real_metrics.keys()) | set(syn_metrics.keys()))
        chapters_name = [chapter_id_to_name.get(ch, f'Chapter {ch}') for ch in chapters]
        real_values = [real_metrics.get(ch, {}).get(metric, 0) for ch in chapters]
        syn_values = [syn_metrics.get(ch, {}).get(metric, 0) for ch in chapters]
        pct_diffs = []
        colors = []
       
        for real_val, syn_val in zip(real_values, syn_values):
            
            if metric == 'completion_rate':
                # print(syn_val, real_val)
                pct_diff = syn_val - real_val
                colors.append(increase_color if pct_diff > 0 else decrease_color)
            else:
                if real_val > 0:
                    pct_diff = ((syn_val - real_val) / real_val) * 100
                    colors.append(increase_color if pct_diff > 0 else decrease_color)
                else:
                    # print(real_values, syn_values, metric, chapters_name)
                    # pct_diff = 100 if syn_val > 0 else 0
                    pct_diff = 0
                    colors.append(increase_color if syn_val > 0 else neutral_color)
            pct_diffs.append(pct_diff)

        fig = go.Figure()

        fig.add_trace(
            go.Bar(
                x=chapters,
                y=pct_diffs,
                marker=dict(color=colors, line=dict(color='#2c3e50', width=0.5)),
                name=ylabel,
                showlegend=False,
                text=[f'{val:+.1f}%' for val in pct_diffs],
                textposition='auto',
                hovertemplate='%{x}<br>%{y:.1f}%<extra></extra>',
                width=0.6
            )
        )

        fig.update_layout(
            title=dict(
                text=ylabel,
                x=0.5,
                xanchor='center',
                font=dict(size=20)
            ),
            font=dict(size=14),
            margin=dict(l=50, r=40, t=60, b=50),
            hovermode="x unified",
            bargap=0.25,
            template='plotly'
            )

        fig.update_xaxes(
            title_text='Topic',
            title_font=dict(size=12),
            tickvals=chapters,
            ticktext=[f'{ch}' for ch in chapters_name],
            tickfont=dict(size=13),
            showgrid=False,
            zeroline=False
        )

        fig.update_yaxes(
            title_text='Change (%)',
            title_font=dict(size=16),
            tickfont=dict(size=13),
            zeroline=True,
            zerolinecolor='white',
            zerolinewidth=1.5,
            showgrid=True,
        )


        plotly_figures.append(fig)

    col1, col2 = st.columns(2)
    with col1:
        st.caption("Total Visits: Number of interactions made by students for every course material group by Topic")
        st.plotly_chart(plotly_figures[0], use_container_width=True)
    with col2:
        st.caption("""Completion Rate: The percentage of course materials that students interacted with within each topic. A value of 100% means the student engaged with every material in that topic at least once.\\
                   For example: if there are 5 course material in a topic, the completion rate would be 80% if they interacted with 4 unique materials.""")
        st.plotly_chart(plotly_figures[1], use_container_width=True)

    col3, col4 = st.columns(2)
    with col3:
        st.caption("Quiz Attempts: Number of quiz attempts made by students for every quiz group by Topic")
        st.plotly_chart(plotly_figures[2], use_container_width=True)
    with col4:
        st.caption("Video Views: Number of video views made by students for every video group by Topic")
        st.plotly_chart(plotly_figures[3], use_container_width=True)
        
def add_new_state(world, trajectories, event_data, week_list=range(7,11), test_ids=None,
                  OUTPUT_DIR='streamlit-assets/dsp-004', trajectories_each_week=None, trajectories_each_week_pass=None, trajectories_each_week_fail=None,
                  history_whatif_pass=None, history_whatif_fail=None, fail_only=False):
    
    print("Adding new state with parameters:", event_data)
    
    results = {
        'week': [],
        'real_trajectories': [],
        'syn_trajectories': [],
        'world': []
    }
    
    for week in week_list:
        print(f"\nProcessing week {week}")
        world._reset_transition_prob_table()
        world._update_transition_prob_table(trajectories)
        # add_feature_matrix = world.designed_features(world.values)

        chapter = event_data['chapter']
        is_problem = True if event_data['event_type'] in 'quiz' else False
        value = float(event_data['difficulty']) if is_problem else float(event_data['duration'])
        
        print(f'Adding event - Topic: {chapter}, Type: {"Quiz" if is_problem else "Video"}, Value: {value}')
        add_feature_matrix = world.add_new_state(chapter=chapter, 
                                                 value=value, 
                                                 is_problem=is_problem, 
                                                 predict=False)
        
        pred = pd.read_csv(f'{OUTPUT_DIR}/predictions_weeks_{week}.csv')
        pred = pred['y_pred']
        if test_ids is None:
            real_student = trajectories_each_week[week-1]
        else:
            real_student = trajectories_each_week[week-1][test_ids]
        if fail_only:
            real_student = [real_student.iloc[i] for i in range(len(real_student)) if pred[i] == 1]
        else:
            real_student = [real_student.iloc[i] for i in range(len(real_student))]
        
        syn_students = irl_helper.make_syn_student_personalized(
            trajectories_pass=[trajectories_each_week_pass[week-1]], 
            trajectories_fail=[trajectories_each_week_fail[week-1]], 
            students=[trajectories_each_week[week-1]],
            history_whatif_pass=[history_whatif_pass[week-1]], 
            history_whatif_fail=[history_whatif_fail[week-1]], 
            feature_matrix=add_feature_matrix, 
            world=world, 
            num_week=1, 
            is_start_state_arr=False, 
            test_labels=pred, 
            new_state=True,
            test_ids=test_ids,
            fail_only=fail_only
        )
        
        syn_students = [syn_students[i][0] for i in range(len(syn_students))]
        if len(week_list) > 1:
            results['week'].append(week)
            results['syn_trajectories'].append(syn_students)
            results['real_trajectories'].append(real_student)
            results['world'].append(world)
        else:
            return syn_students, real_student, world
    return results

def get_top_rewards_per_week(df, top_n=3):
    # Ensure 'Week' is a column to group by
    if 'Week' not in df.columns:
        st.error(f"DataFrame is missing 'Week' column: {df.columns}")
        return pd.DataFrame() 
    df_sorted = df.sort_values(by=['Week', 'Reward'], ascending=[True, False])
    top_n_df = df_sorted.groupby('Week').head(top_n).reset_index(drop=True)
    return top_n_df

def show_highest_rewards(course_id='dsp-002'):
    st.header("🏆 Popular Picks by Passing Students by Week")
    st.markdown("This section displays the top 3 videos and quizzes each week that passing students in previous year are most likely to interact with.")
    df_video = pd.read_csv(f'streamlit-assets/{course_id}/sorted_reward_video.csv')
    df_problem = pd.read_csv(f'streamlit-assets/{course_id}/sorted_reward_problem.csv')

    top_3_video = get_top_rewards_per_week(df_video, top_n=3)
    top_3_problem = get_top_rewards_per_week(df_problem, top_n=3)

    tab_video, tab_problem = st.tabs(["Videos", "Quizzes"])

    # merge df and top_3_video on 'chapter'
    top_3_video = top_3_video.merge(st.session_state.metadata[['Topic', 'Skills']], left_on='chapter', right_on='Topic', how='left')
    top_3_problem = top_3_problem.merge(st.session_state.metadata[['Topic', 'Skills']], left_on='chapter', right_on='Topic', how='left')
    
    with tab_video:
        # st.subheader("Top Videos by Week")
        if not top_3_video.empty:
            weeks = sorted(top_3_video['Week'].unique())
            week_tabs = st.tabs([f"Week {w+1}" for w in weeks])

            for week, week_tab in zip(weeks, week_tabs):
                with week_tab:
                    st.markdown(f"**Top Videos - Week {week+1}**")
                    week_data = top_3_video[top_3_video['Week'] == week]
                    week_data['duration'] = week_data['duration'].apply(lambda x: f"{int(x)//60} min {int(x)%60} sec")
                    week_data['title'] = week_data['title'].apply(extract_signal_name_regex)
                    display_cols = ['title', 'Skills', 'subchapter', 'duration']
                    rename_cols = {'title': 'Name', 'Skills': 'Topic', 'subchapter': 'Sub-Topic','duration': 'Duration'}
                    display_data = week_data[display_cols].dropna(axis=1, how='all').rename(columns=rename_cols)
                    st.dataframe(display_data, use_container_width=False, hide_index=True)
       
        else:
            st.info("No video event data available to display top rewards.")

    with tab_problem:
        # st.subheader("Top Quizzes by Week")
        if not top_3_problem.empty:
            weeks = sorted(top_3_problem['Week'].unique())
            week_tabs = st.tabs([f"Week {w+1}" for w in weeks])

            for week, week_tab in zip(weeks, week_tabs):
                with week_tab:
                    st.markdown(f"**Top Quizzes - Week {week+1}**")
                    week_data = top_3_problem[top_3_problem['Week'] == week]
                    week_data['Difficulty'] = week_data['Difficulty'].apply(lambda x: f"{x:.2f}")
                    display_cols = ['title', 'chapter', 'subchapter', 'grade_max', 'Difficulty']
                    rename_cols = {'title': 'Name', 'chapter': 'Topic', 'subchapter': 'Sub-Topic',
                                'grade_max': 'Max Grade', 'Difficulty': 'Difficulty Level (0-1)'}
                    display_data = week_data[display_cols].dropna(axis=1, how='all').rename(columns=rename_cols)
                    st.dataframe(display_data, use_container_width=False, hide_index=True)
                    
        else:
            st.info("No problem event data available to display top rewards.")

class StateManager:
    def __init__(self, **kwargs):
        self._kwargs = kwargs
        week = kwargs.get('week_list', [6])[0] 
        course_id = 'dsp-004'
        DATA_DIR = f'streamlit-assets/{course_id}/'
        model_path = {
            5: 'lstm-bi-32-64-5-1722490972.1859/model.keras_final_e.keras',
            6: 'lstm-bi-32-64-6-1722494926.4949/model.keras_final_e.keras',
            7: 'lstm-bi-32-64-7-1722499225.71723/model.keras_final_e.keras',
            8: 'lstm-bi-32-64-8-1722504182.3553/model.keras_final_e.keras',
            9: 'lstm-bi-32-64-9-1722511435.7777/model.keras_final_e.keras',
            10: 'lstm-bi-32-64-10-1722519098.62673/model.keras_final_e.keras',
        }
        self.reconstructed_model = tf.keras.models.load_model(f'checkpoints/{model_path[week]}')
        # self.test_ids = np.load(f'{DATA_DIR}/test_students_5.npy')
        self.test_ids = None
        self.x_test = np.load(f'{DATA_DIR}/real-data-early-prediction_{course_id}_1to10_ver2.npy_features.npy')
        print(f"Loaded x_test shape: {self.x_test.shape}")
   
def show_metadata():
    st.subheader("Course Information")
    st.markdown("""On the left is the skill set for each topic; on the right is a visual representation of the prerequisite skill structure.
    The taught skills are divided into three categories:
    
    1. Core skills: Fundamental skills that are essential for understanding the course material.
    2. Applied skills: Skills that apply the core skills in practical scenarios.
    3. Theory-based extension: Advanced skills that build upon the core skills.
    
    """)

    col1, col2 = st.columns([1, 2]) 
    df = pd.read_csv('streamlit-assets/dsp_prerequisite_skills_2.csv')
    st.session_state.metadata = df 
    with col1:
        st.subheader("📋 Topic Table")
        st.dataframe(df[['Topic', 'Skills', 'No. Videos', 'No. Quizzes']], use_container_width=False, hide_index=True)

    with col2:
        st.subheader("Prerequisite skill structure")
        fig = Image.open('streamlit-assets/prerequisite_skills_structure.png')
        st.image(fig, caption='**Blue**: Core skills, **Pink**: Applied skills, **Purple**: Theory-based extension', use_container_width=False)
    
   
if 'current_page' not in st.session_state:
    st.session_state.current_page = 1

def show_navigation():
    """Display navigation buttons"""
    col1, col2, col3 = st.columns(3)
    
    with col1:
        if st.button("Page 1: Course Description", key="nav1"):
            st.session_state.current_page = 1
    
    with col2:
        if st.button("Page 2: Course Analysis", key="nav2"):
            st.session_state.current_page = 2
    
    with col3:
        if st.button("Page 3: What-If Simulation", key="nav3"):
            st.session_state.current_page = 3

def get_reward(history_whatif, whatif_world, feature_matrix):
    transition_probability = whatif_world.p_transition
    for i in range(len(history_whatif)):
        alpha = history_whatif[i]['alpha']
      
        reward = feature_matrix @ alpha
       
        not_valid_state_action = transition_probability.sum(axis=-1) <= 0
        reward[not_valid_state_action] = np.nan 
        reward = (reward - np.nanmean(reward.flatten())) / np.nanstd(reward.flatten())
        reward[not_valid_state_action] = -np.inf 
      
        history_whatif[i]['reward'] = reward
    return history_whatif

def load_data():
    """Load all required data (only once)"""
    if 'data_loaded' not in st.session_state:
        # try:
            course_id = 'dsp-004'
            # Load trajectory data
            with open(f'streamlit-assets/{course_id}/trajectories_each_week.pkl', 'rb') as f:
                st.session_state.trajectories_each_week = pickle.load(f)
            with open(f'streamlit-assets/{course_id}/trajectories_each_week_pass.pkl', 'rb') as f:
                st.session_state.trajectories_each_week_pass = pickle.load(f)
            with open(f'streamlit-assets/{course_id}/trajectories_each_week_fail.pkl', 'rb') as f:
                st.session_state.trajectories_each_week_fail = pickle.load(f)
            with open(f'streamlit-assets/{course_id}/trajectories.pkl', 'rb') as f:
                st.session_state.trajectories = pickle.load(f)
            with open('history_whatif_fail_300.pkl', 'rb') as f:
                history_whatif_fail = pickle.load(f)
            with open('history_whatif_pass_300.pkl', 'rb') as f:
                history_whatif_pass = pickle.load(f)
            
            # Load course data
            # st.session_state.combinedg = pd.read_csv(f'data/{course_id}/combinedg_features_{course_id}_processed.csv')
            
            # Load dictionaries and mappings
            with open(f'streamlit-assets/{course_id}/dict_event.pkl', 'rb') as f:
                st.session_state.dict_event = pickle.load(f)
            with open(f'streamlit-assets/{course_id}/dict_action.pkl', 'rb') as f:
                st.session_state.dict_action = pickle.load(f)
            # with open(f'data/{course_id}/map_week.pkl', 'rb') as f:
            #     st.session_state.map_week = pickle.load(f)
            # with open(f'data/{course_id}/map_event_id.pkl', 'rb') as f:
            #     st.session_state.map_event_id = pickle.load(f)
            # with open(f'data/{course_id}/map_action.pkl', 'rb') as f:
            #     st.session_state.map_action = pickle.load(f)
            with open(f'streamlit-assets/{course_id}/problem_event.pkl', 'rb') as f:
                st.session_state.problem_event = pickle.load(f)
            with open(f'streamlit-assets/{course_id}/video_event.pkl', 'rb') as f:
                st.session_state.video_event = pickle.load(f)
            pred = pd.read_csv(f'streamlit-assets/{course_id}/predictions_weeks_6.csv')
            st.session_state.pred = pred['y_pred']
            # Load schedule and other data
            # st.session_state.schedule = pd.read_csv(f'{datadir}/schedule/{course_id}.csv')
            # st.session_state.values_df, st.session_state.duration = whatif_values(st.session_state.combinedg, 
            #                                         st.session_state.schedule, 
            #                                        map_event_id=st.session_state.map_event_id, 
            #                                        problem_event=st.session_state.problem_event)
            # dump to numpy file and load when needed
            # st.session_state.values_df.to_csv(f'streamlit-assets/value_df_duration_{course_id}.csv')
            st.session_state.values_df = pd.read_csv(f'streamlit-assets/value_df_duration_{course_id}.csv')
            st.session_state.duration = np.load(f'streamlit-assets/duration_{course_id}.npy')
            # print(f"Loaded values shape: {st.session_state.values_df}")
            
            # st.session_state.ids = np.load(f'results/whatif/{course_id}//test_students_5.npy')
            st.session_state.ids = None
            st.session_state.world = Env.ClickstreamWorld(trajectories=st.session_state.trajectories,
                                                        dict_action=st.session_state.dict_action, 
                                                        dict_event=st.session_state.dict_event,
                                                        video_arr=st.session_state.video_event,
                                                        problem_arr=st.session_state.problem_event,
                                                        values=st.session_state.values_df,
                                                        add_state=True)
            # print(len(st.session_state.world.values))
            feature_matrix = st.session_state.world.designed_features(st.session_state.world.values)
            st.session_state.history_whatif_fail = get_reward(history_whatif_fail, st.session_state.world, feature_matrix=feature_matrix)
            st.session_state.history_whatif_pass = get_reward(history_whatif_pass, st.session_state.world, feature_matrix=feature_matrix)
            st.session_state.state_manager = StateManager(world=st.session_state.world, 
                                                        trajectories=st.session_state.trajectories, 
                                                        test_ids=st.session_state.ids,
                                                        trajectories_each_week=st.session_state.trajectories_each_week,
                                                        trajectories_each_week_pass=st.session_state.trajectories_each_week_pass,
                                                        trajectories_each_week_fail=st.session_state.trajectories_each_week_fail,
                                                        history_whatif_pass=st.session_state.history_whatif_pass,
                                                        history_whatif_fail=st.session_state.history_whatif_fail,
                                                        week_list=[6])
            
            st.session_state.data_loaded = True
            

def page_1_description():
    st.title("What-if Classroom for Course Design: Digital Signal Processing (DSP)")
    
    st.markdown(
        """
        <h4>Description:</h4>
        You are teaching a 10-week Digital Signal Processing (DSP) course on a MOOC platform and have just completed the first 5 weeks. 
        You're now considering whether to add new content in week 6.
        
        The platform provides access to detailed clickstream data that tracks how students in previous years interacted with videos and quizzes.
        You want to identify which content is most effective in supporting student success. 
        
        This demo walks you through how to use our simulation model to preview the impact of adding new materials before implementing them in reality.
        """,
        unsafe_allow_html=True
    )
    
    st.markdown("---")
    
    st.header("Course Overview")
    st.markdown("""
    **Course:** Digital Signal Processing (DSP)  
    **Platform:** MOOC (Massive Open Online Course)  
    **Duration:** 10 weeks  
    **Current Status:** Completed first 5 weeks  
    """)
    
    try:
        show_metadata()
    except:
        st.info("Metadata display function not available in this demo.")
    st.markdown("---")
    col1, col2, col3 = st.columns([1, 1, 1])
    with col3:
        if st.button("Next: Course Analysis →", key="next_page_1", type="primary"):
            st.session_state.current_page = 2
            st.rerun()

def page_2_analysis():
    st.title("Course Material Analysis")
    
    if st.session_state.get('data_loaded', False):
        # st.warning("Loading course data...")
        load_data()
        if not st.session_state.get('data_loaded', False):
            st.error("Failed to load data. Please check data files.")
    
    st.header("Top Performing Course Materials")
    
    try:
        show_highest_rewards()
    except:
        st.info("Highest rewards analysis not available in this demo.")
    
    st.header("Current Student Performance Prediction")
    plot_performance(show_image=True)
    # col1, col2 = st.columns([1, 2])
    
    # with col1:
    #     try:
            
    #     except:
    #         st.info("Performance visualization not available.")
    # print(st.session_state.get('data_loaded'))
    # with col2:
    #     if st.session_state.get('data_loaded', True):
    #         try:
    #             fail_num, pass_num, _ = performance_prediction(
    #                 st.session_state.state_manager.reconstructed_model, 
    #                 st.session_state.state_manager.x_test
    #             )
                
    #             st.session_state.real_pred_df = pd.DataFrame({
    #                 'class': ['Predicted Performance Before Intervention'],
    #                 'fail_num': [fail_num],
    #                 'pass_num': [pass_num],
    #                 'fail_percent': [fail_num / (fail_num + pass_num) * 100],
    #                 'pass_percent': [pass_num / (fail_num + pass_num) * 100]
    #             })
                
    #             st.subheader("Performance Metrics")
    #             plot_performance(df=st.session_state.real_pred_df, 
    #                              show_performance=True, col_performance=col2)
                
    #         except Exception as e:
    #             print(f"Error calculating performance prediction: {str(e)}")
    #     else:
    #        st.info("Data not loaded. Cannot display performance metrics.")
    
    st.header("Week 6 Intervention Opportunity")
    st.markdown("""
    Based on the analysis of the first 5 weeks, we can now simulate the impact of introducing 
    new content in week 6. The next page allows you to experiment with different interventions 
    and see their predicted effects on student performance.
    """)
    st.markdown("---")
    col1, col2, col3 = st.columns([1, 1, 1])
    with col1:
        if st.button("← Previous: Course Description", key="prev_page_2"):
            st.session_state.current_page = 1
            st.rerun()
    with col3:
        if st.button("Next: What-If Simulation →", key="next_page_2", type="primary"):
            st.session_state.current_page = 3
            st.rerun()

def page_3_simulation():
    show_navigation()
    st.title("What-If Simulation: Week 6 Interventions")
    st.header("Simulation Capabilities")
    st.markdown("""
    Our simulation model allows you to:
    - Preview the impact of new content before implementation
    - Test different content types (videos, quizzes) and difficulties
    - Target interventions for specific student populations
    - Analyze engagement patterns and predict performance outcomes
    """)
    if not st.session_state.get('data_loaded', False):
        st.warning("Loading course data...")
        load_data()
        if not st.session_state.get('data_loaded', False):
            st.error("Failed to load data. Please check data files.")
            return
    
    # Sidebar controls
    st.header("Introduce New Course Content")
    st.markdown("You are teaching a 10-week Digital Signal Processing (DSP) course on a MOOC platform and have just completed the first 5 weeks. You're now considering whether to add new content in week 6.\
                    You want to identify which content is most effective in supporting student success by trying different interventions.")
    
    # chapter = st.sidebar.slider('Topic:', 1, 10, 1)
    st.write("**Step 1: Please select the topic the new content covers**")

    selected_chapter_name = st.selectbox(
        'Topic:',
        options=st.session_state.metadata['Skills'], # What the user sees in the dropdown
        index=0 # Default selected option (index of 'Introduction to DSP')
    )
    chapter = st.session_state.metadata[st.session_state.metadata['Skills'] == selected_chapter_name]['Topic'].values[0]
   
    st.write("**Step 2: Please select the type of course material you want to add, either a quiz or a video**")
    event_type = st.selectbox('Course Material Type:', ['quiz', 'video'])
    max_duration = st.session_state.duration[0]//60 if 'duration' in st.session_state else 0
    min_duration = st.session_state.duration[1]//60 if 'duration' in st.session_state else 60
    st.write("**Step 3: Please select the difficulty of the quiz or the duration of the video.**\
                     \nDifficulty: value between 0 and 1; easiest 0, hardest: 1.\
                     \nDuration: in minutes; shortest: 2-minute, longest: 25-minute.")
    if event_type == 'quiz':
        value_label = 'Difficulty:'
        value = st.slider(value_label, 0.0, 1.0, 0.5, 0.01)
    else:
        value_label = 'Duration (minutes):'
        value = st.slider(value_label, min_duration, max_duration, min_duration, 0.5)  # Duration in minutes
    
    st.write("**Step 4: Please select whether you want to focus on low-performed students only.**")
    fail_only = st.selectbox('Simulate low-performed students only:', 
                                   [False, True], 
                                   format_func=lambda x: 'Yes' if x else 'No')
    
    # Main content area
    st.subheader("Intervention Settings: ready to run simulation")
    col1, col2 = st.columns([1, 1])
    with col1:
        st.write(f"**Topic:** {chapter}")
        st.write(f"**Content Type:** {event_type.title()}")
    with col2:
        st.write(f"**{value_label}** {value:.2f}")
        st.write(f"**Target Low-Performers Only:** {'Yes' if fail_only else 'No'}")
    
    if st.button('Run Simulation', key='add_state_button'):
        st.session_state.world = Env.ClickstreamWorld(trajectories=st.session_state.trajectories,
                                                        dict_action=st.session_state.dict_action, 
                                                        dict_event=st.session_state.dict_event,
                                                        video_arr=st.session_state.video_event,
                                                        problem_arr=st.session_state.problem_event,
                                                        values=st.session_state.values_df,
                                                        add_state=True)
        event_data = {
            'chapter': chapter,
            'event_type': event_type,
            'difficulty': value,
            'duration':(value - min_duration)/(max_duration - min_duration),
            'fail_only': fail_only
        }
       #  print(len(st.session_state.state_manager.world.values), len(st.session_state.values_df))
        with st.spinner("Analyzing trajectories and simulating intervention..."):
            try:
                results = add_new_state(world=st.session_state.world, 
                                        event_data=event_data, 
                                        week_list=[6], 
                                        test_ids=st.session_state.state_manager.test_ids,
                                        trajectories=st.session_state.trajectories,
                                        trajectories_each_week=st.session_state.trajectories_each_week, 
                                        trajectories_each_week_pass=st.session_state.trajectories_each_week_pass, 
                                        trajectories_each_week_fail=st.session_state.trajectories_each_week_fail,
                                        history_whatif_pass=st.session_state.history_whatif_pass, 
                                        history_whatif_fail=st.session_state.history_whatif_fail,
                                        fail_only=fail_only)
                
                st.session_state.results = results
                st.success("Simulation completed! ")
            except Exception as e:
                st.error(f"Simulation failed: {str(e)}")
    
    # Display results if they exist
    if 'results' in st.session_state:
        st.header("Simulation Results")
        
        try:
            
            import utils.data_helper as data_helper
            if 'real_pred_df' not in st.session_state:
                fail_num, pass_num, _ = performance_prediction(
                    st.session_state.state_manager.reconstructed_model, 
                    st.session_state.state_manager.x_test
                )
                
                st.session_state.real_pred_df = pd.DataFrame({
                    'class': ['Predicted Performance Before Intervention'],
                    'fail_num': [fail_num],
                    'pass_num': [pass_num],
                    'fail_percent': [fail_num / (fail_num + pass_num) * 100],
                    'pass_percent': [pass_num / (fail_num + pass_num) * 100]
                })
    
            if isinstance(st.session_state.results, dict):
                syn = st.session_state.results['syn_trajectories']
                real = st.session_state.results['real_trajectories']
                world = st.session_state.results['world']
            else:
                syn, real, world = st.session_state.results
                
            course_features = data_helper.trajectories_to_features(
                syn, world, num_week=1, syn=True, save_to_disk=False
            )
            if not fail_only:
                x_test = np.concatenate((st.session_state.state_manager.x_test[:, :(st.session_state.state_manager.x_test.shape[0]-100), :], 
                                    course_features), axis=1)
            else:
                x_test = np.concatenate((st.session_state.state_manager.x_test[st.session_state.pred == 1, :(st.session_state.state_manager.x_test.shape[0]-100), :], 
                                    course_features), axis=1)
            
            fail_num, pass_num, _ = performance_prediction(
                st.session_state.state_manager.reconstructed_model, x_test
            )
            
            temp_df = pd.DataFrame({
                'class': ['Predicted Performance After Intervention'],
                'fail_num': [fail_num],
                'pass_num': [pass_num],
                'fail_percent': [fail_num / (fail_num + pass_num) * 100],
                'pass_percent': [pass_num / (fail_num + pass_num) * 100]
            })
            if not fail_only:
                merged_df = pd.concat([temp_df, st.session_state.real_pred_df], ignore_index=True)
                st.session_state.performance_df = merged_df
                st.session_state.num_students = temp_df.iloc[0]['fail_num'] + temp_df.iloc[0]['pass_num']
            else:
                st.session_state.fail_df = pd.DataFrame({
                'class': ['Predicted Performance After Intervention'],
                'fail_num': st.session_state.real_pred_df.fail_num,
                'pass_num': [0],
                'fail_percent': 100,
                'pass_percent': 0
                })
                merged_df = pd.concat([temp_df, st.session_state.fail_df], ignore_index=True)
                st.session_state.performance_df = merged_df
            
            # Performance comparison
            st.subheader("Performance Prediction: Before and After Intervention")
            col1, col2 = st.columns(2)
            
            with col1:
                
                st.metric("Before Intervention - Pass Rate", 
                         f"{st.session_state.performance_df.iloc[-1]['pass_percent']:.1f}%")
                st.metric("Before Intervention - Students Passing", 
                         f"{int(st.session_state.performance_df.iloc[-1]['pass_num'])} / {st.session_state.num_students} Students")

            with col2:
                new_pass_rate = temp_df.iloc[0]['pass_percent']
                old_pass_rate = st.session_state.performance_df.iloc[-1]['pass_percent']
                improvement = new_pass_rate - old_pass_rate
                
                st.metric("After Intervention - Pass Rate", 
                         f"{new_pass_rate:.1f}%", 
                         f"{improvement:+.1f}%")
                st.metric("After Intervention - Students Passing", 
                         f"{int(temp_df.iloc[0]['pass_num'])} / {st.session_state.num_students} Students",
                         int(temp_df.iloc[0]['pass_num'] - st.session_state.performance_df.iloc[-1]['pass_num']))
            
            # Visualization
            try:
                plot_performance(st.session_state.performance_df, show_performance=True)
            except:
                st.info("Performance visualization not available.")
            
            try:
                compare_chapter_engagement(syn, real, world)
            except Exception as e:
                raise e
                st.info(f"Engagement analysis not available: {str(e)}")
                
        except Exception as e:
            raise e
            st.error(f"Error processing simulation results: {str(e)}")
    
    # elif 'real_pred_df' not in st.session_state:
        
    #     st.info("Please visit Page 2 first to load baseline performance predictions.")
    else:
        st.info("Run a simulation using the controls to see results.")
    st.markdown("---")
    col1, col2, col3 = st.columns([1, 1, 1])
    with col1:
        if st.button("← Previous: Course Analysis", key="prev_page_3"):
            st.session_state.current_page = 2
            st.rerun()

def main():
    # Load data once at startup
    load_data()
    
    # Navigation
    # show_navigation()
    # st.markdown("---")
    
    # Display current page
    if st.session_state.current_page == 1:
        page_1_description()
    elif st.session_state.current_page == 2:
        page_2_analysis()
    elif st.session_state.current_page == 3:
        page_3_simulation()
    
    # Page indicator
    st.markdown("---")
    st.markdown(f"**Current Page: {st.session_state.current_page} of 3**")

if __name__ == "__main__":
    main()