#!/usr/bin/env python3

import pandas as pd

# This is the data generated by prepare_zipcodes.py
CONTRIB_PATH = 'contrib-by-zip.csv'
ZIP_DATA_PATH = 'zip_code_database.csv'
FIPS_PATH = 'ZIP-COUNTY-FIPS_2018-03.csv'
OUTPUT_PATH = 'contrib-by-county.csv'

if __name__ == '__main__':
    print('Loading data...')
    cont = pd.read_csv(CONTRIB_PATH)
    zips = pd.read_csv(ZIP_DATA_PATH)[['zip', 'state']]
    zips.rename(columns={'state': 'state_id'}, inplace=True)
    fips = pd.read_csv(FIPS_PATH)
    fips.rename(columns={
        'ZIP': 'zip',
        'STCOUNTYFP': 'county_id',
        'STATE': 'state_id',
        'COUNTYNAME': 'county_name'
    }, inplace=True)
    fips = fips[['zip', 'county_id', 'county_name', 'state_id']]
    fips['zip'].astype(int, inplace=True)
    fips['county_id'].astype(int, inplace=True)

    print('Preparing data...')
    # Fix state ID
    zips = fips[['zip', 'state_id']].drop_duplicates()
    cont.drop(columns=['state_id'], inplace=True)
    cont = cont.groupby(['zip']).sum().reset_index()
    cont = pd.merge(cont, zips, on='zip', how='outer')

    # Add county name and codes to our contribution data.
    fips2 = fips[['zip', 'county_id']].drop_duplicates()
    df = pd.merge(cont, fips2, on='zip', how='outer')

    # Sum contributions grouped by county
    df.drop(columns=['zip'], inplace=True)
    df = df.groupby(['county_id']).sum().reset_index()
    df['county_id'] = df['county_id'].astype(int)
    df['cont_cnt'] = df['cont_cnt'].astype(int)

    # Output to file
    print('Saving data to {}'.format(OUTPUT_PATH))
    df.to_csv(OUTPUT_PATH, index=False)
