import pandas as pd
import plotly.graph_objects as go

school_rename_map = {
    'Dr. Martin Luther King Jr. Literary & Fine Arts School': 'King Arts',
    'Dr Martin Luther King Jr Literary & Fine Arts School': 'King Arts',
    'Dr. Martin Luther King Jr. Literary and Fine Arts School': 'King Arts',
    'Washington Elementary School': 'Washington',
    'Lincoln Elementary School': 'Lincoln',
    'Oakton Elementary School': 'Oakton',
    'Dawes Elementary School': 'Dawes',
    'Dewey Elementary School': 'Dewey',
    'Walker Elementary School': 'Walker',
    'Lincolnwood Elementary School': 'Lincolnwood',
    'Willard Elementary School': 'Willard',
    'Kingsley Elementary School': 'Kingsley',
    'Dr. Bessie Rhodes School of Global Studies': 'Bessie Rhodes',
    'Orrington Elementary School': 'Orrington',
    "Rice Children's Center": "Rice Center",
    'Chute Middle School': 'Chute',
    'Nichols Middle School': 'Nichols',
    'Haven Middle School': 'Haven',
    'Foster School': 'Foster'
}

gen_ed_elementary_schools = [
    'King Arts',
    'Washington',
    'Lincoln',
    'Oakton',
    'Dawes',
    'Dewey',
    'Walker',
    'Lincolnwood',
    'Willard',
    'Kingsley',
    'Bessie Rhodes',
    'Orrington',
    'Foster'
]

staff_directory = pd.read_csv('staff_directory.csv')

# Exclude staff that work in multiple locations for simplicity
staff_directory = staff_directory.dropna(subset=['location'])
filtered_df = staff_directory[~staff_directory['location'].str.contains(',')].copy()

filtered_df['location'] = filtered_df['location'].replace(school_rename_map)
schools = filtered_df['location'].unique().tolist()

# Display staff count and teacher count per school in a dataframe
school_summary = []
for school in schools:
    school_staff = filtered_df[filtered_df['location'] == school]
    total_staff = len(school_staff)
    total_teachers = len(school_staff[school_staff['title'] == 'Teacher'])
    
    school_summary.append({
        'school': school,
        'total_staff': total_staff,
        'total_teachers': total_teachers
    })

# Convert to DataFrame
school_summary_df = pd.DataFrame(school_summary)

# Pull in baseline utilization data to get SY25 enrollment numbers
baseline_utilization = pd.read_csv('0_utilization.csv')
baseline_utilization['Schools'] = baseline_utilization['Schools'].replace(school_rename_map)

# Merge Enroll SY25 from baseline_utilization into school_summary_df_filtered
school_summary_df = school_summary_df.merge(
    baseline_utilization[['Schools', 'Enroll SY25']],
    left_on='school',
    right_on='Schools',
    how='left'
)
# Drop redundant 'Schools' column
school_summary_df = school_summary_df.drop(columns=['Schools'])

school_summary_df['staff_per_100_students'] = school_summary_df['total_staff'] / school_summary_df['Enroll SY25'] * 100
school_summary_df['teacher_per_100_students'] = school_summary_df['total_teachers'] / school_summary_df['Enroll SY25'] * 100

all_locations = school_summary_df.copy()
elementary_schools = school_summary_df[school_summary_df['school'].isin(gen_ed_elementary_schools)].copy()

# Total staff by location
all_locations = all_locations.sort_values('total_staff', ascending=False)

# Plot Total Staff by Location
fig = go.Figure()
fig.add_trace(go.Bar(
    x=all_locations['total_staff'],
    y=all_locations['school'],
    orientation='h',
    text=all_locations['total_staff'],
    textposition='auto'
))
fig.update_layout(
    title='Total Staff by Location',
    xaxis_title='Number of Staff',
    yaxis_title='Location',
    height=800
)
fig.write_html('total_staff_by_location.html')
fig.show()

# Total enrollment by location
all_locations = all_locations.sort_values('Enroll SY25', ascending=False)

# Drop locations with missing enrollment data
enrollment_locations = all_locations[all_locations['Enroll SY25'].notna()]

# Plot Total Staff by Location
fig = go.Figure()
fig.add_trace(go.Bar(
    x=enrollment_locations['Enroll SY25'],
    y=enrollment_locations['school'],
    orientation='h',
    text=enrollment_locations['Enroll SY25'],
    textposition='auto'
))
fig.update_layout(
    title='Total Enrollment by Location',
    xaxis_title='Number of Students',
    yaxis_title='School',
    height=800
)
fig.write_html('total_enrollment_by_location.html')
fig.show()

# Enrollment by school plot
elementary_schools = elementary_schools.sort_values('Enroll SY25', ascending=False)

# Plot Enrollment by School
fig = go.Figure()
fig.add_trace(go.Bar(
    x=elementary_schools['Enroll SY25'],
    y=elementary_schools['school'],
    orientation='h',
    text=elementary_schools['Enroll SY25'],
    textposition='auto'
))
fig.update_layout(
    title='SY 2025 Enrollment by Elementary School',
    xaxis_title='Number of Students',
    yaxis_title='School',
    height=800
)
fig.write_html('elementary_enrollment_by_school.html')
fig.show()

# Staff by school plot
elementary_schools = elementary_schools.sort_values('total_staff', ascending=False)

# Plot Staff by School
fig = go.Figure()
fig.add_trace(go.Bar(
    x=elementary_schools['total_staff'],
    y=elementary_schools['school'],
    orientation='h',
    text=elementary_schools['total_staff'],
    textposition='auto'
))
fig.update_layout(
    title='Staff by Elementary School',
    xaxis_title='Number of Staff',
    yaxis_title='School',
    height=800
)
fig.write_html('elementary_staff_by_school.html')
fig.show()

# Teachers by school plot
elementary_schools = elementary_schools.sort_values('total_teachers', ascending=False)

# Plot Teachers by School
fig = go.Figure()
fig.add_trace(go.Bar(
    x=elementary_schools['total_teachers'],
    y=elementary_schools['school'],
    orientation='h',
    text=elementary_schools['total_teachers'],
    textposition='auto'
))
fig.update_layout(
    title='Teachers by Elementary School',
    xaxis_title='Number of Teachers',
    yaxis_title='School',
    height=800
)
fig.write_html('elementary_teachers_by_school.html')
fig.show()

# Staff per 100 students by school plot
elementary_schools = elementary_schools.sort_values('staff_per_100_students', ascending=False)

# Plot Staff by School
fig = go.Figure()
fig.add_trace(go.Bar(
    x=elementary_schools['staff_per_100_students'],
    y=elementary_schools['school'],
    orientation='h',
    text=elementary_schools['staff_per_100_students'].round(1),
    textposition='auto'
))
fig.update_layout(
    title='Staff per 100 Students by Elementary School',
    xaxis_title='Number of Staff per 100 Students',
    yaxis_title='School',
    height=800
)
fig.write_html('elementary_staff_per_100_students_by_school.html')
fig.show()

# Teachers per 100 students by school plot
elementary_schools = elementary_schools.sort_values('teacher_per_100_students', ascending=False)

# Plot Teachers by School
fig = go.Figure()
fig.add_trace(go.Bar(
    x=elementary_schools['teacher_per_100_students'],
    y=elementary_schools['school'],
    orientation='h',
    text=elementary_schools['teacher_per_100_students'].round(1),
    textposition='auto'
))
fig.update_layout(
    title='Teachers per 100 Students by Elementary School',
    xaxis_title='Number of Teachers per 100 Students',
    yaxis_title='School',
    height=800
)
fig.write_html('elementary_teachers_per_100_students_by_school.html')
fig.show()