%%capture

import warnings
warnings.filterwarnings('ignore')

import altair as alt
import gcsfs
import pandas as pd

from calitp_portfolio import magics
from snapshot_utils import prep_data_utils, _color_palette
from snapshot_utils.project_vars import GCS_FILE_PATH

alt.data_transformers.enable("vegafusion")

#parameters cell
#rtpa = "Sacramento Area Council of Governments"

# Parameters
rtpa = "San Luis Obispo Council of Governments"

%%capture_parameters
rtpa

df = pd.read_parquet(
    f"{GCS_FILE_PATH}annual.parquet",
    # should only certain columns be read in? now this table is much larger
    filesystem=gcsfs.GCSFileSystem(),
    columns = [
        "ntd_id", "source_agency", "agency_status", "source_city", 
        "year",
        "mode", "mode_full_name", "type_of_service", "type_of_service_full_name",
        "reporter_type", "reporting_module", "source_state", "primary_uza_name",
        "unlinked_passenger_trips", "vehicle_revenue_hours", "vehicle_revenue_miles",
        "operating_expenses_total",
        "opex_per_vrh", "opex_per_vrm", "opex_per_upt", 
        "upt_per_vrh", "upt_per_vrm",
        "farebox_recovery_ratio", "fare_revenue",
    ]
).pipe(
    prep_data_utils.merge_with_crosswalk
).query(
    f'rtpa_name == "{rtpa}"'
).dropna(
    subset="unlinked_passenger_trips"
).rename(columns = {"source_agency": "agency"})

agency_df = prep_data_utils.calculate_efficiency_metrics_by_group(
    df, 
    ["ntd_id", "agency", "year", "rtpa_name"]
)

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

WIDTH = 400
HEIGHT = 250

def readable(word: str) -> str:
    """
    Coerce words that are abbreviated into readable labels for viz.
    """
    ABBREV_TO_FULL = {
        "opex_per_upt": "Operating Expense per Unlinked Passenger Trip",
        "opex_per_vrh": "Operating Expense per Vehicle Revenue Hour",
        "opex_per_vrm": "Operating Expense per Vehicle Revenue Mile",
        "upt_per_vrh": "Unlinked Passenger Trips per Vehicle Revenue Hour",
        "upt_per_vrm": "Unlinked Passenger Trips per Vehicle Revenue Mile",
    }

    if word in ABBREV_TO_FULL.keys():
        word = ABBREV_TO_FULL[word]

    return word.replace("_full_name", "").replace("_", " ").title().replace("Of", "of")

# Define all shared chart functions here
def title_by_group(
    group_col: str, 
    x_col: str = None,
    y_col: str = None,
):
    """
    Set title here for consistency.
    """
    readable_group = readable(group_col)
    
    if group_col == "reporter_type":
        readable_group = f"NTD {readable_group}"

    # For scatterplots
    if x_col and y_col:
        return f"Scatter: {readable(y_col)} by {readable(x_col)}"
    # For line plots by year
    if not x_col and y_col: 
        return f"{readable(y_col)} by {readable_group}"
    else:
        return readable_group
        
def tooltip_by_group(group_col: str): 
    """
    Consistent set of tooltip columns.
    """
    return ["year", group_col, "opex_per_upt", "opex_per_vrh", "opex_per_vrm", "upt_per_vrh", "upt_per_vrm", "rtpa_name"]

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.

    This one is similar to annual report.
    """
    selection = alt.selection_point(fields=[color_col], bind='legend')

    chart = (
        alt.Chart(df)
        .encode(
            x=alt.X("year:O"),    
            y=alt.Y(
                y_col, title=y_col, 
                scale=alt.Scale(zero=False, clamp=True)
            ),
            color=alt.Color(
                color_col,
                scale=alt.Scale(range=color_scale),
                #legend=None
            ),
            opacity=alt.when(selection).then(alt.value(1)).otherwise(alt.value(0.02)),
            tooltip=tooltip_by_group(color_col),
        ).properties(width=WIDTH, height=HEIGHT)
        .interactive()
    ).add_params(selection)

    return chart

def combined_line_charts_by_year(
    df, 
    color_col: str
):
    """
    Cost-efficiency and service-effectiveness charts
    can be titled, labeled separately, but ultimately,
    want to display as 2 groups.
    Group them together here, and add the legend selector again,
    it loses the legend selector by the end.
    """
    cost_efficiency_chart_list = [
        make_base_chart(
            df, 
            y_col = m,
            color_col = color_col
        ).mark_line(
            point=alt.OverlayMarkDef(filled=False, fill="white")
        ).properties(
            title=title_by_group(color_col, x_col= None, y_col = m)
        )
        for m in cost_efficiency_metrics
    ]

    # Loop over service-effectiveness metrics and create charts by agency
    service_effectiveness_chart_list = [
        make_base_chart(
            df, 
            y_col = m,
            color_col = color_col
        ).mark_line(
            point=alt.OverlayMarkDef(filled=False, fill="white")
        ).properties(
            title=title_by_group(color_col, x_col= None, y_col = m)
        )
        for m in service_effectiveness_metrics
    ]

    selection = alt.selection_point(fields=[color_col], bind='legend')
    
    chart1 = alt.hconcat(*cost_efficiency_chart_list).properties(
        title=f"Cost Efficiency by {readable(color_col)}"
    )
    
    chart2 = alt.hconcat(*service_effectiveness_chart_list).properties(
        title=f"Service Effectiveness by {readable(color_col)}"
    )

    combined_chart = alt.vconcat(chart1, chart2).add_params(selection) # add legend selector here again

    return combined_chart 

# Loop over the cost-efficiency metrics and create charts by agency
cost_efficiency_metrics = ["opex_per_vrh", "opex_per_vrm", "opex_per_upt"]
service_effectiveness_metrics = ["upt_per_vrh", "upt_per_vrm"]

combined_line_charts_by_year(agency_df, "agency")

mode_df = prep_data_utils.calculate_efficiency_metrics_by_group(
    df, 
    ["mode", "mode_full_name", "year", "rtpa_name"]
)

combined_line_charts_by_year(mode_df, "mode_full_name")

# scatterplot with all the years and agencies?
# The regression lines might indicate marginal cost of providing that service
# steeper lines show higher marginal cost?
def make_scatterplot_with_line(
    df: pd.DataFrame,
    x_col: str,
    y_col: str,
    color_col: str,
) -> alt.Chart:
    """
    Add scatterplot with regression line.
    Since scatterplot is log-log, the regression line we use 
    can be linear.
    Only upt is not log(upt).
    If we did regular_x, regular_y, we would have to use power curve?
    """
    scatter_selection = alt.selection_point(fields=[color_col], bind='legend')

    scatter = alt.Chart(df).mark_point().encode(
        y=alt.Y(y_col, title = f"log({readable(y_col)})").scale(type="log"),
        color=alt.Color(
            color_col, 
            title=readable(color_col),
            scale=alt.Scale(range=color_scale)
        ),
        tooltip = [color_col, "year", "ntd_id",
                   x_col, y_col, "rtpa_name"],
        opacity=alt.when(scatter_selection).then(alt.value(1)).otherwise(alt.value(0.02)),
    )

    # Be explicit here around which ones need log scale (exceptions listed out, just upt)
    if x_col == "unlinked_passenger_trips":
        scatter = scatter.encode(
            x=alt.X(x_col, title=readable(x_col))
        )
    else:
        scatter = scatter.encode(
            x=alt.X(x_col, title=f"log({readable(x_col)})").scale(type="log")
        )
    
    scatter_line = scatter.transform_regression(
        regression=x_col, 
        on = y_col,
        method="linear",
        groupby = [color_col] # this will draw individual lines for this group
    ).mark_line()

    chart = (scatter_line + scatter).properties(
        width=WIDTH, height=HEIGHT,
        title = title_by_group(color_col, x_col, y_col)
    ).add_params(scatter_selection)
    
    return chart

cost_efficiency_scatter_list = [
    make_scatterplot_with_line(
        df, 
        m, 
        "operating_expenses_total", 
        "mode_full_name",
    ) 
    for m in ["unlinked_passenger_trips", "vehicle_revenue_hours", "vehicle_revenue_miles"]
]

cost_efficiency_scatter = alt.hconcat(*cost_efficiency_scatter_list).resolve_scale(
    x='independent', y='shared'
)

service_effectiveness_scatter_list = [
    make_scatterplot_with_line(
        df, 
        m, 
        "unlinked_passenger_trips", 
        "mode_full_name",
    ) 
    for m in ["vehicle_revenue_hours", "vehicle_revenue_miles"]
]


service_effectiveness_scatter = alt.hconcat(
    *service_effectiveness_scatter_list
).resolve_scale(
    x='independent', y='shared'
)

alt.vconcat(cost_efficiency_scatter, service_effectiveness_scatter).interactive()

tos_df = prep_data_utils.calculate_efficiency_metrics_by_group(
    df, 
    ["type_of_service", "type_of_service_full_name", "year", "rtpa_name",]
)

combined_line_charts_by_year(tos_df, "type_of_service_full_name")

cost_efficiency_scatter_list = [
    make_scatterplot_with_line(
        df, 
        m, 
        "operating_expenses_total", 
        "type_of_service_full_name",
    )
    for m in ["vehicle_revenue_hours", "vehicle_revenue_miles"]
]

cost_efficiency_scatter = alt.hconcat(*cost_efficiency_scatter_list).resolve_scale(
    x='independent', y='shared'
)

service_effectiveness_scatter_list = [
    make_scatterplot_with_line(
        df, 
        m, 
        "unlinked_passenger_trips", 
        "type_of_service_full_name",
    )
    for m in ["vehicle_revenue_hours", "vehicle_revenue_miles"]
]


service_effectiveness_scatter = alt.hconcat(*service_effectiveness_scatter_list).resolve_scale(
    x='independent', y='shared'
)

alt.vconcat(cost_efficiency_scatter, service_effectiveness_scatter)