"""Reproduce the US rates portfolio exhibits from the bundled FRED snapshots.

Run: python analysis.py
No network calls. Source CSVs are immutable inputs; chart/table outputs may be
regenerated. Rates in input files are percentages, not decimal yields.
"""

from pathlib import Path
import hashlib
import json

import matplotlib
matplotlib.use("Agg")
import matplotlib.dates as mdates
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd


ROOT = Path(__file__).resolve().parent
SERIES = ("DGS2", "DGS10", "DFII10", "SOFR", "IORB")
DATES = ("2024-06-28", "2024-09-16", "2024-12-31")
INK, BLUE, RUST = "#272c30", "#315e78", "#a45231"


def load_rates() -> pd.DataFrame:
    """Check provenance, schema and ordering; retain missing observations."""
    manifest = json.loads((ROOT / "data/manifest.json").read_text())
    frames = []
    for name in SERIES:
        path = ROOT / "data" / f"{name}.csv"
        digest = hashlib.sha256(path.read_bytes()).hexdigest()
        if digest != manifest["series"][name]["sha256"]:
            raise ValueError(f"Source checksum changed: {name}")
        frame = pd.read_csv(path, na_values=["."])
        if list(frame.columns) != ["observation_date", name]:
            raise ValueError(f"Unexpected columns: {name}")
        frame["observation_date"] = pd.to_datetime(
            frame["observation_date"], format="%Y-%m-%d", errors="raise"
        )
        frame[name] = pd.to_numeric(frame[name], errors="raise")
        frame = frame.set_index("observation_date")
        if frame.index.has_duplicates or not frame.index.is_monotonic_increasing:
            raise ValueError(f"Duplicate or unsorted dates: {name}")
        if not frame.index.to_series().between("2024-01-01", "2024-12-31").all():
            raise ValueError(f"Observation outside the declared sample: {name}")
        if not np.isfinite(frame[name].dropna()).all() or frame[name].notna().sum() == 0:
            raise ValueError(f"Invalid or empty series: {name}")
        frames.append(frame)
    return pd.concat(frames, axis=1, sort=True).sort_index()


def dollar_dv01(market_value: float, modified_duration: float) -> float:
    """Positive dollar sensitivity per 1 bp; long-bond P&L has the opposite sign."""
    if market_value <= 0 or modified_duration <= 0:
        raise ValueError("Market value and duration must be positive")
    return market_value * modified_duration * 0.0001


def carry_roll_bps(yield_pct, funding_pct, horizon_years, duration,
                   roll_yield_bp, shock_yield_bp=0.0):
    """First-order near-par excess return, in bp of starting market value.

    Income uses an annualized yield proxy. Negative roll_yield_bp means a lower
    yield at the shorter maturity. No convexity, haircut, or execution costs.
    """
    income = (yield_pct - funding_pct) * 100 * horizon_years
    roll = -duration * roll_yield_bp
    shock = -duration * shock_yield_bp
    return income, roll, shock, income + roll + shock


def spread_excess_bps(spread_bp, extra_funding_bp, horizon_years,
                      spread_duration, spread_change_bp, round_trip_cost_bp):
    """Stylized Treasury-hedged bullet, normalized to initial market value."""
    return ((spread_bp - extra_funding_bp) * horizon_years
            - spread_duration * spread_change_bp - round_trip_cost_bp)


def verify_financial_units():
    """Guard against the sign and percent/bp mistakes that change conclusions."""
    assert np.isclose(dollar_dv01(10_000_000, 1.9), 1900)
    assert np.isclose(carry_roll_bps(4.2, 4.6, .25, 4.4, -10)[-1], 34)
    assert np.isclose(carry_roll_bps(4.2, 5.6, .25, 4.4, -10)[-1], 9)
    assert np.isclose(carry_roll_bps(4.2, 4.6, .25, 4.4, -10, 25)[-1], -76)
    assert np.isclose(spread_excess_bps(40, 10, .25, 4.5, 10, 6), -43.5)


def save_table(name, frame, index=False):
    frame.to_csv(ROOT / "tables" / f"{name}.csv", index=index, float_format="%.6f")


def style_axis(ax, ylabel):
    ax.set_ylabel(ylabel)
    ax.grid(axis="y", color="#deddd7", linewidth=.6)
    ax.spines[["top", "right"]].set_visible(False)
    ax.spines[["left", "bottom"]].set_color("#bcbcb4")
    ax.margins(x=.01)


def finish_chart(fig, axes, filename, source):
    locator = mdates.AutoDateLocator(minticks=4, maxticks=7)
    axes[-1].xaxis.set_major_locator(locator)
    axes[-1].xaxis.set_major_formatter(mdates.ConciseDateFormatter(locator))
    fig.subplots_adjust(left=.12, right=.97, top=.88, bottom=.15, hspace=.35)
    fig.text(.12, .035, source, fontsize=8, color="#555b5d")
    for ext in ("svg", "png"):
        fig.savefig(ROOT / "charts" / f"{filename}.{ext}", dpi=180,
                    facecolor="#faf9f5")
    plt.close(fig)


def main():
    (ROOT / "charts").mkdir(exist_ok=True)
    (ROOT / "tables").mkdir(exist_ok=True)
    verify_financial_units()
    raw = load_rates()
    save_table("aligned-rates", raw, index=True)

    # Align each exhibit only on its required inputs. Do not forward-fill.
    curve = raw[["DGS2", "DGS10"]].dropna().copy()
    curve["slope_bp"] = (curve.DGS10 - curve.DGS2) * 100
    tips = raw[["DGS10", "DFII10"]].dropna().copy()
    tips["breakeven_pct"] = tips.DGS10 - tips.DFII10
    funding = raw[["SOFR", "IORB"]].dropna().copy()
    funding["sofr_minus_iorb_bp"] = (funding.SOFR - funding.IORB) * 100
    assert np.allclose(tips.DGS10, tips.DFII10 + tips.breakeven_pct)

    save_table("curve-daily", curve, index=True)
    save_table("tips-daily", tips, index=True)
    save_table("funding-daily", funding, index=True)
    endpoints = raw.loc[list(DATES), ["DGS2", "DGS10", "DFII10"]].copy()
    if endpoints.isna().any().any():
        raise ValueError("A required article endpoint is missing")
    endpoints["slope_bp"] = (endpoints.DGS10 - endpoints.DGS2) * 100
    endpoints["breakeven_pct"] = endpoints.DGS10 - endpoints.DFII10
    save_table("historical-endpoints", endpoints, index=True)
    episodes = []
    for start, end in zip(DATES[:-1], DATES[1:]):
        change = endpoints.loc[end] - endpoints.loc[start]
        episodes.append({"start": start, "end": end,
                         "two_year_change_bp": change.DGS2 * 100,
                         "ten_year_change_bp": change.DGS10 * 100,
                         "slope_change_bp": change.slope_bp,
                         "real_yield_change_bp": change.DFII10 * 100,
                         "breakeven_change_bp": change.breakeven_pct * 100})
    save_table("historical-episodes", pd.DataFrame(episodes))

    dv01 = dollar_dv01(10_000_000, 1.9)
    hedge_value = dv01 / (8.1 * .0001)
    curve_scenarios = []
    for label, two_bp, ten_bp in [
        ("Bull steepening", -25, -5), ("Bear steepening", 5, 25),
        ("Parallel selloff", 25, 25), ("Flattening", 5, -5)
    ]:
        pnl = -dv01 * two_bp + dv01 * ten_bp
        curve_scenarios.append({"scenario": label, "two_year_change_bp": two_bp,
                                "ten_year_change_bp": ten_bp,
                                "long_two_year_market_value_usd": 10_000_000,
                                "short_ten_year_market_value_usd": hedge_value,
                                "dv01_per_leg_usd_per_bp": dv01, "pnl_usd": pnl})
    assert np.isclose(curve_scenarios[2]["pnl_usd"], 0)
    save_table("curve-scenarios", pd.DataFrame(curve_scenarios))

    carry_rows = []
    for label, rate, roll, shock in [
        ("Base case", 4.6, -10, 0), ("No roll benefit", 4.6, 0, 0),
        ("Funding +100 bp", 5.6, -10, 0), ("Yield shock +25 bp", 4.6, -10, 25)
    ]:
        income, roll_return, shock_return, net = carry_roll_bps(4.2, rate, .25, 4.4, roll, shock)
        carry_rows.append({"scenario": label, "yield_pct": 4.2,
                           "funding_pct": rate, "horizon_years": .25,
                           "modified_duration": 4.4, "roll_yield_bp": roll,
                           "shock_yield_bp": shock, "net_income_return_bp": income,
                           "roll_return_bp": roll_return, "shock_return_bp": shock_return,
                           "excess_return_bp": net})
    save_table("carry-scenarios", pd.DataFrame(carry_rows))
    price = 100 * (1 - .045 * 90 / 360)
    save_table("bill-conventions", pd.DataFrame([{
        "days": 90, "bank_discount_pct": 4.5, "price_per_100": price,
        "act_360_investment_pct": (100 / price - 1) * 360 / 90 * 100,
        "act_365_investment_pct": (100 / price - 1) * 365 / 90 * 100
    }]))
    spread_rows = []
    for label, shock in [("Tightens 10 bp", -10), ("Unchanged", 0),
                         ("Widens 10 bp", 10), ("Widens 20 bp", 20)]:
        spread_rows.append({"scenario": label, "spread_change_bp": shock,
                            "starting_spread_bp": 40, "extra_funding_bp": 10,
                            "horizon_years": .25, "spread_duration": 4.5,
                            "round_trip_cost_bp": 6,
                            "excess_return_bp": spread_excess_bps(40, 10, .25, 4.5, shock, 6)})
    save_table("agency-scenarios", pd.DataFrame(spread_rows))

    plt.rcParams.update({"font.family": "DejaVu Sans", "font.size": 11,
                         "text.color": INK, "axes.labelcolor": INK,
                         "xtick.color": INK, "ytick.color": INK,
                         "axes.facecolor": "#faf9f5", "figure.facecolor": "#faf9f5",
                         "svg.fonttype": "path"})
    fig, axs = plt.subplots(2, 1, figsize=(9, 6.8), sharex=True)
    fig.suptitle("Two kinds of steepening in 2024", x=.12, ha="left", fontsize=17)
    axs[0].plot(curve.index, curve.DGS2, color=BLUE, label="2-year Treasury", lw=1.8)
    axs[0].plot(curve.index, curve.DGS10, color=RUST, label="10-year Treasury", lw=1.8)
    axs[0].legend(frameon=False, loc="upper right", fontsize=10)
    style_axis(axs[0], "Par yield (%)")
    axs[1].plot(curve.index, curve.slope_bp, color=BLUE, lw=1.8)
    axs[1].axhline(0, color=INK, lw=.7, ls="--")
    style_axis(axs[1], "10-year minus 2-year (bp)")
    for ax in axs:
        ax.axvline(pd.Timestamp("2024-09-16"), color="#999d9b", ls=":", lw=1)
    finish_chart(fig, axs, "treasury-curve-2024",
                 "Source: FRED / Federal Reserve H.15, DGS2 and DGS10. Daily observations, Jan–Dec 2024.")

    fig, axs = plt.subplots(2, 1, figsize=(9, 6.8), sharex=True)
    fig.suptitle("Separating real yields and inflation compensation", x=.12, ha="left", fontsize=16)
    axs[0].plot(tips.index, tips.DGS10, color=BLUE, label="10-year nominal yield", lw=1.8)
    axs[0].plot(tips.index, tips.DFII10, color=RUST, label="10-year real yield", lw=1.8)
    axs[0].legend(frameon=False, loc="center right", fontsize=10)
    style_axis(axs[0], "Par yield (%)")
    axs[1].plot(tips.index, tips.breakeven_pct, color=BLUE, lw=1.8)
    style_axis(axs[1], "Nominal minus real (%)")
    finish_chart(fig, axs, "tips-decomposition-2024",
                 "Source: FRED / Federal Reserve H.15, DGS10 and DFII10. Par-yield difference; not a zero-coupon measure.")

    fig, ax = plt.subplots(figsize=(9, 4.8))
    fig.suptitle("Treasury repo relative to the reserve rate", x=.12, ha="left", fontsize=17)
    ax.plot(funding.index, funding.sofr_minus_iorb_bp, color=BLUE, lw=1.5)
    ax.axhline(0, color=INK, lw=.7, ls="--")
    ax.axvline(pd.Timestamp("2024-11-25"), color=RUST, lw=1, ls=":", label="SOFR methodology change: Nov 25")
    ax.legend(frameon=False, loc="upper left", fontsize=9)
    style_axis(ax, "SOFR minus IORB (bp)")
    finish_chart(fig, [ax], "funding-spread-2024",
                 "Source: FRED, SOFR (New York Fed) and IORB (Federal Reserve Board). Aligned by observation date.")

    summary = {"source_rows_with_values": {s: int(raw[s].notna().sum()) for s in SERIES},
               "curve_observations": len(curve), "tips_observations": len(tips),
               "funding_observations": len(funding), "forward_fill": False,
               "ten_year_hedge_market_value_usd": hedge_value,
               "bill_price_per_100": price, "financial_unit_checks": "passed"}
    (ROOT / "tables/validation.json").write_text(json.dumps(summary, indent=2) + "\n")
    print(json.dumps(summary, indent=2))


if __name__ == "__main__":
    main()
