%%capture
import warnings
warnings.filterwarnings("ignore")

import calitp_data_analysis.magics

import pandas as pd
import _new_ct_report_utils as utils
from update_vars import GTFS_DATA_DICT, analysis_month, file_name,  previous_month
# Maps
import geopandas as gpd
import branca.colormap as cm
from shared_utils import webmap_utils

# Display
from great_tables import GT
import gt_extras as gte
from IPython.display import display

from functools import cache
from calitp_data_analysis.gcs_pandas import GCSPandas
from calitp_data_analysis.gcs_geopandas import GCSGeoPandas

@cache
def gcs_pandas():
    return GCSPandas()

@cache
def gcs_geopandas():
    return GCSGeoPandas()

# district = '07-Los Angeles / Ventura'

# Parameters
district = "09-Bishop"

%%capture_parameters
district

district_number = int(district[:2])

GCS_PATH = f"{GTFS_DATA_DICT.gcs_paths.DIGEST_GCS}processed/"

# Files for webmaps
boundary_file = f"district_{district_number}_boundary"
transit_routes_file = f"district_{district_number}_transit_routes"
shn_file = f"district_{district_number}_shn"
transit_shn_file = f"district_{district_number}_transit_routes_shn"

try:
    transit_route_shs_gdf = utils.load_shn_transit_routes(
    district = district,
    pct = 15,
    month = analysis_month
    )
except:
    transit_route_shs_gdf = utils.load_shn_transit_routes(
    district = district,
    pct = 15,
    month = previous_month
    )

fct_monthly_routes_url = f"{GCS_PATH}{GTFS_DATA_DICT.gtfs_digest_rollup.route_map}_{file_name}.parquet"

operator_summary_url = f"{GCS_PATH}{GTFS_DATA_DICT.gtfs_digest_rollup.operator_summary}_{file_name}.parquet"

operator_df = gcs_pandas().read_parquet(
    operator_summary_url,
    filters=[
        ("Caltrans District", "==", district),
        ("Date", "==", pd.Timestamp(analysis_month)),
        ("Day Type", "==", "Weekday")
    ],
)

fct_monthly_route_df = gcs_geopandas().read_parquet(
    fct_monthly_routes_url,
    filters=[[("Caltrans District", "==", district)]]
)

fct_monthly_route_df = utils.prep_gdf(fct_monthly_route_df)

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"
    ).drop(columns = "Caltrans District")
except:
    pass

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

color_map = cm.linear.Spectral_11.scale()

try:
    color_map = cm.LinearColormap(
    colors=color_map.colors[7:], vmin=0, vmax=fct_monthly_route_df.Number.max()
)
except:
    pass

try:
    district_map = webmap_utils.set_state_export(
        district_gdf,
        subfolder = "caltrans_district_digest/",
        filename=boundary_file,
        map_title="District Map",
        overwrite = True
    )
except:
    pass

try:
    transit_routes = webmap_utils.set_state_export(
        fct_monthly_route_df,
        subfolder = "caltrans_district_digest/",
        filename=transit_routes_file,
        map_title="Transit Routes",
        cmap=color_map,
        color_col="Number",
        existing_state=district_map,
        overwrite = True    
)
except:
    pass

try:
    webmap_utils.render_spa_link(transit_routes["spa_link"], text="Open Routes for all Operators Map")
except:
    pass

try:
    webmap_utils.display_spa_map(transit_routes["spa_link"])
except:
    pass

color_map2 = cm.linear.RdYlBu_11.scale()

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

try:
    shn_map = webmap_utils.set_state_export(
        shn_gdf,
        subfolder = "caltrans_district_digest/",
        filename=shn_file,
        map_title="State Highway Network Map",
        map_type='state_highway_network',
        overwrite = True
    )
except:
    pass

try:
    transit_shn_map = webmap_utils.set_state_export(
        transit_route_shs_gdf,
        subfolder = "caltrans_district_digest/",
        filename=transit_shn_file,
        map_title="Transit Routes on the State Highway Network",
        cmap=color_map2,
        color_col="Percentage of Transit Route on SHN Across All Districts",
        existing_state=shn_map,
        legend_url="https://storage.googleapis.com/calitp-map-tiles/transit_route_pct.svg",
        overwrite = True    
    )
except:
    pass

try:
    webmap_utils.render_spa_link(transit_shn_map["spa_link"], text="Open Routes on State Highway System Map")
except:
    pass

try:
    webmap_utils.display_spa_map(transit_shn_map["spa_link"])
except:
    pass

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],
    )
))
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