import matplotlib.pyplot as plt
import io
import streamlit as st
from graphviz import Digraph

class Visualization:
    def plot_intervention_statistics(self, intervention_stats):
        sessions_held = intervention_stats['Intervention Sessions Held'].values[0]
        sessions_not_held = intervention_stats['Intervention Sessions Not Held'].values[0]

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

        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)

        ax.set_ylabel('Dosage')
        ax.set_title('Intervention Dosage', fontsize=16)
        handles, labels = ax.get_legend_handles_labels()
        ax.legend(handles[::-1], labels[::-1])

        ax.spines['top'].set_visible(False)
        ax.spines['right'].set_visible(False)
        
        st.pyplot(fig)
        return fig

    def plot_student_metrics(self, student_metrics_df, attendance_avg_stats, engagement_avg_stats):
        fig, ax = plt.subplots(figsize=(10, 6))
        bar_width = 0.35
        index = range(len(student_metrics_df))

        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
        )

        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')

        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')

        # Get the current x-axis limits
        x_min, x_max = ax.get_xlim()

        # Highlight the range 80-89 for "Engagement Threshold" and 90-100 for "Attendance Threshold"
        ax.fill_betweenx(y=[90, 100], x1=x_min, x2=x_max, color='#005288', alpha=0.2, label='Attendance Threshold: ≥ 90%')
        ax.fill_betweenx(y=[80, 90], x1=x_min, x2=x_max, color='#3AB0FF', alpha=0.2, label='Engagement Threshold: ≥ 80%')

        # Add horizontal lines for thresholds and averages
        # ax.axhline(y=90, color='#005288', linestyle=':', linewidth=1.5, label=f'Attendance Threshold: ≥ 90%')
        # ax.axhline(y=80, color='#3AB0FF', linestyle=':', linewidth=1.5, label=f'Engagement Threshold: ≥ 80%')
        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 the x-limits to ensure the fill covers the entire plot width
        ax.set_xlim(left=x_min, right=x_max)
        
        ax.set_xlabel('Student')
        ax.set_ylabel('Percentage (%)')
        ax.set_title('Student Attendance and Engagement Metrics')
        ax.legend(loc='upper right', frameon=False)
        ax.set_xticks(index)
        ax.set_xticklabels(student_metrics_df['Student'], rotation=0, ha='right')
        ax.set_ylim(0, 140)
        ax.yaxis.set_ticks(range(0, 124, 25))

        ax.spines['top'].set_visible(False)
        ax.spines['right'].set_visible(False)

        plt.tight_layout()
        st.pyplot(fig)
        return fig

    def build_tree_diagram(self, row):
        dot = Digraph()
    
        # Add the student's name as the first node
        student_name = row['Student']  # Assuming 'Student' is the column name for student names
        dot.node("Student", student_name)
    
        # Connect the student's name to the first question
        dot.edge("Student", "Q1")
    
        # Add the first question node
        dot.node("Q1", "Has the student attended ≥ 90% of interventions?")
    
        if row["Attended ≥ 90%"] == "No":
            # If attendance is less than 90%, add the corresponding action node
            dot.node("A1", "Address Attendance", shape="box")
            dot.edge("Q1", "A1", label="No")
        else:
            # If attendance is 90% or more, proceed to the next question
            dot.node("Q2", "Has the student been engaged ≥ 80% of intervention time?")
            dot.edge("Q1", "Q2", label="Yes")
    
            if row["Engagement ≥ 80%"] == "No":
                # If engagement is less than 80%, add the corresponding action node
                dot.node("A2", "Address Engagement", shape="box")
                dot.edge("Q2", "A2", label="No")
            else:
                # If engagement is 80% or more, add the final action node
                dot.node("A3", "Consider barriers, fidelity, and progress monitoring", shape="box")
                dot.edge("Q2", "A3", label="Yes")
    
        return dot
    
    def download_chart(self, fig, filename):
        buffer = io.BytesIO()
        fig.savefig(buffer, format='png')
        buffer.seek(0)
        st.download_button(label="Download Chart", data=buffer, file_name=filename, mime='image/png', icon="📊", use_container_width=True)