@register_hack("dataset", "us", "vehicles_json")
def add_token_dollar_to_investments(vehicles_df: pd.DataFrame) -> pd.DataFrame:
    """
    Adding a token dollar to all vehicles investments in order to get optimization results
    for vehicles with missing invesmtents for the years 2023, 2024.
    """
    vehicles_df.loc[vehicles_df.index.get_level_values("week_date").year.isin([2023, 2024])] += 1 / 52
    return vehicles_df

@register_hack("dataset", "us", "vehicles_json")
def fix_nba_mul_sponsorship_spread(vehicles_df: pd.DataFrame) -> pd.DataFrame:
    """
    We believe that the percentage spread of spends across calendar year 2021 should be the same as
    it was in 2023.
    """
    if ("MUL" not in vehicles_df.index.unique("brand_code")) or (
        vehicles_df.index.unique("week_date").min().year > 2021
    ):
        return vehicles_df
    league_brand = vehicles_df.xs("MUL", level="brand_code")[["league_NBA"]].droplevel("product_code")
    league_brand_year = league_brand[league_brand.index.year == 2021]
    league_brand_correct_year = league_brand[league_brand.index.year == 2023]
    corrected_league_brand_year = (
        league_brand_correct_year / league_brand_correct_year.sum()
    ) * league_brand_year.iloc[: len(league_brand_correct_year)].sum()
    truncated_index = league_brand_year.index[: len(corrected_league_brand_year)]
    vehicles_df.loc[
        pd.IndexSlice[:, "MUL", truncated_index.tolist()], "league_NBA"
    ] = corrected_league_brand_year.values
    return vehicles_df