import logging

import numpy as np
import pandas as pd
from xbbg import blp

from tulip.data.dataquery import DataQueryClient as dq
from tulip.plots import plot_lines, switch_trace_to_secondary_axis
from tulip.utils.notebook_related import display_two_charts

loggers = [logging.getLogger(name) for name in logging.root.manager.loggerDict]
for logger in loggers:
    logger.setLevel(logging.CRITICAL + 1)


def style_curve_frame(v):
    if pd.isna(v):
        return "background-color: white;"
    # Gradient from white (0) to black (20)
    v_clipped = max(0, min(20, v))
    grey = int(255 - (v_clipped / 20) * 255)
    bg_color = f"rgb({grey},{grey},{grey})"
    font_color = "white" if v_clipped > 12 else "black"
    if v < 0:
        font_color = "red"
    return f"background-color: {bg_color}; color: {font_color};"


time_in_months_numeric = {
    "Today": 0,
    "1M": 1,
    "3M": 3,
    "6M": 6,
    "1Y": 12,
    "2Y": 24,
    "3Y": 36,
    "5Y": 60,
    "7Y": 84,
    "10Y": 120,
    "20Y": 240,
    "30Y": 360,
}

tenor = "1D"
countries = {
    "USD": "USA",
    "CAD": "CAN",
    "GBP": "GBR",
    "EUR": "EUR",
    "JPY": "JPN",
    "SEK": "SWE",
    "BRL": "BRA",
    "MXN": "MEX",
    "PLN": "POL",
}

curves = {
    "USD": "YCGT0025",
    "CAD": "YCGT0007",
    "GBP": "YCGT0022",
    "EUR": "YCGT0013",
    "JPY": "YCGT0018",
    "SEK": "YCGT0021",
    "BRL": "YCGT0393",
    "MXN": "YCGT0251",
    "PLN": "YCGT0177",
}

overnight_rates = {
    "USD": "SOFRRATE Index",
    "GBP": "SONIO/N Index",
    "EUR": "ESTRON Index",
    "JPY": "MUTSCALM Index",
    "CAD": "CAONREPO Index",
    "SEK": "STIB1D Index",
    "BRL": "BZSELICA Index",
    "MXN": "MXIBTIEF Index",
    "PLN": "WIRON Index",
}

today = (
    blp.bdp(tickers=[overnight_rates[x] for x in countries.keys()], flds=["PX_LAST"])
    .rename(index={v: k for k, v in overnight_rates.items()})
    .squeeze()
    .rename("Today")
)

wpi = {}
curve = {}

for ccy, ctry in countries.items():
    try:
        wpi[ccy] = dq.wpi(ccy, tenor=tenor)
    except Exception as e:
        print(f"Warning: Could not fetch WPI for {ccy}: {e}")
    curve[ccy] = blp.bds(tickers=f"{curves[ccy]} Index", flds=["CURVE_TENOR_RATES"])

curve_frame = pd.concat(curve)
curve_frame.index.names = ["ccy", "curve"]
curve_frame = (
    curve_frame.droplevel("curve")
    .set_index("tenor", append=True)["mid_yield"]
    .unstack("tenor")
)
month_cols = sorted(
    [c for c in curve_frame.columns if c.endswith("M")], key=lambda x: int(x[:-1])
)
year_cols = sorted(
    [c for c in curve_frame.columns if c.endswith("Y")], key=lambda x: int(x[:-1])
)
curve_frame = curve_frame[month_cols + year_cols]
curve_frame["Today"] = today

curve_frame = curve_frame.loc[
    countries.keys(),
    ["Today", "1M", "3M", "6M", "1Y", "2Y", "3Y", "5Y", "7Y", "10Y", "20Y", "30Y"],
]

styled_curve_frame = curve_frame.style.format("{:.2f}").map(style_curve_frame)
discounting_df = (
    curve_frame.rename(columns=time_in_months_numeric)
    .T.interpolate("slinear")
    .round(3)
    .T
)


def plot_discounting(curve_df, title="<b>Easing/Tightenings Priced In The Curve<b>"):
    fig = plot_lines(
        curve_df.sub(curve_df.iloc[:, 0], axis=0).T,
        title=title,
        show_0=True,
    )

    y_range = fig.layout.yaxis.range or [None, None]
    if y_range[0] is not None and y_range[1] is not None:
        y_ticks = np.arange(
            np.floor(y_range[0] * 2) / 2, np.ceil(y_range[1] * 2) / 2, 0.5
        )
        fig.update_yaxes(tickvals=y_ticks)
    else:
        y_ticks = fig.layout.yaxis.tickvals

    switch_trace_to_secondary_axis(fig, "USD")

    if 360 in curve_df.columns:
        x_ticks = ["" if k.endswith("M") else k for k in time_in_months_numeric.keys()]
    else:
        x_ticks = list(time_in_months_numeric.keys())

    fig.update_layout(
        yaxis=dict(
            tickvals=y_ticks,
            ticktext=[f"{val * 100:.0f}" for val in y_ticks],
            title="Absolute change (bps)",
        ),
        yaxis2=dict(
            anchor="y",
            range=y_range,
            overlaying="y",
            side="right",
            title="# of 25bps Hikes/Cuts",
            tickvals=y_ticks,
            ticktext=[f"{val / 0.25:.0f}" for val in y_ticks],
        ),
        xaxis=dict(
            title="Tenor",
            tickmode="array",
            tickvals=list(time_in_months_numeric.values()),
            ticktext=x_ticks,
        ),
    )
    return fig


fig_all = plot_discounting(discounting_df)
fig_sub_1y = plot_discounting(discounting_df.loc[:, discounting_df.columns <= 12])

import matplotlib.pyplot as plt
import math


def plot_forward_curve(df):
    # Extract metadata for title
    date = (
        df.index[0].strftime("%Y-%m-%d")
        if hasattr(df.index[0], "strftime")
        else str(df.index[0])
    )
    ccy = df["currency"].iloc[0]
    fwd = df["forward_life"].iloc[0]

    # Get numeric columns only
    numeric_cols = df.select_dtypes(include="number").columns
    yields = df[numeric_cols].iloc[0]

    # Plot
    fig, ax = plt.subplots(figsize=(10, 5))
    bars = ax.bar(yields.index, yields.values, color="black")
    ax.set_ylabel("Yield (%)")
    ax.set_title(f"$\\mathbf{{{ccy}}}$ {fwd} Forward Curve - {date}")
    ax.yaxis.grid(True, linestyle="--", alpha=0.7)
    ax.set_axisbelow(True)
    ax.set_facecolor("white")
    fig.patch.set_facecolor("white")

    # Remove top and right spines
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)

    # Set y-axis upper limit to nearest integer
    ax.set_ylim(0, math.ceil(yields.max()))

    # Add value labels on top of bars
    for bar, val in zip(bars, yields.values):
        ax.text(
            bar.get_x() + bar.get_width() / 2,
            bar.get_height() + 0.05,
            f"{val:.2f}",
            ha="center",
            va="bottom",
            fontsize=10,
            color="dimgrey",
        )

    plt.tight_layout()
    return fig, ax