import pandas as pd
import pprint
import requests

state_codes = ['FL', 'MI', 'PA']

def collapse_results_by_party(results_by_candidate):
    candidates = {
        'trumpd': 'rep',
        'bidenj': 'dem'
    }
    results_by_party = {}
    for candidate, count in results_by_candidate.iteritems():
        party = candidates.get(candidate, 'trd')
        results_by_party[party] = results_by_party.get(candidate, 0) + (count or 0)

    return results_by_party

def flatten_results_to_dict(d, key, results_by_candidate):
    results_by_party = collapse_results_by_party(results_by_candidate)
    for party in ['dem', 'rep', 'trd']:
        d['{}_{}'.format(key, party)] = results_by_party[party]

all_results = {}
for state_code in state_codes:
    print 'Downloading {}'.format(state_code)
    state_result = requests.get('https://static01.nyt.com/elections-assets/2020/data/api/2020-11-03/precincts/{}GeneralConcatenator-latest.json'.format(state_code)).json()
    all_results[state_code] = state_result

print 'Processing precincts'
precinct_records = []
for state, state_result in all_results.iteritems():
    print 'Processing {}'.format(state)
    precinct_totals = None
    if state == 'FL' or state == 'PA':
        precinct_totals = state_result['precinct_totals']
    elif state == 'MI':
        precinct_totals = state_result['precincts']

    for precinct_result in precinct_totals:
        flatten_results_to_dict(precinct_result, 'votes', precinct_result['results'])
        precinct_result['state'] = state
        precinct_result.pop('results')
        precinct_records.append(precinct_result)

print 'Processing precinct vote types'
precinct_vote_type_records = []
for state, state_result in all_results.iteritems():
    print 'Processing {}'.format(state)
    if state == 'MI':
        continue

    precinct_vote_types = state_result['precinct_by_vote_type']

    for precinct_vote_type in precinct_vote_types:
        flatten_results_to_dict(precinct_vote_type, 'votes', precinct_vote_type['results'])
        precinct_vote_type['state'] = state
        print precinct_vote_type
        precinct_vote_type.pop('results')
        precinct_vote_type_records.append(precinct_vote_type)

p_df = pd.DataFrame.from_records(precinct_records)
pvt_df = pd.DataFrame.from_records(precinct_vote_type_records)
p_df.to_csv('data/nyt_precinct.csv', encoding='utf-8')
pvt_df.to_csv('data/nyt_precinct_vote_types.csv', encoding='utf-8')