#!/usr/bin/env python3
"""Stitch a continuous futures series locally from individual contracts (transparent methodology,
and a fallback when the v2 /continuous endpoint is unavailable).

  stitch_local.py --root CL --start 2023-01-01 [--roll volume|open_interest|calendar] [--roll-days 5]
                  [--adjust none|back_adjusted|ratio] [--out cl_local.csv]

Costs one request per contract in the range (≈ 12/yr for CL, 4/yr for ES) — set HFMD_API_KEY.
"""
from __future__ import annotations

import argparse
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


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--root", required=True)
    ap.add_argument("--start", default="2023-01-01")
    ap.add_argument("--end")
    ap.add_argument("--roll", default="volume", choices=["volume", "open_interest", "calendar"])
    ap.add_argument("--roll-days", type=int, default=5, help="calendar rule: business days before expiration")
    ap.add_argument("--adjust", default="ratio", choices=["none", "back_adjusted", "ratio"])
    ap.add_argument("--out")
    a = ap.parse_args()
    root = a.root.upper()

    cons = hfmd.contracts(root)
    if cons.empty:
        print("no contracts returned", file=sys.stderr)
        return 1
    cons["expiration_date"] = pd.to_datetime(cons["expiration_date"])
    cons["last_data_date"] = pd.to_datetime(cons.get("last_data_date"))
    lo = pd.Timestamp(a.start) - pd.Timedelta(days=120)
    hi = pd.Timestamp(a.end) if a.end else pd.Timestamp.today()
    sel = cons[(cons["expiration_date"] >= lo) & (cons["first_data_date"].pipe(pd.to_datetime) <= hi)].sort_values("expiration_date") if "first_data_date" in cons else cons[cons["expiration_date"] >= lo].sort_values("expiration_date")
    print(f"{root}: {len(sel)} contracts in scope ({sel['symbol'].iloc[0]} … {sel['symbol'].iloc[-1]})")

    panels: dict[str, pd.DataFrame] = {}
    for sym in sel["symbol"]:
        df = hfmd.contract_bars(sym, "1day", a.start and str((pd.Timestamp(a.start) - pd.Timedelta(days=120)).date()), a.end)
        if df.empty:
            hfmd.log(f"  {sym}: no bars (gap reported, not filled)")
            continue
        panels[sym] = df.set_index("datetime")
    if not panels:
        return 1
    order = [s for s in sel["symbol"] if s in panels]
    expiry = dict(zip(sel["symbol"], sel["expiration_date"]))

    # --- choose the active contract on each date ---------------------------------------------------
    dates = sorted(set().union(*[p.index for p in panels.values()]))
    dates = [d for d in dates if d >= pd.Timestamp(a.start)]
    active_idx, active, rolls = 0, order[0], []
    chosen: list[tuple[pd.Timestamp, str]] = []
    for d in dates:
        # skip contracts already expired
        while active_idx + 1 < len(order) and expiry[active] < d:
            active_idx += 1
            active = order[active_idx]
            rolls.append((d, active, "expired"))
        if active_idx + 1 < len(order):
            nxt = order[active_idx + 1]
            switch = False
            if a.roll == "calendar":
                switch = d >= expiry[active] - pd.tseries.offsets.BDay(a.roll_days)
            elif d in panels[active].index and d in panels[nxt].index:
                col = "volume" if a.roll == "volume" else "open_interest"
                if col in panels[nxt].columns and col in panels[active].columns:
                    switch = float(panels[nxt].loc[d, col] or 0) > float(panels[active].loc[d, col] or 0)
            if switch:
                active_idx += 1
                active = nxt
                rolls.append((d, active, a.roll))
        chosen.append((d, active))

    # --- assemble + adjust --------------------------------------------------------------------------
    rows = []
    for d, sym in chosen:
        if d in panels[sym].index:
            r = panels[sym].loc[d]
            rows.append({"datetime": d, "contract": sym, "open": r["open"], "high": r["high"], "low": r["low"], "close": r["close"], "volume": r.get("volume"), "open_interest": r.get("open_interest")})
        else:
            rows.append({"datetime": d, "contract": sym, "open": np.nan, "high": np.nan, "low": np.nan, "close": np.nan, "volume": np.nan, "open_interest": np.nan})
    out = pd.DataFrame(rows).set_index("datetime")

    if a.adjust != "none":
        adj = out.copy()
        # walk backwards: at each roll, compute the gap between the new and old contract on the roll date
        roll_dates = [d for d, _, _ in rolls if d in out.index]
        for d in reversed(roll_dates):
            i = out.index.get_loc(d)
            new_sym, old_sym = out["contract"].iloc[i], out["contract"].iloc[i - 1] if i > 0 else None
            if not old_sym or d not in panels[old_sym].index or d not in panels[new_sym].index:
                continue
            new_c, old_c = float(panels[new_sym].loc[d, "close"]), float(panels[old_sym].loc[d, "close"])
            before = adj.index < d
            for col in ("open", "high", "low", "close"):
                if a.adjust == "back_adjusted":
                    adj.loc[before, col] = adj.loc[before, col] + (new_c - old_c)
                else:
                    adj.loc[before, col] = adj.loc[before, col] * (new_c / old_c if old_c else 1.0)
        out = adj

    print(f"\nrolls ({a.roll}{' ' + str(a.roll_days) + 'bd' if a.roll == 'calendar' else ''}, adjust={a.adjust}): {len(rolls)}")
    for d, sym, why in rolls:
        print(f"  {d.date()}  → {sym}  ({why})")
    print(f"\n{len(out):,} rows {out.index[0].date()} → {out.index[-1].date()} · missing closes: {int(out['close'].isna().sum())}")
    print(f"last close (adjusted={a.adjust}): {out['close'].iloc[-1]:.4f} · contract {out['contract'].iloc[-1]}")
    if a.out:
        out.to_csv(a.out)
        print(f"wrote {a.out}")
    return 0


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