#!/usr/bin/env python3
"""Futures term structure: fetch, quantify (contango/backwardation, spreads, roll yield), compare, plot.

  term_structure.py --root CL [--as-of YYYY-MM-DD] [--compare YYYY-MM-DD] [--depth 12] [--plot out.png] [--json out.json]
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import numpy as np
import pandas as pd

sys.path.insert(0, str(Path(__file__).parent))
import hfmd  # noqa: E402

PRICE_COLS = ("settle", "close", "last", "price")


def price_col(df: pd.DataFrame) -> str | None:
    return next((c for c in PRICE_COLS if c in df.columns), None)


def curve(root: str, as_of: str | None, depth: int | None) -> tuple[pd.DataFrame, dict]:
    """v2 term-structure endpoint, with a local rebuild fallback from contracts + bars."""
    try:
        df, meta = hfmd.term_structure(root, as_of)
        if depth:
            df = df.head(depth)
        return df, meta
    except hfmd.HfmdError as e:
        if e.status != 404:
            raise
        hfmd.log("term-structure endpoint unavailable (404) — rebuilding from contracts + bars")
    try:
        cons = hfmd.contracts(root)
    except hfmd.HfmdError as e:
        if e.status == 404:
            raise SystemExit(f"{e}\n→ the v2 futures module (term-structure, contracts) is not deployed on this server yet; nothing to rebuild from.")
        raise
    if cons.empty:
        raise SystemExit("no contracts available")
    ref = pd.Timestamp(as_of) if as_of else pd.Timestamp.today().normalize()
    cons["expiration_date"] = pd.to_datetime(cons["expiration_date"])
    live = cons[cons["expiration_date"] >= ref].sort_values("expiration_date").head(depth or 12)
    rows, used = [], None
    for _, c in live.iterrows():
        b = hfmd.contract_bars(c["symbol"], "1day", str((ref - pd.Timedelta(days=7)).date()), str(ref.date()))
        if b.empty:
            rows.append({"symbol": c["symbol"], "expiration_date": c["expiration_date"], "close": np.nan, "volume": np.nan, "open_interest": np.nan})
            continue
        last = b.iloc[-1]
        used = max(used or last["datetime"], last["datetime"])
        rows.append({"symbol": c["symbol"], "expiration_date": c["expiration_date"], "close": last["close"], "volume": last.get("volume"), "open_interest": last.get("open_interest")})
    return pd.DataFrame(rows), {"as_of": str(used.date()) if used is not None else as_of, "source": "rebuilt-locally"}


def analyse(df: pd.DataFrame, as_of: str | None) -> tuple[pd.DataFrame, dict]:
    pc = price_col(df)
    if pc is None:
        raise SystemExit(f"no price column in {list(df.columns)}")
    d = df.copy()
    d["expiration_date"] = pd.to_datetime(d.get("expiration_date", d.get("expiry")))
    d = d.sort_values("expiration_date").reset_index(drop=True)
    ref = pd.Timestamp(as_of) if as_of else pd.Timestamp.today().normalize()
    d["dte"] = (d["expiration_date"] - ref).dt.days
    front = float(d[pc].dropna().iloc[0]) if d[pc].notna().any() else np.nan
    d["spread_vs_front"] = d[pc] - front
    d["spread_pct"] = d["spread_vs_front"] / front * 100
    prev_p, prev_dte = d[pc].shift(1), d["dte"].shift(1)
    gap_days = (d["dte"] - prev_dte).replace(0, np.nan)
    d["leg_roll_yield_ann_pct"] = -(d[pc] - prev_p) / prev_p * 365 / gap_days * 100
    legs = d["leg_roll_yield_ann_pct"].dropna()
    up = int((d[pc].diff().dropna() > 0).sum())
    down = int((d[pc].diff().dropna() < 0).sum())
    if up and not down:
        shape = "contango (monotonic upward)"
    elif down and not up:
        shape = "backwardation (monotonic downward)"
    elif up + down == 0:
        shape = "flat"
    else:
        shape = f"mixed / humped ({up} rising legs, {down} falling legs)"
    n6 = min(6, len(d) - 1)
    summary = {
        "shape": shape,
        "front": d["symbol"].iloc[0] if "symbol" in d else None,
        "front_price": front,
        "front_to_%dth_pct" % (n6 + 1): float(d["spread_pct"].iloc[n6]) if n6 > 0 else np.nan,
        "avg_leg_roll_yield_ann_pct": float(legs.mean()) if len(legs) else np.nan,
        "front_leg_roll_yield_ann_pct": float(legs.iloc[0]) if len(legs) else np.nan,
        "legs_without_price": int(d[pc].isna().sum()),
    }
    return d, summary


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--root", required=True)
    ap.add_argument("--as-of")
    ap.add_argument("--compare", help="second date to compare the curve with")
    ap.add_argument("--depth", type=int, default=12)
    ap.add_argument("--plot")
    ap.add_argument("--json")
    a = ap.parse_args()
    root = a.root.upper()

    df, meta = curve(root, a.as_of, a.depth)
    if df.empty:
        print("empty curve", file=sys.stderr)
        return 1
    used = meta.get("as_of") or a.as_of
    d, s = analyse(df, used)
    pc = price_col(d)
    cols = [c for c in ("symbol", "expiration_date", "dte", pc, "spread_vs_front", "spread_pct", "leg_roll_yield_ann_pct", "volume", "open_interest") if c in d.columns]
    print(f"{root} term structure as of {used}{' (' + meta['source'] + ')' if 'source' in meta else ''}")
    with pd.option_context("display.width", 200, "display.float_format", "{:.3f}".format):
        print(d[cols].to_string(index=False))
    print("\nverdict:", s["shape"])
    for k, v in s.items():
        if k != "shape":
            print(f"  {k}: {v:.3f}" if isinstance(v, float) and np.isfinite(v) else f"  {k}: {v}")

    d2 = s2 = None
    if a.compare:
        df2, meta2 = curve(root, a.compare, a.depth)
        used2 = meta2.get("as_of") or a.compare
        d2, s2 = analyse(df2, used2)
        print(f"\ncompared with {used2}: {s2['shape']} · front {s2['front_price']:.3f} → {s['front_price']:.3f} ({(s['front_price'] / s2['front_price'] - 1) * 100:+.2f}%)")
        k6 = [k for k in s if k.startswith("front_to_")][0]
        print(f"  slope front→{k6.split('_')[2]}: {s2.get(k6, np.nan):+.2f}% → {s.get(k6, np.nan):+.2f}% (twist {s.get(k6, np.nan) - s2.get(k6, np.nan):+.2f} pp)")
        m = d.merge(d2, on="symbol", suffixes=("", "_prev")) if "symbol" in d and "symbol" in d2 else pd.DataFrame()
        if not m.empty:
            m["chg_pct"] = (m[pc] / m[f"{pc}_prev"] - 1) * 100
            print("  per-contract change (%):", ", ".join(f"{r.symbol} {r.chg_pct:+.1f}" for r in m.itertuples()))

    if a.json:
        Path(a.json).write_text(json.dumps({"root": root, "as_of": used, "summary": s, "curve": json.loads(d[cols].to_json(orient="records", date_format="iso")),
                                            **({"compare_as_of": a.compare, "compare_summary": s2} if s2 else {})}, indent=1, default=str))
        print(f"wrote {a.json}")

    if a.plot:
        try:
            import matplotlib
            matplotlib.use("Agg")
            import matplotlib.pyplot as plt
        except ImportError:
            print("matplotlib not installed", file=sys.stderr)
            return 0
        fig, ax = plt.subplots(figsize=(10, 5))
        ax.plot(d["expiration_date"], d[pc], marker="o", lw=1.4, label=f"{used} — {s['shape'].split(' (')[0]}")
        for _, r in d.iterrows():
            if pd.notna(r[pc]) and "symbol" in r:
                ax.annotate(r["symbol"], (r["expiration_date"], r[pc]), textcoords="offset points", xytext=(0, 7), fontsize=7, ha="center")
        if d2 is not None:
            ax.plot(d2["expiration_date"], d2[pc], marker="s", lw=1.0, alpha=0.7, label=f"{a.compare} — {s2['shape'].split(' (')[0]}")
        ax.set_title(f"{root} futures term structure")
        ax.set_xlabel("contract expiry")
        ax.set_ylabel("price")
        ax.grid(alpha=0.25)
        ax.legend()
        fig.autofmt_xdate()
        fig.tight_layout()
        fig.savefig(a.plot, dpi=130)
        print(f"wrote {a.plot}")
    return 0


if __name__ == "__main__":
    sys.exit(main())