#!/usr/bin/env python3
"""Walk-forward moving-average crossover backtest on HF Market Data bars.

  ma_crossover.py --asset etf --ticker SPY --start 2010-01-01 [--grid 10,20,50 --grid-slow 100,150,200]
                  [--train-years 3 --test-years 1] [--cost-bps 5] [--allow-short] [--fixed 50,200] [--plot out.png]

The reported metrics are OUT-OF-SAMPLE: for each test block the (fast, slow) pair is chosen on the
preceding training window only. Positions apply to the next bar (no look-ahead).
"""
from __future__ import annotations

import argparse
import itertools
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 signal_ma_cross(close: pd.Series, fast: int, slow: int, allow_short: bool) -> pd.Series:
    """Position in {0,1} (or {-1,0,1}) from past data only. Replace to test another idea."""
    f, s = close.rolling(fast).mean(), close.rolling(slow).mean()
    pos = (f > s).astype(float)
    if allow_short:
        pos = pos - (f < s).astype(float)
    return pos


def run(close: pd.Series, pos: pd.Series, cost_bps: float) -> pd.DataFrame:
    ret = close.pct_change().fillna(0.0)
    p = pos.shift(1).fillna(0.0)  # decide at t, hold from t+1
    trades = p.diff().abs().fillna(0.0)
    strat = p * ret - trades * cost_bps / 1e4
    return pd.DataFrame({"ret": ret, "pos": p, "strat": strat, "trades": trades})


def periods_per_year(dt: pd.Series) -> int:
    step = dt.diff().dropna().median()
    if step >= pd.Timedelta(days=1):
        return 252 if dt.dt.dayofweek.max() <= 4 else 365
    per_day = int(pd.Timedelta(hours=6.5) / step) if step < pd.Timedelta(hours=1) else int(pd.Timedelta(hours=24) / step)
    return max(per_day, 1) * 252


def metrics(r: pd.Series, ppy: int) -> dict:
    if len(r) < 2:
        return {"cagr": np.nan, "vol": np.nan, "sharpe": np.nan, "max_dd": np.nan}
    eq = (1 + r).cumprod()
    years = len(r) / ppy
    return {
        "cagr": eq.iloc[-1] ** (1 / years) - 1 if years > 0 else np.nan,
        "vol": r.std() * np.sqrt(ppy),
        "sharpe": r.mean() / r.std() * np.sqrt(ppy) if r.std() > 0 else np.nan,
        "max_dd": (eq / eq.cummax() - 1).min(),
    }


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--asset", required=True, choices=["stock", "etf", "crypto", "index", "fx", "futures"])
    ap.add_argument("--ticker", required=True)
    ap.add_argument("--timeframe", default="1day")
    ap.add_argument("--start", default="2010-01-01")
    ap.add_argument("--end")
    ap.add_argument("--adjustment", help="futures: prefer contin_adj_ratio")
    ap.add_argument("--grid", default="10,20,50", help="fast MA candidates")
    ap.add_argument("--grid-slow", default="100,150,200", help="slow MA candidates")
    ap.add_argument("--fixed", help="fast,slow — skip optimisation, evaluate one pair on the whole history")
    ap.add_argument("--train-years", type=float, default=3.0)
    ap.add_argument("--test-years", type=float, default=1.0)
    ap.add_argument("--cost-bps", type=float, default=5.0, help="cost per side per unit of position change")
    ap.add_argument("--allow-short", action="store_true")
    ap.add_argument("--show-insample", action="store_true", help="also print the best pair on the full history (for contrast)")
    ap.add_argument("--plot")
    ap.add_argument("--out-csv", help="write the out-of-sample equity curve")
    a = ap.parse_args()

    adj = a.adjustment or ("contin_adj_ratio" if a.asset == "futures" else None)
    df = hfmd.bars(a.asset, a.ticker.upper(), a.timeframe, a.start, a.end, adj)
    if len(df) < 300:
        print(f"only {len(df)} bars — too short for a walk-forward", file=sys.stderr)
        return 1
    df = df.set_index("datetime")
    close = df["close"].astype(float)
    ppy = periods_per_year(pd.Series(df.index))
    print(f"{a.ticker.upper()} {a.timeframe} {adj or ''}: {len(close):,} bars {close.index[0].date()} → {close.index[-1].date()} · periods/year={ppy} · cost={a.cost_bps} bps/side")

    fasts = [int(x) for x in a.grid.split(",")]
    slows = [int(x) for x in a.grid_slow.split(",")]
    pairs = [(f, s) for f, s in itertools.product(fasts, slows) if f < s]

    if a.fixed:
        f, s = (int(x) for x in a.fixed.split(","))
        res = run(close, signal_ma_cross(close, f, s, a.allow_short), a.cost_bps).iloc[s:]
        blocks = pd.DataFrame([{"test_start": res.index[0].date(), "test_end": res.index[-1].date(), "fast": f, "slow": s, **{f"oos_{k}": v for k, v in metrics(res["strat"], ppy).items()}}])
    else:
        train_n, test_n = int(a.train_years * ppy), int(a.test_years * ppy)
        if train_n < max(slows) + 20:
            print(f"train window ({train_n} bars) too short for slow MA {max(slows)}", file=sys.stderr)
            return 1
        pieces, rows = [], []
        start = train_n
        while start + 20 < len(close):
            tr = close.iloc[start - train_n:start]
            best, best_sh = None, -np.inf
            for f, s in pairs:
                sh = metrics(run(tr, signal_ma_cross(tr, f, s, a.allow_short), a.cost_bps)["strat"].iloc[s:], ppy)["sharpe"]
                if np.isfinite(sh) and sh > best_sh:
                    best, best_sh = (f, s), sh
            f, s = best or pairs[0]
            # compute the signal on history + test block so the MAs are warm at the block start
            seg = close.iloc[max(0, start - s - 5):start + test_n]
            res = run(seg, signal_ma_cross(seg, f, s, a.allow_short), a.cost_bps).loc[close.index[start]:]
            pieces.append(res)
            m = metrics(res["strat"], ppy)
            rows.append({"test_start": res.index[0].date(), "test_end": res.index[-1].date(), "fast": f, "slow": s, "train_sharpe": round(best_sh, 2), **{f"oos_{k}": v for k, v in m.items()}})
            start += test_n
        if not pieces:
            print("not enough data for one test block", file=sys.stderr)
            return 1
        res = pd.concat(pieces)
        blocks = pd.DataFrame(rows)

    strat, bench = metrics(res["strat"], ppy), metrics(res["ret"], ppy)
    n_trades = int((res["trades"] > 0).sum())
    print(f"\nOUT-OF-SAMPLE {res.index[0].date()} → {res.index[-1].date()} ({len(res):,} bars, {len(blocks)} block(s))")
    print(f"{'':14}{'strategy':>12}{'buy&hold':>12}")
    for k in ("cagr", "vol", "sharpe", "max_dd"):
        print(f"{k:<14}{strat[k]:>12.3f}{bench[k]:>12.3f}")
    print(f"{'exposure':<14}{res['pos'].abs().mean():>12.2f}")
    print(f"{'trades':<14}{n_trades:>12d}   turnover/yr={res['trades'].sum() / (len(res) / ppy):.1f}")
    print("\nper block (parameters chosen on the preceding training window):")
    with pd.option_context("display.width", 160, "display.float_format", "{:.3f}".format):
        print(blocks.to_string(index=False))

    if a.show_insample and not a.fixed:
        best = max(pairs, key=lambda p: metrics(run(close, signal_ma_cross(close, *p, a.allow_short), a.cost_bps)["strat"].iloc[p[1]:], ppy)["sharpe"] or -np.inf)
        m = metrics(run(close, signal_ma_cross(close, *best, a.allow_short), a.cost_bps)["strat"].iloc[best[1]:], ppy)
        print(f"\nIN-SAMPLE best pair on full history (optimistic, for contrast): {best} sharpe={m['sharpe']:.2f} cagr={m['cagr']:.3f} max_dd={m['max_dd']:.3f}")

    if a.out_csv:
        out = res.copy()
        out["equity"] = (1 + out["strat"]).cumprod()
        out["benchmark"] = (1 + out["ret"]).cumprod()
        out.to_csv(a.out_csv)
        print(f"wrote {a.out_csv}")

    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
        eq, bm = (1 + res["strat"]).cumprod(), (1 + res["ret"]).cumprod()
        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(11, 6.5), sharex=True, gridspec_kw={"height_ratios": [3, 1]})
        ax1.plot(eq.index, eq, lw=1.3, label=f"MA cross walk-forward (Sharpe {strat['sharpe']:.2f})")
        ax1.plot(bm.index, bm, lw=1.0, alpha=0.7, label=f"buy & hold (Sharpe {bench['sharpe']:.2f})")
        if not a.fixed:
            for _, b in blocks.iterrows():
                ax1.axvline(pd.Timestamp(b["test_start"]), color="grey", alpha=0.25, lw=0.8)
                ax1.text(pd.Timestamp(b["test_start"]), ax1.get_ylim()[0], f"{b['fast']}/{b['slow']}", fontsize=7, rotation=90, va="bottom", alpha=0.7)
        ax1.set_yscale("log")
        ax1.grid(alpha=0.25)
        ax1.legend(loc="upper left")
        ax1.set_title(f"{a.ticker.upper()} — out-of-sample equity (log), cost {a.cost_bps} bps/side")
        ax2.fill_between(eq.index, (eq / eq.cummax() - 1) * 100, 0, alpha=0.4)
        ax2.set_ylabel("drawdown %")
        ax2.grid(alpha=0.25)
        fig.tight_layout()
        fig.savefig(a.plot, dpi=130)
        print(f"wrote {a.plot}")
    return 0


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