%%capture

import altair as alt
import gcsfs
import pandas as pd

from calitp_portfolio import magics
from IPython.display import HTML, Markdown, display
from snapshot_utils import _color_palette, prep_data_utils
from snapshot_utils.project_vars import GCS_FILE_PATH 
from update_vars import min_year

alt.data_transformers.enable("vegafusion")

# Parameters
rtpa = "Stanislaus Council of Governments"

# parameters cell for local
#rtpa = "Metropolitan Transportation Commission"

%%capture_parameters
rtpa, min_year

# TODO: should the columns get subset?
# if publishing all columns, then we can leave it all here
not_published_cols = [
    "key", "legacy_ntd_id", "fta_region",
    "upt_prior_month", "upt_change_1mo", "upt_pct_change_1mo",
] 
df = pd.read_parquet(
    f"{GCS_FILE_PATH}monthly_with_crosswalk.parquet",
    filesystem = gcsfs.GCSFileSystem(),
    filters = [[("rtpa", "==", rtpa)]]
).drop(columns = not_published_cols)

# no rows were found, but just in case
# does this need to affect the aggregations?
#df = df[(df.upt ==0) & (df.upt_change_1yr==0)].reset_index(drop=True)

URL = "https://console.cloud.google.com/storage/browser/calitp-publish-data-analysis"
PUBLIC_FILENAME = df.month_first_day.max().strftime("%Y_%B")

display(
    HTML(
        f"""
        <a href={URL}>
        <b>Download the data: </b> ntd_monthly_ridership/{PUBLIC_FILENAME}.zip</a>
        """
    )
)

# this is total upt since 2018, which is a parameter in the query
# might need to set this in update_vars, otherwise if it updates, 
# we don't know and caption is wrong
# agg by agency
agency_agg_yr = df.pipe(prep_data_utils.proportion_of_upt_by_agency)
total_upt = agency_agg_yr.total_upt.sum()
agency_count = agency_agg_yr.agency.nunique()

Markdown(f"""
Within {rtpa}:
- Number of Reporters: <b>{agency_count}</b>.
- Total Unlinked Passenger Trips since {min_year}: <b>{total_upt:,}</b>.
- Individual Reporters ridership breakdown:
""")

# new chart stuff - keep
# these chart sizes are different than annual
WIDTH = 325
HEIGHT = 150

color_scale = _color_palette.CALITP_CATEGORY_BRIGHT_COLORS + _color_palette.CALITP_CATEGORY_BOLD_COLORS

# https://altair-viz.github.io/gallery/layered_chart_with_dual_axis.html
bar_selection = alt.selection_point(fields=['agency'], bind='legend')

# everything here is shared, only y-axis differs for dual-axis chart
bar_chart_base = (
    alt.Chart(agency_agg_yr)
    .mark_bar()
    .encode(
        x=alt.X("agency", title = "Agency", sort=None, axis=alt.Axis(labelFontSize=8)), 
        # set this sorting to None to favor y sorting
        color=alt.Color(
            "agency", title = "Agency", 
            scale=alt.Scale(range=color_scale),
            legend=alt.Legend(labelFontSize=8)
        ), # default font is 10, so slightly smaller
        tooltip=["agency", "total_upt", "pct_of_total_upt"],
    )
)

chart1 = bar_chart_base.encode(
    y=alt.Y("total_upt", title = "UPT", sort="-y"),
)

chart2 = bar_chart_base.encode(
    y=alt.Y("pct_of_total_upt", title = "Percent", sort="-y"),
)

alt.layer(chart1, chart2).encode(
    opacity=alt.when(bar_selection).then(alt.value(1)).otherwise(alt.value(0.02))
).resolve_scale(
    y = 'independent', x='shared'
).add_params(
    bar_selection
).properties(
    width=WIDTH*1.6, height=HEIGHT*1.2, 
    title={
        "text": "Total Unlinked Passenger Trips (UPT) per Agency in RTPA", 
        "subtitle": f"{min_year} - present"
    }
).interactive()

# Define all shared chart functions here
# annual has reporter_type, remove that for monthly
# tooltip switched for monthly

def title_by_group(group_col: str, y_col: str):
    """
    Set title here for consistency.
    """
    readable_group = group_col.replace("_", " ").replace("_full_name", "").title()

    if y_col=="upt":
        return f"Annual Unlinked Passenger Trips by {readable_group}"

    elif y_col =="upt_change_1yr":
        return  f"Yearly Change in Unlinked Passenger Trips by {readable_group}"

    
def tooltip_by_group(group_col: str): 
    """
    Consistent set of tooltip columns.
    """
    return ["month_first_day", "year", "month", "upt", "upt_change_1yr", group_col, "rtpa"]

CHANGE_SUBTITLE = "Change in UPT from same month, prior year. (Jan 2026 compared to Jan 2025)"

# set scaling to zero=True, removed clamp (not sure what this does)
def make_base_chart(
    df: pd.DataFrame,
    y_col: str,
    color_col: str,
) -> alt.Chart:
    """
    Use 1 base chart function. 
    year is always x-axis, make it ordinal for better display.
    tooltip is standardized with function to populate as much as we can.

    Everything else, such as title, even .mark_line(), .mark_bar() 
    can be layered on top of this function.
    """
    if y_col == "upt_change_1yr":
        y_col_title = "change"
    else:
        y_col_title = y_col
        
    chart = (
        alt.Chart(df)
        .encode(
            x=alt.X("yearmonth(month_first_day):T", title="Date"),    
            y=alt.Y(
                y_col, title=y_col_title, 
                scale=alt.Scale(zero=True)
            ),
            color=alt.Color(
                color_col,
                title="",
                scale=alt.Scale(range=color_scale),
                legend=None
            ),
            tooltip=tooltip_by_group(color_col),
        ).properties(width=WIDTH, height=HEIGHT)
        .interactive()
    )

    return chart

agency_df = pd.read_parquet(
    f"{GCS_FILE_PATH}monthly/agency.parquet",
    filesystem = gcsfs.GCSFileSystem(),
    filters = [[("rtpa", "==", rtpa)]]
)

agency_line = make_base_chart(
    agency_df, 
    y_col = "upt", 
    color_col = "agency"
).mark_line().facet(
    "agency", columns = 2, title = ""
).properties(
    title=title_by_group("agency", "upt"),
).resolve_scale(x="independent", y="independent") 
# independent x-scale helps zooming for specific agencies, esp if we want to focus on a month
# annual report does x='shared'

agency_line

agency_bar = make_base_chart(
    agency_df, 
    y_col="upt_change_1yr",
    color_col = "agency", 
).mark_bar().facet(
    "agency", columns = 2, title = ""
).properties(
    title={
        "text": title_by_group("agency", "upt_change_1yr"), 
        "subtitle": CHANGE_SUBTITLE} ,
).resolve_scale(x="independent", y="independent")

agency_bar

mode_df = pd.read_parquet(
    f"{GCS_FILE_PATH}monthly/mode.parquet",
    filesystem = gcsfs.GCSFileSystem(),
    filters = [[("rtpa", "==", rtpa)]]
)

mode_line = make_base_chart(
    mode_df, 
    y_col = "upt", 
    color_col = "mode_full_name"
).mark_line().facet(
    "mode_full_name", columns = 2, title = ""
).properties(
    title=title_by_group("mode_full_name", "upt"),
).resolve_scale(x="independent", y="independent") 

mode_line

mode_bar = make_base_chart(
    mode_df, 
    y_col="upt_change_1yr",
    color_col = "mode_full_name", 
).mark_bar().facet(
    "mode_full_name", columns = 2, title = ""
).properties(
    title={
        "text": title_by_group("mode_full_name", "upt_change_1yr"), 
        "subtitle": CHANGE_SUBTITLE} ,
).resolve_scale(x="independent", y="independent")

mode_bar

tos_df = pd.read_parquet(
    f"{GCS_FILE_PATH}monthly/type_of_service.parquet",
    filesystem = gcsfs.GCSFileSystem(),
    filters = [[("rtpa", "==", rtpa)]]
)

tos_line = make_base_chart(
    tos_df, 
    y_col = "upt", 
    color_col = "type_of_service_full_name"
).mark_line().facet(
    "type_of_service_full_name", columns = 2, title = ""
).properties(
    title=title_by_group("type_of_service_full_name", "upt"),
).resolve_scale(x="independent", y="independent") 

tos_line

tos_bar = make_base_chart(
    tos_df, 
    y_col="upt_change_1yr",
    color_col = "type_of_service_full_name", 
).mark_bar().facet(
    "type_of_service_full_name", columns = 2, title = ""
).properties(
    title={
        "text": title_by_group("type_of_service_full_name", "upt_change_1yr"), 
        "subtitle": CHANGE_SUBTITLE} ,
).resolve_scale(x="independent", y="independent")

tos_bar