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

transportation_1A = pd.read_csv('data/1A_transportation_idot_d65.csv')
transportation_1A = transportation_1A.rename(columns={'new_attendance_area_school': 'Schools'})
transportation_1A['Schools'] = transportation_1A['Schools'].replace(school_rename_map)
transportation_1A['Enroll Total'] = transportation_1A['Bus'] + transportation_1A['Hazard'] + transportation_1A['Walk <=.75 miles'] + transportation_1A['Walk >.75 miles']
transportation_1A['walk_pct'] = transportation_1A['Walk <=.75 miles'] / transportation_1A['Enroll Total'] * 100

transportation_1B = pd.read_csv('data/1B_transportation_idot_d65.csv')
transportation_1B = transportation_1B.rename(columns={'new_attendance_area_school': 'Schools'})
transportation_1B['Schools'] = transportation_1B['Schools'].replace(school_rename_map)
transportation_1B['Enroll Total'] = transportation_1B['Bus'] + transportation_1B['Hazard'] + transportation_1B['Walk <=.75 miles'] + transportation_1B['Walk >.75 miles']
transportation_1B['walk_pct'] = transportation_1B['Walk <=.75 miles'] / transportation_1B['Enroll Total'] * 100

transportation_2FR = pd.read_csv('data/2FR_transportation_idot_d65.csv')
transportation_2FR['Schools'] = transportation_2FR['Schools'].replace(school_rename_map)
transportation_2FR['Enroll Total'] = transportation_2FR['Bus'] + transportation_2FR['Hazard'] + transportation_2FR['Walk <=.75 miles'] + transportation_2FR['Walk >.75 miles']
transportation_2FR['walk_pct'] = transportation_2FR['Walk <=.75 miles'] / transportation_2FR['Enroll Total'] * 100

transportation_2DR = pd.read_csv('data/2DR_transportation_idot_d65.csv')
transportation_2DR['Schools'] = transportation_2DR['Schools'].replace(school_rename_map)
transportation_2DR['Enroll Total'] = transportation_2DR['Bus'] + transportation_2DR['Hazard'] + transportation_2DR['Walk <=.75 miles'] + transportation_2DR['Walk >.75 miles']
transportation_2DR['walk_pct'] = transportation_2DR['Walk <=.75 miles'] / transportation_2DR['Enroll Total'] * 100

elementary_transportation_1A = transportation_1A[transportation_1A['Schools'].isin(gen_ed_elementary_schools)].copy()
elementary_transportation_1B = transportation_1B[transportation_1B['Schools'].isin(gen_ed_elementary_schools)].copy()
elementary_transportation_2FR = transportation_2FR[transportation_2FR['Schools'].isin(gen_ed_elementary_schools)].copy()
elementary_transportation_2DR = transportation_2DR[transportation_2DR['Schools'].isin(gen_ed_elementary_schools)].copy()

feeder_map = {
    'Washington': 'Nichols',
    'Lincoln': 'Nichols',
    'Oakton': 'Chute',
    'Dawes': 'Chute',
    'Dewey': 'Nichols',
    'Walker': 'Chute',
    'Lincolnwood': 'Haven',
    'Willard': 'Haven',
    'Kingsley': 'Haven',
    'Orrington': 'Haven',
    'Foster': 'Haven'
}

# Add column for feeder middle school
elementary_transportation_2DR['Feeder'] = elementary_transportation_2DR['Schools'].map(feeder_map)
elementary_transportation_2FR['Feeder'] = elementary_transportation_2FR['Schools'].map(feeder_map)
elementary_transportation_1A['Feeder'] = elementary_transportation_1A['Schools'].map(feeder_map)
elementary_transportation_1B['Feeder'] = elementary_transportation_1B['Schools'].map(feeder_map)

# Group by feeder middle school and sum the Walk and Enroll Total columns
feeder_walk_2DR = elementary_transportation_2DR.groupby('Feeder').agg({
    'Walk <=.75 miles': 'sum',
    'Enroll Total': 'sum'
}).reset_index()
feeder_walk_2DR['walk_pct'] = feeder_walk_2DR['Walk <=.75 miles'] / feeder_walk_2DR['Enroll Total'] * 100

feeder_walk_2FR = elementary_transportation_2FR.groupby('Feeder').agg({
    'Walk <=.75 miles': 'sum',
    'Enroll Total': 'sum'
}).reset_index()
feeder_walk_2FR['walk_pct'] = feeder_walk_2FR['Walk <=.75 miles'] / feeder_walk_2FR['Enroll Total'] * 100   

feeder_walk_1A = elementary_transportation_1A.groupby('Feeder').agg({
    'Walk <=.75 miles': 'sum',
    'Enroll Total': 'sum'
}).reset_index()
feeder_walk_1A['walk_pct'] = feeder_walk_1A['Walk <=.75 miles'] / feeder_walk_1A['Enroll Total'] * 100

feeder_walk_1B = elementary_transportation_1B.groupby('Feeder').agg({
    'Walk <=.75 miles': 'sum',
    'Enroll Total': 'sum'
}).reset_index()
feeder_walk_1B['walk_pct'] = feeder_walk_1B['Walk <=.75 miles'] / feeder_walk_1B['Enroll Total'] * 100

# Create a side by side bar chart comparing walkability percentages across scenarios for each elementary school
fig = go.Figure()
fig.add_trace(go.Bar(
    y=elementary_transportation_1A['Schools'],
    x=elementary_transportation_1A['walk_pct'],
    name='Scenario 1A',
    text=elementary_transportation_1A['walk_pct'].round(1),
    textposition='auto',
    orientation='h'
))
fig.add_trace(go.Bar(
    y=elementary_transportation_1B['Schools'],
    x=elementary_transportation_1B['walk_pct'],
    name='Scenario 1B',
    text=elementary_transportation_1B['walk_pct'].round(1),
    textposition='auto',
    orientation='h'
))
fig.add_trace(go.Bar(
    y=elementary_transportation_2FR['Schools'],
    x=elementary_transportation_2FR['walk_pct'],
    name='Scenario 2FR',
    text=elementary_transportation_2FR['walk_pct'].round(1),
    textposition='auto',
    orientation='h'
))
fig.add_trace(go.Bar(
    y=elementary_transportation_2DR['Schools'],
    x=elementary_transportation_2DR['walk_pct'],
    name='Scenario 2DR',
    text=elementary_transportation_2DR['walk_pct'].round(1),
    textposition='auto',
    orientation='h'
))
fig.update_layout(
    title='Walkability Percentage by Elementary School Across Scenarios',
    xaxis_title='Percentage of Students Walking (%)',
    yaxis_title='School',
    barmode='group',
    height=1600
)
fig.write_html('assets/elementary_walkability_comparison.html')
fig.show()

# Plot walkability percentages by feeder middle school
fig = go.Figure()
fig.add_trace(go.Bar(
    y=feeder_walk_1A['Feeder'],
    x=feeder_walk_1A['walk_pct'],
    name='Scenario 1A',
    text=feeder_walk_1A['walk_pct'].round(1),
    textposition='auto',
    orientation='h'
))
fig.add_trace(go.Bar(
    y=feeder_walk_1B['Feeder'],
    x=feeder_walk_1B['walk_pct'],
    name='Scenario 1B',
    text=feeder_walk_1B['walk_pct'].round(1),
    textposition='auto',
    orientation='h'
))
fig.add_trace(go.Bar(
    y=feeder_walk_2FR['Feeder'],
    x=feeder_walk_2FR['walk_pct'],
    name='Scenario 2FR',
    text=feeder_walk_2FR['walk_pct'].round(1),
    textposition='auto',
    orientation='h'
))
fig.add_trace(go.Bar(
    y=feeder_walk_2DR['Feeder'],
    x=feeder_walk_2DR['walk_pct'],
    name='Scenario 2DR',
    text=feeder_walk_2DR['walk_pct'].round(1),
    textposition='auto',
    orientation='h'
))
fig.update_layout(
    title='Walkability Percentage of Elementary Schools Grouped by Feeder Pattern Across Scenarios',
    xaxis_title='Percentage of Students Walking (%)',
    yaxis_title='Feeder Pattern',
    barmode='group',
    height=800
)
fig.write_html('assets/feeder_walkability_comparison.html')
fig.show()