%%capture

import altair as alt
import gcsfs
import pandas as pd

from IPython.display import HTML, Markdown, display
from snapshot_utils import _color_palette, prep_data_utils
from calitp_portfolio import magics

from snapshot_utils.project_vars import GCS_FILE_PATH 
from update_vars import min_year

alt.data_transformers.enable("vegafusion")

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

# Parameters
rtpa = "Los Angeles County Department of Public Works"

%%capture_parameters
rtpa, min_year

URL = "https://console.cloud.google.com/storage/browser/calitp-publish-data-analysis"

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

# read in data
df = pd.read_parquet(
    f"{GCS_FILE_PATH}annual_with_crosswalk.parquet",
    filesystem=gcsfs.GCSFileSystem(),
    filters = [[("rtpa", "==", rtpa)]],
    columns = [
        "ntd_id", "agency", "year", 
        "mode", "mode_full_name", 
        "type_of_service", "type_of_service_full_name",
        "upt", "upt_prior_year", "upt_change_1yr", "upt_pct_change_1yr",
        "agency_status",
        "rtpa"
    ]
)

# 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
def proportion_of_upt_by_agency(df: pd.DataFrame):
    initial_agg = (
        df
        .groupby("agency")
        .agg(
            total_upt=("upt", "sum")
        ).reset_index()
        .astype({"total_upt": "int64"})
        .sort_values(by="total_upt", ascending=False)
    )
     # % total columns
    initial_agg["pct_of_total_upt"] = ((
        initial_agg["total_upt"] / initial_agg["total_upt"].sum()
    ) * 100).round(decimals=2)

    return initial_agg

# agg by agency
agency_agg_yr = df.pipe(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 the beginning of this report: <b>{total_upt:,}</b>.
- Individual Reporters ridership breakdown:
"""
)

# set some chart variables
color_scale = _color_palette.CALITP_CATEGORY_BRIGHT_COLORS + _color_palette.CALITP_CATEGORY_BOLD_COLORS

WIDTH = 300
HEIGHT = 150

# 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

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 group_col == "reporter_type":
        readable_group = f"NTD {readable_group}"
    
    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 ["year", "upt", "upt_change_1yr", group_col, "rtpa"]

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("year:O"),    
            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}annual/agency.parquet",
    filesystem = gcsfs.GCSFileSystem(),
    filters = [[("rtpa", "==", rtpa)]]
)

agency_line = make_base_chart(
    agency_df,
    y_col="upt",
    color_col="agency",
).mark_line(point=True).facet(
    "agency", columns = 2, title = ""
).properties(
    title=title_by_group("agency", "upt"),
).resolve_scale(x="shared", y="independent")

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=title_by_group("agency", "upt_change_1yr"),
).resolve_scale(x="shared", y="independent")

agency_bar

mode_df = pd.read_parquet(
    f"{GCS_FILE_PATH}annual/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(point=True).facet(
    "mode_full_name", columns = 2, title = ""
).properties(
    title=title_by_group("mode_full_name", "upt"),
).resolve_scale(x="shared", 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=title_by_group("mode_full_name", "upt_change_1yr"),
).resolve_scale(x="shared", y="independent")

mode_bar

tos_df = pd.read_parquet(
    f"{GCS_FILE_PATH}annual/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(point=True).facet(
    "type_of_service_full_name", columns = 2, title = ""
).properties(
    title=title_by_group("type_of_service_full_name", "upt"),
).resolve_scale(x="shared", 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=title_by_group("type_of_service_full_name", "upt_change_1yr"),
).resolve_scale(x="shared", y="independent")

tos_bar

reporter_type_df = pd.read_parquet(
    f"{GCS_FILE_PATH}annual/reporter_type.parquet",
    filesystem = gcsfs.GCSFileSystem(),
    filters = [[("rtpa", "==", rtpa)]]
)

reporter_line = make_base_chart(
    reporter_type_df,
    y_col="upt",
    color_col="reporter_type",
).mark_line(point=True).facet(
    "reporter_type", columns = 2, title = ""
).properties(
    title=title_by_group("reporter_type", "upt"),
).resolve_scale(x="shared", y="independent")

reporter_line

reporter_bar = make_base_chart(
    reporter_type_df, 
    y_col="upt_change_1yr",
    color_col = "reporter_type", 
).mark_bar().facet(
    "reporter_type", columns = 2, title = ""
).properties(
    title=title_by_group("reporter_type", "upt_change_1yr"),
).resolve_scale(x="shared", y="independent")

reporter_bar