"""
US Treasury Yield Curve Spread — CalcFi Open Data.

Data loaded live from the CalcFi Open Data dataset on Hugging Face.
"""

import gradio as gr
import pandas as pd

DATASET_BASE = (
    "https://huggingface.co/datasets/iizy/calcfi-open-data/resolve/main/datasets"
)

SERIES = [
    {"slug": "federal-funds-rate", "label": "Federal Funds Rate"},
    {"slug": "2-year-treasury", "label": "2-Year Treasury Yield"},
    {"slug": "10-year-treasury", "label": "10-Year Treasury Yield"},
]


def load_series(slug: str, label: str) -> pd.DataFrame:
    url = f"{DATASET_BASE}/{slug}/data.csv"
    df = pd.read_csv(url, comment="#", parse_dates=["date"])
    df["series"] = label
    return df[["date", "value", "series"]]


def load_all() -> pd.DataFrame:
    frames = []
    for s in SERIES:
        try:
            frames.append(load_series(s["slug"], s["label"]))
        except Exception as e:
            print(f"failed to load {s['slug']}: {e}")
    if not frames:
        return pd.DataFrame(columns=["date", "value", "series"])
    return pd.concat(frames, ignore_index=True).sort_values("date")


def filter_view(start_year: int, end_year: int, selected_series: list[str]) -> pd.DataFrame:
    df = load_all()
    if start_year > end_year:
        start_year, end_year = end_year, start_year
    mask = (df["date"].dt.year >= start_year) & (df["date"].dt.year <= end_year)
    df = df[mask]
    if selected_series:
        df = df[df["series"].isin(selected_series)]
    return df


INITIAL_DF = load_all()
MIN_YEAR = int(INITIAL_DF["date"].dt.year.min()) if not INITIAL_DF.empty else 1970
MAX_YEAR = int(INITIAL_DF["date"].dt.year.max()) if not INITIAL_DF.empty else 2026
ALL_LABELS = [s["label"] for s in SERIES]


with gr.Blocks(
    title="US Treasury Yield Curve Spread — CalcFi Open Data",
    theme=gr.themes.Soft(),
) as demo:
    gr.Markdown("# US Treasury Yield Curve Spread")
    gr.Markdown("Visualize the 2-year vs 10-year Treasury yield spread alongside the Federal Funds Rate. When the 10Y drops below the 2Y, the curve inverts — historically a leading indicator of US recessions with a 12-24 month lead time. Released under CC BY 4.0. Source dataset: [calcfi-open-data](https://huggingface.co/datasets/iizy/calcfi-open-data).")

    with gr.Row():
        with gr.Column(scale=1):
            start_year = gr.Slider(
                MIN_YEAR, MAX_YEAR, value=max(MIN_YEAR, MAX_YEAR - 25),
                step=1, label="Start year",
            )
            end_year = gr.Slider(
                MIN_YEAR, MAX_YEAR, value=MAX_YEAR, step=1, label="End year",
            )
            series_picker = gr.CheckboxGroup(
                ALL_LABELS, value=ALL_LABELS, label="Series",
            )

    chart = gr.LinePlot(
        filter_view(max(MIN_YEAR, MAX_YEAR - 25), MAX_YEAR, ALL_LABELS),
        x="date",
        y="value",
        color="series",
        title="US Treasury Yield Curve Spread",
        x_title="Date",
        y_title="Value",
        height=500,
        width=900,
        tooltip=["date", "value", "series"],
    )

    inputs = [start_year, end_year, series_picker]
    for inp in inputs:
        inp.change(filter_view, inputs, chart)


if __name__ == "__main__":
    demo.launch()