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

import calitp_data_analysis.magics

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

# Display
from great_tables import GT
from IPython.display import HTML, Image, Markdown, display, display_html
from slugify import slugify

import google.auth
credentials, project = google.auth.default()
import gcsfs

fs = gcsfs.GCSFileSystem()

from functools import cache
from calitp_data_analysis.gcs_pandas import GCSPandas

@cache
def gcs_pandas():
    return GCSPandas()

# district = '07-Los Angeles / Ventura'

# Parameters
district = "10-Stockton"

%%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"

transit_route_shs_gdf = (
    utils.load_shn_transit_routes(district = district,
                                 pct = 15,
                                 month = analysis_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 = gpd.read_parquet(fct_monthly_routes_url,
                                        filters=[[("Caltrans District", "==", district)]],
                                        storage_options={"token": credentials.token})

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)
except:
    pass

try:
    display(
    GT(operator_df2
        )
        .tab_header(title=f"District {district} GTFS summary stats")
    )
except:
    pass

color_map = cm.linear.Spectral_11.scale()

try:
    color_map = branca.colormap.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 = branca.colormap.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)
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=f"District {district}",
        subtitle="Daily GTFS schedule statistics for Weekday by operator",
    )
    .cols_align(
        columns=[
            c
            for c in gtfs_table.columns
        ],
        align="center",
    )
))
except:
    pass