%%capture

import warnings

warnings.filterwarnings("ignore")
import calitp_portfolio.magics

import folium
import geopandas as gpd
import gcsfs
import google.auth
import pandas as pd

from great_tables import GT
import gt_extras as gte

import _new_ct_report_utils as utils
from update_vars import DIGEST_DICT, PROCESSED_GCS, abbrev_month

credentials, _ = google.auth.default()

# Comment out and leave this cell right below pandas
#district = "AD 03"

# Parameters
district = "AD 41"

def readable_district_name(district_abbrev: str):

    return (
        district_abbrev
        .replace('AD', 'Assembly District')
        .replace('SD', 'Senate District')
    )

district_full_name = readable_district_name(district)

# %%capture_parameters
district, district_full_name

crosswalk_url = f"{PROCESSED_GCS}{DIGEST_DICT.crosswalk_legislative}_{abbrev_month}.parquet"

crosswalk_legislative = pd.read_parquet(
    crosswalk_url,
    filters=[
        ("legislative_district", "==", district)
    ],
    filesystem = gcsfs.GCSFileSystem()
)

list_of_operators = crosswalk_legislative.analysis_name.unique().tolist()

operator_summary_url = f"{PROCESSED_GCS}{DIGEST_DICT.operator_summary}_{abbrev_month}.parquet"

operator_df = pd.read_parquet(
    operator_summary_url,
    filesystem = gcsfs.GCSFileSystem(),
    filters=[
        ("Day Type", "==", "Weekday"),
        ("Analysis Name", "in", list_of_operators),
    ],
).merge(
    crosswalk_legislative.rename(columns = {"analysis_name": "Analysis Name"}),
    on = "Analysis Name",
    how = "inner"
)

fct_monthly_routes_url = f"{PROCESSED_GCS}{DIGEST_DICT.route_map}_{abbrev_month}.parquet"

fct_monthly_route_df = gpd.read_parquet(
    fct_monthly_routes_url,
    storage_options = {"token": credentials.token},
    filters=[[("Analysis Name", "in", list_of_operators)]]
)

fct_monthly_route_df = utils.prep_gdf(fct_monthly_route_df)

district_gdf = utils.load_legislative_district(district)

try:
    operator_df2 = utils.create_summary_table(
        operator_df, district_col = "legislative_district"
    )
except:
    pass

try:
    display(
        GT(operator_df2)
        .tab_header(
            title = "GTFS Summary Stats",
            subtitle=f"District {district}"
        )
    )
except:
    pass

# Plot both layers here
# rgb to hex (58, 25, 79)
m = district_gdf.explore(
    color = "#3a194f",
    tiles = "CartoDB Positron",
    name = "Legislative District Boundary",
    style_kwds = {"weight": 1, "fillOpacity": 0},
    highlight_kwds = {"fillOpacity": 0.2},
    tooltip=False
)

#color_map = cm.linear.Spectral_11.scale()
fct_monthly_route_df.explore(
    "Route Name",
    m = m,
    name = "Transit Routes in Legislative District",
    categorical = True,
    cmap = "Spectral",
    legend = False
    #marker_kwds={"fill": True},
    #style_kwds={"opacity": 0.5, "fillOpacity": 0.3}
)

folium.LayerControl().add_to(m)
m

try:
    gtfs_table = utils.create_operator_table(
        operator_df, district_col = "legislative_district"
    ).drop(columns = "legislative_district")
except:
    pass

try:
    display((
    GT(gtfs_table.sort_values("Daily Trips", ascending=False))
    .fmt_integer(
        columns=[c for c in gtfs_table.columns if c not in ["Operator"]]
    ).tab_header(
        title="Daily Weekday GTFS stats by operator",
        subtitle=f"District {district}",
    ).cols_align(
        columns=[c for c in gtfs_table.columns if c != "Operator"],
        align="center",
    ).pipe(
        gte.gt_color_box, 
        columns=["Daily Trips", "# Routes", "# Shapes", "# Stops", 
                 "Daily Arrivals", "Arrivals per Stop"],
        palette="YlGnBu", 
    )
))
except:
    pass