#!/usr/bin/env python
# coding: utf-8

# # NYC Subway With BigQuery #

# In[1]:


import geopandas as gpd
from shapely import wkt
from google.oauth2 import service_account
from google.cloud import bigquery

from lets_plot import *
LetsPlot.setup_html()


# In[2]:


def load_gdf(client, *, geometry_name, query):
    df = client.query(query).result().to_dataframe()
    df[geometry_name] = df[geometry_name].apply(wkt.loads)
    return gpd.GeoDataFrame(df, geometry=geometry_name)


# In[3]:


# You should add your own bigquery_credentials.json file
# See https://cloud.google.com/docs/authentication/getting-started
credentials = service_account.Credentials.from_service_account_file(
    '../data/bigquery_credentials.json',
    scopes=['https://www.googleapis.com/auth/cloud-platform'],
)
client = bigquery.Client(credentials=credentials, project=credentials.project_id)


# In[4]:


boroughs_gdf = load_gdf(client, geometry_name='borough_geom', query='''
    SELECT borough_name, borough_geom
    FROM `bigquery-public-data.new_york_subway.geo_nyc_borough_boundaries`
''')


# In[5]:


stations_gdf = load_gdf(client, geometry_name='station_geom', query='''
    SELECT station_name, line, station_geom
    FROM `bigquery-public-data.new_york_subway.stations`
''')


# In[6]:


lines_gdf = load_gdf(client, geometry_name='line_geom', query='''
    SELECT ny_s.line, ST_MakeLine(ARRAY(
        SELECT station_geom
        FROM `bigquery-public-data.new_york_subway.stations`
        WHERE line = ny_s.line
        )) AS line_geom
    FROM(
        SELECT line 
        FROM `bigquery-public-data.new_york_subway.stations`
        ) AS ny_s
    GROUP BY line
''')


# In[7]:


ggplot() + \
    geom_polygon(aes(fill='borough_name'), data=boroughs_gdf, color='black', alpha=.15, \
                 tooltips=layer_tooltips().line('@borough_name')) + \
    geom_path(aes(color='line'), data=lines_gdf, size=2) + \
    geom_point(data=stations_gdf, shape=1, size=1.5, color='black', \
               tooltips=layer_tooltips().line('@station_name')) + \
    ggtitle('NYC Subway') + \
    ggsize(800, 600) + \
    theme(legend_position='none', axis_title='blank', axis_text='blank', axis_ticks='blank', \
          axis_line='blank', axis_tooltip='blank')