%%capture

import warnings

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

import branca.colormap as cm
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, analysis_month

credentials, _ = google.auth.default()

# parameters cell
#district = '07 - Los Angeles / Ventura'

# Parameters
district = "06 - Fresno / Bakersfield"

%%capture_parameters
district_number = int(district[:2])
district, district_number

transit_route_shs_gdf = utils.load_shn_transit_routes(
    district = district,
    pct = 15,
)

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

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

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

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,
    filters=[[("Analysis Name", "in", list_of_operators)]],
    storage_options = {"token": credentials}
).reset_index(drop=True).pipe(
    utils.prep_gdf
)

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=[
        ("Caltrans District", "==", district),
        ("Date", "==", pd.to_datetime(analysis_month)),
        ("Day Type", "==", "Weekday")
    ],
)

district_gdf = utils.load_ct_district(district_number)

shn_gdf =  utils.load_buffered_shn_map(district_number)

try:
    operator_df2 = utils.create_summary_table(
        operator_df, district_col = "Caltrans 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 = "Caltrans 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 Caltrans 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

color_map2 = cm.linear.RdYlBu_11.scale()
color_map2 = cm.LinearColormap(
    colors=color_map2.colors[7:], vmin=0, vmax=100
)

# rgb to hex (155, 155, 155)
shn_map = shn_gdf.explore(
    color = "#9b9b9b",
    tiles = "CartoDB Positron",
    name = "State Highway Network",
    legend=True,
    style_kwds = {"weight": 1, "fillOpacity": 0},
    highlight_kwds = {"fillOpacity": 0.5},
    tooltip=True
)

transit_route_shs_gdf.explore(
    "Percentage of Transit Route on SHN Across All Districts",
    m=shn_map,
    name="Percentage of Transit Route on SHN Across All Districts",
    cmap=color_map2,
    tooltip=True,
)

folium.LayerControl().add_to(shn_map)
shn_map

try:
    display(
        GT(transit_route_shs_gdf.drop(columns = ["geometry"]).sort_values(
            by=[
            "Analysis Name",
            "Percentage of Transit Route on SHN Across All Districts",
        ], ascending=[True, False])
    ).cols_move_to_start(
        ["Analysis Name"]
    ).tab_stub(
        groupname_col="Analysis Name" # creates groups within table, routes grouped by operator
    )
    .pipe(
        gte.gt_color_box,
        columns=["Percentage of Transit Route on SHN Across All Districts"],
        palette="YlGnBu", 
    )    
    ).cols_width(cases={
        "Route Name": "40%",
        "State Highway Network Route": "30%",
    }).tab_options(
        table_font_size="11px", 
    )
    #.pipe(gte.gt_plt_dot, category_col="Analysis Name", 
    #      data_col="Percentage of Transit Route on SHN Across All Districts", domain=[0, 0])
    # plt_dot adds dot ahead of Analysis Name, but spacing gets too crowded
except:
    pass

try:
    gtfs_table = utils.create_operator_table(
        operator_df, district_col = "Caltrans District"
    ).drop(columns = "Caltrans 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