#!/usr/bin/python

import geopandas as gpd
import os
import fiona
import sys
import csv

states_to_include = ["MN"]

state_fipses_to_include = []
county_lookup = {}
print("Reading county census data...")
with open('national_cousub2020.txt') as csvfile:
    reader = csv.DictReader(csvfile, delimiter='|')
    for row in reader:
        if row['STATE'] in states_to_include:
            state_fipses_to_include.append(row['STATEFP'])
        county_lookup[row['STATEFP'] + row['COUNTYFP']] = row

input_file = sys.argv[1]

print("Reading input gdb...")
gdf = gpd.read_file(input_file)
gdf = gdf[gdf['STATEFIPS'].isin(state_fipses_to_include)]

print("Reprojecting...")
gdf = gdf.to_crs("EPSG:4326")

print("Simplifying geometry...")
gdf['geometry'] = gdf['geometry'].simplify(0.0001, preserve_topology=True)

print("Calculating FULLFIPS...")
gdf['FULLFIPS'] = gdf['STATEFIPS'].astype(str) + gdf['CNTYFIPS'].astype(str)

print("Finding unique FULLFIPS...")
counties = gdf['FULLFIPS'].unique()

# TODO: Trim down which fields are included
#
# "CSBID", "CSBYEARS", "CSBACRES",
# "CDL2016", "CDL2017", "CDL2018", "CDL2019", "CDL2020", "CDL2021", "CDL2022", "CDL2023",
# "STATEFIPS", "STATEASD", "ASD", "CNTY", "CNTYFIPS",
# "INSIDE_X", "INSIDE_Y", "Shape_Length", "Shape_Area", "FULLFIPS"

for i, county in enumerate(counties, 1):
    print(f"Processing county ({county}): {i}/{len(counties)}")
    county_gdf = gdf[gdf['FULLFIPS'] == county]
    output_file = f"{county}.geojson"
    county_gdf.to_file(os.path.join("data", output_file), driver="GeoJSON", COORDINATE_PRECISION=5)

with open('data/counties.js', 'w') as f:
    for county in counties:
        f.write(f"import county{county} from './{county}.geojson?url';\n")
    f.write('\nexport default {\n')
    for county in counties:
        county_name = county_lookup[county]['COUNTYNAME']
        state_name = county_lookup[county]['STATE']
        f.write(f"    '{county_name}, {state_name}': county{county},\n")
    f.write("};\n")