#! usr/bin/env python

# This script is for testing without having to host the flask app.

import folium
import pandas as pd
import os
from sqlalchemy import create_engine
import geojson

DATABASE_URL = os.environ["DATABASE_URL"]
STATES_GEOJSON_PATH = "static/us-states.json"

engine = create_engine(DATABASE_URL)
with engine.connect() as db:

    query = "Select state, count(*) From registry Group By state;"
    df = pd.read_sql_query(query, db)

with open(STATES_GEOJSON_PATH, "r") as file:
    gj = geojson.load(file)

# Folium choropleth requires a one-to-one correspondence between GeoJSON
#  features (state definitions) and shade values, so we will make a new
#  GeoJSON object that is a FeatureCollection of only the states that we
#  have data for.
relevant_features = [feature for feature in gj["features"]
                     if ("id" in feature) and
                     (feature["id"] in df["state"].values)]

gj_relevant = geojson.FeatureCollection(relevant_features)

geo_str = geojson.dumps(gj_relevant)

base_map = folium.Map([43, -100], zoom_start=5)

base_map.choropleth(
    geo_str=geo_str,
    data=df,
    columns=['state', 'count'],
    key_on='feature.id',
    fill_color='PuBuGn',
)


base_map.save("map.html")
