countries = [
    "US",  # United States
    "EU",  # Euro
    "GB",  # United Kingdom
    "JP",  # Japan
    "SE",  # Sweden
    "AU",  # Australia
    "CA",  # Canada
    "KR",  # South Korea
    "CN",  # China
    "IN",  # India
    "BR",  # Brazil
    # "DE",  # Germany
    # "FR",  # France
    "MX",  # Mexico
    "ZA",  #  South Africa
    # "ES",  # Spain
    # "IT",  # Italy
]

today_str = pd.Timestamp("today").strftime("%Y%m%d")
week_ago_str = (pd.Timestamp("today") - pd.offsets.BusinessDay(n=5)).strftime("%Y%m%d")
wirp_period = (pd.Timestamp("today") + pd.offsets.YearEnd()).strftime("%b%Y").upper()

bonds_markets = {
    k: {
        "pycountry": pycountry.countries.get(alpha_2=k),
        "currency": COUNTRY_CURRENCY[k],
        "WIRP": WIRP[k] if k in WIRP.keys() else None,
        "Hikes/Cuts ticker": (
            WIRP[k].replace("BPR", "BNM") if k in WIRP.keys() else None
        ),
        "Implied Rate ticker": (
            WIRP[k].replace("BPR", "BFR") if k in WIRP.keys() else None
        ),
        "Hikes/Cuts by EoY ticker": (
            WIRP[k].replace("BPR", "BNM").replace(" ", " " + wirp_period + " ")
            if k in WIRP.keys()
            else None
        ),
        "Implied Rate by EoY ticker": (
            WIRP[k].replace("BPR", "BFR").replace(" ", " " + wirp_period + " ")
            if k in WIRP.keys()
            else None
        ),
        "policy_rate_ticker": POLICY_RATES[k],
        "policy_rate": bb.bdp(POLICY_RATES[k], fields=["PX_LAST"])["px_last"].squeeze(),
        "overnight_rate_ticker": OVERNIGHT_RATES[k],
        "overnight_rate": bb.bdp(OVERNIGHT_RATES[k], fields=["PX_LAST"])[
            "px_last"
        ].squeeze(),
        "sovereign_curve_ticker": f"{CURVES['sovereign'][k]} Index",
        "sovereign_curve": blp.bds(
            f"{CURVES['sovereign'][k]} Index", flds=["CURVE_TENOR_RATES"]
        ),
        "sovereign_curve_1_week_ago": blp.bds(
            f"{CURVES['sovereign'][k]} Index",
            flds=["CURVE_TENOR_RATES"],
            curve_date=week_ago_str,
        ),
        #  "growth_ticker_QoQ":bb.bdp(f"EHGD{k} Index", fields=['PX_LAST'])['px_last'].squeeze(),
    }
    for k in countries
}

for ctry in countries:
    try:
        bonds_markets[ctry]["Hikes/Cuts by EoY"] = bb.bdp(
            bonds_markets[ctry]["Hikes/Cuts by EoY ticker"], fields=["PX_LAST"]
        )["px_last"].squeeze()
    except:
        bonds_markets[ctry]["Hikes/Cuts by EoY"] = None
    try:
        bonds_markets[ctry]["Implied Rate by EoY"] = bb.bdp(
            bonds_markets[ctry]["Implied Rate by EoY ticker"], fields=["PX_LAST"]
        )["px_last"].squeeze()
    except:
        bonds_markets[ctry]["Implied Rate by EoY"] = None

for ctry in countries:
    bonds_markets[ctry]["central_bank_wip"] = pd.Series(
        {
            "Policy Rate": bonds_markets[ctry]["policy_rate"],
            "Hike/Cuts by EOY": bonds_markets[ctry]["Hikes/Cuts by EoY"],
            "Implied Rate by EOY": bonds_markets[ctry]["Implied Rate by EoY"],
        },
        name=f"{ctry} CB expectations",
    )

for ctry in countries:
    bonds_markets[ctry]["central_bank_wip_styled"] = (
        bonds_markets[ctry]["central_bank_wip"].to_frame().style.format("{:.3f}")
    )

    title = (
        f"<b>{bonds_markets[ctry]['currency']} Sovereign Curve</b>"
        if bonds_markets[ctry]["pycountry"] is None
        else f"<b>{bonds_markets[ctry]['pycountry'].name} Sovereign Curve</b>"
    )

    bonds_markets[ctry]["curve_data"] = curve = (
        pd.DataFrame(
            {
                "Today": bonds_markets[ctry]["sovereign_curve"].set_index("tenor")[
                    "mid_yield"
                ],
                "1W ago": bonds_markets[ctry]["sovereign_curve_1_week_ago"].set_index(
                    "tenor"
                )["mid_yield"],
            }
        )
        .round(4)
        .replace(0, np.nan)
    )

    bonds_markets[ctry]["curve_fig"] = plot_curve(curve, tick_suffix="%", title=title)
    bonds_markets[ctry]["curve_fig"].add_hline(
        bonds_markets[ctry]["overnight_rate"],
        line_color="darkred",
        line_dash="dash",
        layer="above",
        opacity=0.8,
        line_width=1,
        annotation_text="Overnight Rate",
    )

overnight_rates = pd.Series({k: v["overnight_rate"] for k, v in bonds_markets.items()})

def transform_tenor(x):
    if x < 12:
        return ("Months", int(x))
    else:
        return ("Years", int(x / 12))


def style_scan(df):
    def hide_nan_css(v):
        return "visibility: hidden;" if pd.isna(v) else ""

    styler = df.style.background_gradient(axis=0)
    styler = styler.map(hide_nan_css).format(na_rep="", precision=2)
    return styler


todays_yields = pd.concat(
    {k: v["curve_data"]["Today"] for k, v in bonds_markets.items()}, axis=1
).sort_index()
todays_yields.index = todays_yields.index.map(lambda x: round(x, 2))
todays_yields = todays_yields.drop(1.38)

todays_yields.index = pd.MultiIndex.from_tuples(
    [transform_tenor(x) for x in todays_yields.index], names=["Period", "Value"]
)

style_scan(todays_yields.T)

spreadsto1s = todays_yields.T.loc[:, pd.IndexSlice["Years", [10, 20, 30]]].sub(
    todays_yields.T.loc[:, pd.IndexSlice["Years", 1]].squeeze(), axis=0
)
spreadsto1s.columns = pd.MultiIndex.from_product(
    [["Spreads1sXs"], [10, 20, 30]], names=spreadsto1s.columns.names
)
spreadsto2s = todays_yields.T.loc[:, pd.IndexSlice["Years", [10, 20, 30]]].sub(
    todays_yields.T.loc[:, pd.IndexSlice["Years", 2]].squeeze(), axis=0
)
spreadsto2s.columns = pd.MultiIndex.from_product(
    [["Spreads2sXs"], [10, 20, 30]], names=spreadsto2s.columns.names
)
style_scan(pd.concat([spreadsto1s, spreadsto2s], axis=1))

for ctry in countries:
    display_table_and_chart(
        table=bonds_markets[ctry]["central_bank_wip_styled"],
        chart=bonds_markets[ctry]["curve_fig"],
    )

from tulip_mania.notebook_related import notebook_updated

notebook_updated()