visualization.py

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('Frequency')
        ax.set_title('Intervention Sessions Held vs Not Held', 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')

        ax.axhline(y=90, color='#48BB88', linestyle='--', linewidth=1.5, label=f'Attendance Threshold: ≥ 90%')
        ax.axhline(y=80, color='#235E44', 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}%')

        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, 125)
        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(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 addressing logistical barriers", shape="box")
                dot.edge("Q2", "A3", label="Yes")
    
        return dot
    
    def download_tree_diagram(dot, filename):
        # Render the diagram to a file
        filepath = dot.render(filename, format='png', cleanup=True)
        return filepath
        
    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)