# CHARTS + DOWNLOAD + NO NAMES
# intervention_analysis_app.py

#------------------------------------------------------------------------
# Import Modules
#------------------------------------------------------------------------
import streamlit as st
import pandas as pd
import matplotlib.pyplot as plt
import io
import re
# from transformers import pipeline
from huggingface_hub import InferenceClient
import os
from pathlib import Path
from dotenv import load_dotenv

load_dotenv()

#------------------------------------------------------------------------
# Configurations
#------------------------------------------------------------------------
# Streamlit page setup
st.set_page_config(
    page_title="Intervention Program Analysis", 
    page_icon=":bar_chart:", 
    layout="centered", 
    initial_sidebar_state="auto",
    menu_items={
        'Get Help': 'mailto:clevesseur@mimtss.org',
        'About': "This app is built to support spreadsheet analysis"
    }
)

#------------------------------------------------------------------------
# Sidebar
#------------------------------------------------------------------------
# logo
main_body_logo = "mimtss.png"
sidebar_logo = "mimtss_small.png"
st.logo(sidebar_logo, icon_image=main_body_logo)

with st.sidebar:
    # Password input field
    # password = st.text_input("Enter Password:", type="password")
    
    # Set the desired width in pixels
    image_width = 300  
    # Define the path to the image
    image_path = "mimtss.png"
    # Display the image
    st.image(image_path, width=image_width)
    
    # Toggle for Help and Report a Bug
    with st.expander("Need help and report a bug"):
        st.write("""
        **Contact**: Cheyne LeVesseur, PhD  
        **Email**: clevesseur@mimtss.org
        """)
    st.divider()
    st.subheader('User Instructions')
    
    # Principles text with Markdown formatting
    User_Instructions = """

    - **Step 1**: Upload your Excel file.
    - **Step 2**: Anonymization – student names are replaced with initials for privacy.
    - **Step 3**: Review anonymized data.
    - **Step 4**: View **intervention session statistics**.
    - **Step 5**: Review **student attendance and engagement metrics**.
    - **Step 6**: Review AI-generated **insights and recommendations**.

    ### **Privacy Assurance**
    - **No full names** are ever displayed or sent to the AI model—only initials are used.
    - This ensures that sensitive data remains protected throughout the entire process.

    ### **Detailed Instructions**

    #### **1. Upload Your Excel File**
    - Start by uploading an Excel file that contains intervention data. 
    - Click on the **“Upload your Excel file”** button and select your `.xlsx` file from your computer.

    **Note**: Your file should have columns like "Did the intervention happen today?" and "Student Attendance [FirstName LastName]" for the analysis to work correctly.

    #### **2. Automated Name Anonymization**
    - Once the file is uploaded, the app will **automatically replace student names with initials** in the "Student Attendance" columns.
      - For example, **"Student Attendance [Cheyne LeVesseur]"** will be displayed as **"Student Attendance [CL]"**.
      - If the student only has a first name, like **"Student Attendance [Cheyne]"**, it will be displayed as **"Student Attendance [C]"**.
    - This anonymization helps to **protect student privacy**, ensuring that full names are not visible or sent to the AI language model.

    #### **3. Review the Uploaded Data**
    - You will see the entire table of anonymized data to verify that the information has been uploaded correctly and that names have been replaced with initials.

    #### **4. Intervention Session Statistics**
    - The app will calculate and display statistics related to intervention sessions, such as:
      - **Total Number of Days Available**
      - **Intervention Sessions Held**
      - **Intervention Sessions Not Held**
      - **Intervention Frequency (%)**
    - A **stacked bar chart** will be shown to visualize the number of sessions held versus not held.
    - If you need to save the visualization, click the **“Download Chart”** button to download it as a `.png` file.

    #### **5. Student Metrics Analysis**
    - The app will also calculate metrics for each student:
      - **Attendance (%)** – The percentage of intervention sessions attended.
      - **Engagement (%)** – The level of engagement during attended sessions.
    - These metrics will be presented in a **line graph** that shows attendance and engagement for each student.
    - You can click the **“Download Chart”** button to download the visualization as a `.png` file.

    #### **6. Generate AI Analysis and Recommendations**
    - The app will prepare data from the student metrics to provide notes, key takeaways, and suggestions for improving outcomes using an **AI language model**.
    - You will see a **spinner** labeled **“Generating AI analysis…”** while the AI processes the data.
      - This step may take a little longer, but the spinner ensures you know that the system is working.
    - Once the analysis is complete, the AI's recommendations will be displayed under **"AI Analysis"**.
    - You can click the **“Download LLM Output”** button to download the AI-generated recommendations as a `.txt` file for future reference.

    """
    st.markdown(User_Instructions)

#------------------------------------------------------------------------
# Functions
#------------------------------------------------------------------------
# Set the Hugging Face API key
# Retrieve Hugging Face API key from environment variables
hf_api_key = os.getenv('HF_API_KEY')
if not hf_api_key:
    raise ValueError("HF_API_KEY not set in environment variables")

# Create the Hugging Face inference client
client = InferenceClient(api_key=hf_api_key)

# Constants
INTERVENTION_COLUMN = 'Did the intervention happen today?'
ENGAGED_STR = 'Engaged (Respect, Responsibility, Effort)'
PARTIALLY_ENGAGED_STR = 'Partially Engaged (about 50%)'
NOT_ENGAGED_STR = 'Not Engaged (less than 50%)'

def safe_convert_to_time(series, format_str='%I:%M %p'):
    try:
        converted = pd.to_datetime(series, format='%H:%M:%S', errors='coerce')
        if format_str:
            return converted.dt.strftime(format_str)
        return converted
    except Exception as e:
        print(f"Error converting series to time: {e}")
        return series

def safe_convert_to_datetime(series, format_str=None):
    try:
        # Attempt to convert to datetime, ignoring errors
        converted = pd.to_datetime(series, errors='coerce')
        if format_str:
            # Format if a format string is provided
            return converted.dt.strftime(format_str)
        return converted
    except Exception as e:
        print(f"Error converting series to datetime: {e}")
        return series

def format_session_data(df):
    # Format "Date of Session" and "Timestamp" columns with safe conversion
    df['Date of Session'] = safe_convert_to_datetime(df['Date of Session'], '%m/%d/%Y')
    df['Timestamp'] = safe_convert_to_datetime(df['Timestamp'], '%I:%M %p')
    df['Session Start Time'] = safe_convert_to_time(df['Session Start Time'], '%I:%M %p')
    df['Session End Time'] = safe_convert_to_time(df['Session End Time'], '%I:%M %p')

    # Reorder columns
    df = df[['Date of Session', 'Timestamp'] + [col for col in df.columns if col not in ['Date of Session', 'Timestamp']]]

    return df

def main():
    st.title("Intervention Program Analysis")

    # File uploader
    uploaded_file = st.file_uploader("Upload your Excel file", type=["xlsx"])

    if uploaded_file is not None:
        try:
            # Read the Excel file into a DataFrame
            df = pd.read_excel(uploaded_file)

            # Format the session data
            df = format_session_data(df)

            # Replace student names with initials
            df = replace_student_names_with_initials(df)

            st.subheader("Uploaded Data")
            st.write(df)

            # Ensure expected column is available
            if INTERVENTION_COLUMN not in df.columns:
                st.error(f"Expected column '{INTERVENTION_COLUMN}' not found.")
                return

            # Clean up column names
            df.columns = df.columns.str.strip()

            # Compute Intervention Session Statistics
            intervention_stats = compute_intervention_statistics(df)
            st.subheader("Intervention Session Statistics")
            st.write(intervention_stats)

            # Two-column layout for the visualization and intervention frequency
            col1, col2 = st.columns([3, 1])  # Set the column width ratio
            
            with col1:
                intervention_fig = plot_intervention_statistics(intervention_stats)
            
            with col2:
                intervention_frequency = intervention_stats['Intervention Frequency (%)'].values[0]
                # Display the "Intervention Frequency (%)" text
                st.markdown("<h3 style='color: #358E66;'>Intervention Frequency</h3>", unsafe_allow_html=True)
                # Display the frequency value below it
                st.markdown(f"<h1 style='color: #358E66;'>{intervention_frequency}%</h1>", unsafe_allow_html=True)

            # Add download button for Intervention Session Statistics chart
            download_chart(intervention_fig, "intervention_statistics_chart.png")

            # Compute Student Metrics
            student_metrics_df = compute_student_metrics(df)
            st.subheader("Student Metrics")
            st.write(student_metrics_df)

            # Compute Student Metric Averages
            attendance_avg_stats, engagement_avg_stats = compute_average_metrics(student_metrics_df)

            # Visualization for Student Metrics
            student_metrics_fig = plot_student_metrics(student_metrics_df, attendance_avg_stats, engagement_avg_stats)

            # # Two-column layout for the visualization and intervention frequency
            # col1, col2 = st.columns([3, 1])  # Set the column width ratio
            
            # with col1:
            #     student_metrics_fig = plot_student_metrics(student_metrics_df, attendance_avg_stats, engagement_avg_stats)
                      
            # with col2:
            #     # Display the "Attendance Average (%)" text and value
            #     st.markdown("<h3 style='color: #358E66;'>Attendance Average (%)</h3>", unsafe_allow_html=True)
            #     if attendance_avg_stats is not None:
            #         st.markdown(f"<h2 style='color: #358E66;'>{attendance_avg_stats}%</h2>", unsafe_allow_html=True)
            #     else:
            #         st.markdown("<h2 style='color: #358E66;'>N/A</h2>", unsafe_allow_html=True)
                    
            #     # Display the "Engagement Average (%)" text and value
            #     st.markdown("<h3 style='color: #358E66;'>Engagement Average (%)</h3>", unsafe_allow_html=True)
            #     if engagement_avg_stats is not None:
            #         st.markdown(f"<h2 style='color: #358E66; margin-top: 0px; margin-bottom: 0px;'>{engagement_avg_stats}%</h2>", unsafe_allow_html=True)
            #     else:
            #         st.markdown("<h2 style='color: #358E66;'>N/A</h2>", unsafe_allow_html=True)

            # Add download button for Student Metrics chart
            download_chart(student_metrics_fig, "student_metrics_chart.png")

            # Prepare input for the language model
            llm_input = prepare_llm_input(student_metrics_df)

            # Generate Notes and Recommendations using Hugging Face LLM
            with st.spinner("Generating AI analysis..."):
                recommendations = prompt_response_from_hf_llm(llm_input)

            st.subheader("AI Analysis")
            st.markdown(recommendations)

            # Add download button for LLM output
            download_llm_output(recommendations, "llm_output.txt")

        except Exception as e:
            st.error(f"Error reading the file: {str(e)}")

def replace_student_names_with_initials(df):
    """Replace student names in column headers with initials."""
    updated_columns = []
    for col in df.columns:
        if col.startswith('Student Attendance'):
            # Extract the name from the column header
            match = re.match(r'Student Attendance \[(.+?)\]', col)
            if match:
                name = match.group(1)
                # Split the name into parts (first and last name)
                name_parts = name.split()
                # Convert the name to initials
                if len(name_parts) == 1:
                    initials = name_parts[0][0]  # Just take the first letter
                else:
                    initials = ''.join([part[0] for part in name_parts])  # Take the first letter of each part
                # Update the column name
                updated_columns.append(f'Student Attendance [{initials}]')
            else:
                updated_columns.append(col)
        else:
            updated_columns.append(col)
    df.columns = updated_columns
    return df

def compute_intervention_statistics(df):
    # Total Number of Days Available
    total_days = len(df)

    # Intervention Sessions Held
    sessions_held = df[INTERVENTION_COLUMN].str.strip().str.lower().eq('yes').sum()

    # Intervention Sessions Not Held
    sessions_not_held = df[INTERVENTION_COLUMN].str.strip().str.lower().eq('no').sum()

    # Intervention Frequency (%)
    intervention_frequency = (sessions_held / total_days) * 100 if total_days > 0 else 0
    intervention_frequency = round(intervention_frequency, 0)

    # Reorder columns as specified
    stats = {
        'Intervention Frequency (%)': [intervention_frequency],
        'Intervention Sessions Held': [sessions_held],
        'Intervention Sessions Not Held': [sessions_not_held],
        'Total Number of Days Available': [total_days]
    }
    stats_df = pd.DataFrame(stats)
    return stats_df

def plot_intervention_statistics(intervention_stats):
    # Create a stacked bar chart for sessions held and not held
    sessions_held = intervention_stats['Intervention Sessions Held'].values[0]
    sessions_not_held = intervention_stats['Intervention Sessions Not Held'].values[0]

    fig, ax = plt.subplots()
    # Plot "Held" on the bottom
    ax.bar(['Intervention Sessions'], [sessions_held], label='Held', color='#358E66')
    # Plot "Not Held" on top of "Held"
    ax.bar(['Intervention Sessions'], [sessions_not_held], bottom=[sessions_held], label='Not Held', color='#91D6B8')

    # Display values on the bars
    ax.text(0, sessions_held / 2, str(sessions_held), ha='center', va='center', color='white', 
            fontweight='bold', fontsize=14)
    ax.text(0, sessions_held + sessions_not_held / 2, str(sessions_not_held), ha='center', va='center', color='black', 
            fontweight='bold', fontsize=14)

    # Update chart settings
    ax.set_ylabel('Frequency')
    ax.set_title('Intervention Sessions Held vs Not Held', fontsize=16)  # Optional: Increased title font size
    # Reverse the legend order to match the new stacking order
    handles, labels = ax.get_legend_handles_labels()
    ax.legend(handles[::-1], labels[::-1])

    # Hide the top and right spines
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    
    st.pyplot(fig)

    return fig

def compute_student_metrics(df):
    # Filter DataFrame for sessions where intervention happened
    intervention_df = df[df[INTERVENTION_COLUMN].str.strip().str.lower() == 'yes']
    intervention_sessions_held = len(intervention_df)

    # Get list of student columns
    student_columns = [col for col in df.columns if col.startswith('Student Attendance')]

    student_metrics = {}

    for col in student_columns:
        student_name = col.replace('Student Attendance [', '').replace(']', '').strip()
        # Get the attendance data for the student
        student_data = intervention_df[[col]].copy()

        # Treat blank entries as 'Absent'
        student_data[col] = student_data[col].fillna('Absent')

        # Assign attendance values
        attendance_values = student_data[col].apply(lambda x: 1 if x in [
            ENGAGED_STR,
            PARTIALLY_ENGAGED_STR,
            NOT_ENGAGED_STR
        ] else 0)

        # Number of Sessions Attended
        sessions_attended = attendance_values.sum()

        # Attendance (%)
        attendance_pct = (sessions_attended / intervention_sessions_held) * 100 if intervention_sessions_held > 0 else 0
        attendance_pct = round(attendance_pct)  # Round to whole number

        # Calculate the number of students in each engagement category
        engagement_counts = {
            'Engaged': 0,
            'Partially Engaged': 0,
            'Not Engaged': 0,
            'Absent': 0
        }

        # Count the engagement states
        for x in student_data[col]:
            if x == ENGAGED_STR:
                engagement_counts['Engaged'] += 1
            elif x == PARTIALLY_ENGAGED_STR:
                engagement_counts['Partially Engaged'] += 1
            elif x == NOT_ENGAGED_STR:
                engagement_counts['Not Engaged'] += 1
            else:
                engagement_counts['Absent'] += 1  # Count as Absent if not engaged

        # Calculate percentages for engagement states
        total_sessions = sum(engagement_counts.values())
        
        # Engagement (%)
        engagement_pct = (engagement_counts['Engaged'] / total_sessions * 100) if total_sessions > 0 else 0
        engagement_pct = round(engagement_pct)

        engaged_pct = (engagement_counts['Engaged'] / total_sessions * 100) if total_sessions > 0 else 0
        engaged_pct = round(engaged_pct)

        partially_engaged_pct = (engagement_counts['Partially Engaged'] / total_sessions * 100) if total_sessions > 0 else 0
        partially_engaged_pct = round(partially_engaged_pct)

        not_engaged_pct = (engagement_counts['Not Engaged'] / total_sessions * 100) if total_sessions > 0 else 0
        not_engaged_pct = round(not_engaged_pct)

        absent_pct = (engagement_counts['Absent'] / total_sessions * 100) if total_sessions > 0 else 0
        absent_pct = round(absent_pct)

        # Store metrics in the required order
        student_metrics[student_name] = {
            'Attendance (%)': attendance_pct,
            'Attendance #': sessions_attended,  # Raw number of sessions attended
            'Engagement (%)': engagement_pct,
            'Engaged (%)': engaged_pct,
            'Partially Engaged (%)': partially_engaged_pct,
            'Not Engaged (%)': not_engaged_pct,
            'Absent (%)': absent_pct
        }

    # Create a DataFrame from student_metrics
    student_metrics_df = pd.DataFrame.from_dict(student_metrics, orient='index').reset_index()
    student_metrics_df.rename(columns={'index': 'Student'}, inplace=True)
    return student_metrics_df

def compute_average_metrics(student_metrics_df):
    # Calculate the attendance and engagement average percentages across students
    attendance_avg_stats = student_metrics_df['Attendance (%)'].mean()  # Calculate the average attendance percentage
    engagement_avg_stats = student_metrics_df['Engagement (%)'].mean()  # Calculate the average engagement percentage
    
    # Round the averages to make them whole numbers
    attendance_avg_stats = round(attendance_avg_stats)
    engagement_avg_stats = round(engagement_avg_stats)
    
    return attendance_avg_stats, engagement_avg_stats

def plot_student_metrics(student_metrics_df, attendance_avg_stats, engagement_avg_stats):
    # Create the figure and axis
    fig, ax = plt.subplots(figsize=(10, 6))  # Increased figure size for better readability

    # Width for the bars
    bar_width = 0.35  # Width of the bars

    index = range(len(student_metrics_df))  # Index for each student

    # Plot Attendance and Engagement bars side by side
    attendance_bars = ax.bar([i - bar_width / 2 for i in index], 
                              student_metrics_df['Attendance (%)'], 
                              width=bar_width, label='Attendance (%)', 
                              color='#005288', alpha=0.7)
    engagement_bars = ax.bar([i + bar_width / 2 for i in index], 
                              student_metrics_df['Engagement (%)'], 
                              width=bar_width, label='Engagement (%)', 
                              color='#3AB0FF', alpha=0.7)

    # Add labels to each bar
    for bar in attendance_bars:
        height = bar.get_height()
        ax.text(bar.get_x() + bar.get_width() / 2, height, 
                f'{height:.0f}%', ha='center', va='bottom', color='black')  # No decimal for integer percentage

    for bar in engagement_bars:
        height = bar.get_height()
        ax.text(bar.get_x() + bar.get_width() / 2, height, 
                f'{height:.0f}%', ha='center', va='bottom', color='black')  # No decimal for integer percentage

    # Add average lines for attendance and engagement
    ax.axhline(
        y=attendance_avg_stats,
        color='#005288',
        linestyle='--',
        linewidth=1.5,
        label=f'Attendance Average: {attendance_avg_stats}%'
    )
    ax.axhline(
        y=engagement_avg_stats,
        color='#3AB0FF',
        linestyle='--',
        linewidth=1.5,
        label=f'Engagement Average: {engagement_avg_stats}%'
    )
    
    # Set labels, title, and legend
    ax.set_xlabel('Student')
    ax.set_ylabel('Percentage (%)')
    ax.set_title('Student Attendance and Engagement Metrics')
    # ax.legend()
    # ax.legend(loc='upper right', bbox_to_anchor=(1.25, 1), borderaxespad=0.)
    ax.legend(loc='upper right', frameon=False)
    ax.set_xticks(index)  # Set x-ticks to the index
    ax.set_xticklabels(student_metrics_df['Student'], rotation=0, ha='right')  # Set student names as x-tick labels
    # Set the y-axis limits and tick locations
    ax.set_ylim(0, 119)  # Range from 0 to 100
    ax.yaxis.set_ticks(range(0, 119, 20))  # Increments of 20

    # Hide the top and right spines
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)

    # Display the plot
    plt.tight_layout()  # Adjust layout to fit elements
    # plt.show()  # Show the plot in a script environment (use st.pyplot(fig) in Streamlit)
    st.pyplot(fig)  # This line displays the plot

    return fig

def download_chart(fig, filename):
    # Create a buffer to hold the image data
    buffer = io.BytesIO()
    # Save the figure to the buffer
    fig.savefig(buffer, format='png')
    # Set the file pointer to the beginning
    buffer.seek(0)
    # Add a download button to Streamlit
    st.download_button(label="Download Chart", data=buffer, file_name=filename, mime='image/png', icon="📊", use_container_width=True)

def download_llm_output(content, filename):
    # Create a buffer to hold the text data
    buffer = io.BytesIO()
    buffer.write(content.encode('utf-8'))
    buffer.seek(0)
    # Add a download button to Streamlit
    st.download_button(label="Download AI Output", data=buffer, file_name=filename, mime='text/plain', icon="✏️", use_container_width=True)

def prepare_llm_input(student_metrics_df):
    # Convert the student metrics DataFrame to a string
    metrics_str = student_metrics_df.to_string(index=False)
    llm_input = f"""
Based on the following student metrics:

{metrics_str}

Provide:

1. Notes and Key Takeaways: Summarize the data, highlight students with the lowest and highest attendance and engagement percentages, identify students who may need adjustments to their intervention due to low attendance or engagement, and highlight students who are showing strong performance. 

2. Recommendations and Next Steps: Provide interpretations based on the analysis and suggest possible next steps or strategies to improve student outcomes.
"""
    return llm_input

    
# def prompt_response_from_hf_llm(llm_input):
#     # Generate the refined prompt using Hugging Face API
#     response = client.chat.completions.create(
#         model="meta-llama/Llama-3.1-70B-Instruct",
#         messages=[
#             {"role": "user", "content": llm_input}
#         ],
#         stream=True,
#         temperature=0.5,
#         max_tokens=1024,
#         top_p=0.7
#     )

def prompt_response_from_hf_llm(llm_input):
    # Define a system prompt to guide the model's responses
    system_prompt = """
    <Persona> An expert Implementation Specialist at Michigan's Multi-Tiered System of Support Technical Assistance Center (MiMTSS TA Center) with deep expertise in SWPBIS, SEL, Structured Literacy, Science of Reading, and family engagement practices.</Persona> 
    <Task> Analyze educational data and provide evidence-based recommendations for improving student outcomes across multiple tiers of support, drawing from established frameworks in behavioral interventions, literacy instruction, and family engagement.</Task> 
    <Context> Operating within Michigan's educational system to support schools in implementing multi-tiered support systems, with access to student metrics data and knowledge of state-specific educational requirements and MTSS frameworks. </Context> 
    <Format> Deliver insights through clear, actionable recommendations supported by data analysis, incorporating technical expertise while maintaining accessibility for educators and administrators at various levels of MTSS implementation.</Format>
    """
    
    # Generate the refined prompt using Hugging Face API
    response = client.chat.completions.create(
        model="meta-llama/Llama-3.1-70B-Instruct",
        messages=[
            {"role": "system", "content": system_prompt},  # Add system prompt here
            {"role": "user", "content": llm_input}
        ],
        stream=True,
        temperature=0.5,
        max_tokens=1024,
        top_p=0.7
    )
    
    # Combine messages if response is streamed
    response_content = ""
    for message in response:
        response_content += message.choices[0].delta.content

    return response_content.strip()

if __name__ == '__main__':
    main()