Py.Cafe

from dash import Dash, html, dcc, callback, Output, Input
import dash_ag_grid as dag
import plotly.graph_objects as go
import pandas as pd
import numpy as np

# Import data
df = pd.read_csv('data/2025-26/EQA_2025-26.csv', header=1)
df['district_label'] = df['district_name'] + ' (' + df['county_name'] + ')'
for col in ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'T', 'U', 'V', 'W', 'base_per_pupil', 'special_education_per_pupil', 'speech_only_per_pupil']:
    df[col] = df[col].str.replace(r'[$,]', '', regex=True).astype(float)

df = df.set_index(['county_code', 'district_code'])

df_cat = pd.read_csv('data/2025-26/CAT_2025-26.csv', header=1)
df_cat['B'] = df_cat['B'].str.replace(r'[$,]', '', regex=True).astype(float)
df_cat['E'] = df_cat['E'].str.replace(r'[$,]', '', regex=True).astype(float)
df_cat = df_cat.set_index(['county_code', 'district_code'])

df_trn = pd.read_csv('data/2025-26/TRN_2025-26.csv', header=1)
df_trn['C'] = df_trn['C'].str.replace(r'[$,]', '', regex=True).astype(float)
df_trn = df_trn.set_index(['county_code', 'district_code'])

df_lim = pd.read_csv('data/2025-26/LIM_2025-26.csv', header=1)
df_lim['B'] = df_lim['B'].str.replace(r'[$,]', '', regex=True).astype(float)
df_lim = df_lim.set_index(['county_code', 'district_code'])

df_net = pd.read_csv('data/2025-26/NET_2025-26.csv', header=1)
df_net['A'] = df_net['A'].str.replace(r'[$,]', '', regex=True).astype(float)
df_net['A-2'] = df_net['A-2'].str.replace(r'[$,]', '', regex=True).astype(float)
df_net['A-3'] = df_net['A-3'].str.replace(r'[$,]', '', regex=True).astype(float)
df_net['A-4'] = df_net['A-4'].str.replace(r'[$,]', '', regex=True).astype(float)
df_net['A-7'] = df_net['A-7'].str.replace(r'[$,]', '', regex=True).astype(float)
df_net['A-8'] = df_net['A-8'].str.replace(r'[$,]', '', regex=True).astype(float)
df_net = df_net.set_index(['county_code', 'district_code'])

# GCA test
def remove_gca(df, remove):
    if remove:
        adequacy = (
            df['A'] * df['base_per_pupil']
            + df['C'] * df['base_per_pupil']
            + df['E'] * df['base_per_pupil']
            + df['G'] * df['base_per_pupil']
            + df['I'] * df['special_education_per_pupil'] * 2/3
            + df['K'] * df['speech_only_per_pupil']
        )
        lfs = np.where(df['T'] > 0, df['S'] * adequacy, df['P'])
        ret = (adequacy - lfs).clip(lower=0)
        return ret
    else:
        return df['W']
    
# Calculate general fund aid
def calc_general_fund_aid(df_eqa, df_cat, df_trn, df_lim, df_net):
    df = df_eqa[['district_label', 'change']].copy()
    print(df.loc[(25, 3260), 'district_label'])
    
    # Remove GCA from CAT
    df['cat_b'] = df_cat['B'] / df_cat['gca']
    df['cat_e'] = df_cat['E'] / df_cat['gca']

    df['uncapped_aid'] = df['change'] + df['cat_b'] + df['cat_e']
    df['uncapped_aid'] = (df['uncapped_aid'] + df_trn['C']).where(df_net['A-4'] != 0, df['uncapped_aid'])
    print(df.loc[(25, 3260), 'change'])
    print(df_cat.loc[(25, 3260), 'B'])
    print(df_cat.loc[(25, 3260), 'E'])
    print(df.loc[(25, 3260), 'uncapped_aid'])
    df['sfra_change'] = df['uncapped_aid']-df_lim['B']
    print(df.loc[(25, 3260), 'sfra_change'])
    df['sfra_adj_dollars'] = ((df_lim['B'] * 0.06) - df['sfra_change']).clip(upper=0).where(df['sfra_change'] > 0, ((df_lim['B'] * -0.03) - df['sfra_change']).clip(lower=0))
    print(df.loc[(25, 3260), 'sfra_adj_dollars'])
    df['sfra_adj_percent'] = df['sfra_adj_dollars'] / df['uncapped_aid'] + 1
    df['new_general_fund_aid'] = (
        df['change'] * df['sfra_adj_percent']
        + df_net['A-2']
        + df_net['A-3']
        + df['cat_b'] * df['sfra_adj_percent']
        + df['cat_e'] * df['sfra_adj_percent']
        + df_net['A-7']
    )
    df['new_general_fund_aid'] = (df['new_general_fund_aid'] + df_trn['C'] * df['sfra_adj_percent']).where(df_net['A-4'] != 0, df['new_general_fund_aid'])
    df['new_general_fund_aid'] = (df['new_general_fund_aid'] + df_net['A-8']).where(df_net['A-8'].notna(), df['new_general_fund_aid'])
    print(df.loc[(25, 3260), 'new_general_fund_aid'])
    df['general_fund_aid'] = df_net['A']
    df['general_fund_aid_change'] = df['new_general_fund_aid'] - df['general_fund_aid']
    df['general_fund_aid_change_percent'] = df['general_fund_aid_change'] / df['general_fund_aid']
    return df[['general_fund_aid', 'new_general_fund_aid', 'general_fund_aid_change', 'general_fund_aid_change_percent']]


# Initialize the Dash app
app = Dash()

display_controls = html.Div(
    children=[
        html.Span('Change from Original: '),
        dcc.RadioItems(options=['Dollars', 'Percentage'], value='Dollars', id='change-type'),
        html.Br(),
        html.Span('Sort By: '),
        dcc.Dropdown(['District Name', 'Change Amount'], 'Change Amount', id='sort-by'),
        dcc.RadioItems(options=['asc','desc'], value='asc', id='sort-order'),
        html.Br(),
    ]
)


app.layout = [
    html.H1(children='Draft SFRA Funding Vizualizations and Calculators'),
    html.Div(
        children=[
            html.Div(
                children=[
                    display_controls,
                    html.H3('Controls'),
                    dcc.RadioItems(options=['Standard GCA', 'All GCA = 1.0', 'Any GCA <1 set to 1.0'], value='Standard GCA', id='butterfly-mode'),
                    html.H3('Total Dollar Change:'),
                    html.Div(id='total-dollar-change')
                ],
                style={'padding': '10px'}
            ),
            html.Div(
                children=[
                    dcc.Graph(figure={}, id='butterfly-graph')
                ],
                style={'padding': '10px'}
            )
        ],
        style={
            'display': 'grid',
            'gridTemplateColumns': '20% 70%',
            'gap': '20px'
        }
    ),
    # dag.AgGrid(
    #     rowData=df.to_dict('records'),
    #     columnDefs=[{"field": col} for col in df.columns],
    #     defaultColDef={"sortable": True, "filter": True, "resizable": True}
    # ),
]

@callback(
    Output(component_id='butterfly-graph', component_property='figure'),
    Output(component_id='total-dollar-change', component_property='children'),
    Input(component_id='change-type', component_property='value'),
    Input(component_id='sort-by', component_property='value'),
    Input(component_id='sort-order', component_property='value'),
    Input(component_id='butterfly-mode', component_property='value')
)
def update_graph(
        change_type,
        sort_by, sort_order,
        mode
    ):
    global df, df_cat, df_trn, df_lim, df_net
    # Create graph
    fig = go.Figure()
    if mode == 'Standard GCA':
        df['change'] = 0
        total_dollar_change = 0
    elif mode == 'All GCA = 1.0':
        df['change'] = remove_gca(df, remove=True)
        df_new = calc_general_fund_aid(df, df_cat, df_trn, df_lim, df_net)
        if change_type == 'Percentage':
            df['change'] = df_new['general_fund_aid_change_percent']
        else:
            df['change'] = df_new['general_fund_aid_change']
        total_dollar_change = df_new['general_fund_aid_change'].sum()
    else:  # Any GCA <1 set to 1.0
        df['change'] = remove_gca(df, remove=True)
        df_new = calc_general_fund_aid(df, df_cat, df_trn, df_lim, df_net)
        if change_type == 'Percentage':
            df['change'] = df_new['general_fund_aid_change_percent'].where(df['gca'] < 1.0, 0)
        else:
            df['change'] = df_new['general_fund_aid_change'].where(df['gca'] < 1.0, 0)
        total_dollar_change = df_new['general_fund_aid_change'].where(df['gca'] < 1.0, 0).sum()
    
    # X axis range
    max_x_range = max(1_000_000, df['change'].abs().max()+1_000_000) if change_type == 'Dollars' else max(.2, df['change'].abs().max()+.1)
    tickformat = '$,' if change_type == 'Dollars' else '.1%'
    
    # Sort
    if sort_by == 'District Name':
        df = df.sort_values(by='district_label', ascending=(sort_order=='asc'))
    else:
        df = df.sort_values(by='change', ascending=(sort_order=='asc'))

    fig.add_bar(y=df['district_label'], x=df['change'], orientation='h')
    fig.update_layout(
        title='Change in General Fund Aid by District',
        yaxis_title='School District',
        barmode='relative',
        height=len(df) * 30 + 150
    )
    fig.update_xaxes(
        title=f"Change in General Fund Aid ({change_type})",
        side='top',
        zeroline=True,
        zerolinecolor='black',
        zerolinewidth=2,
        tickformat=tickformat,
        range=[-max_x_range, max_x_range]
    )

    # return the fig 
    return fig, f'${total_dollar_change:,.0f}'

if __name__ == "__main__":
    app.run(debug=True)