#!/usr/bin/env python3
"""
ASX Momentum Reference Strategy — research engine + daily signal generator.

Strategy (reference design, long-only, daily bars, monthly rebalance):
  Universe   : S&P/ASX 200 constituents (loaded from universe.csv), liquidity screened
  Signal     : 12-1 momentum = total return from t-252 to t-21 trading days,
               divided by 252-day realised volatility (risk-adjusted momentum)
  Regime     : invest only when S&P/ASX 200 (^AXJO) close > its 200-day SMA;
               otherwise 100% cash
  Portfolio  : top N (default 20) by signal, equal weight, position cap
  Buffer     : existing holdings are kept while they remain in the top 2N (turnover control)
  Rebalance  : first trading day of each month, executed at the next session's close
               (proxy for the closing single price auction, CSPA)
  Costs      : commission + slippage per side, default 0.20% total

Modes:
  python asx_momentum.py backtest  --start 2010-01-01
  python asx_momentum.py signals   (prints today's target book as JSON for the executor)

Honesty notes baked in: the universe is the CURRENT ASX 200 list, so backtests carry
survivorship bias (they flatter results). Treat every number as an upper bound and
re-run against a point-in-time constituent history before believing any of it.
"""
import argparse, json, os, sys, math, time
import numpy as np
import pandas as pd

HERE = os.path.dirname(os.path.abspath(__file__))
DATA = os.path.join(HERE, "data")
os.makedirs(DATA, exist_ok=True)

P = dict(
    top_n=20, buffer_mult=2, lookback=252, skip=21, vol_window=252,
    regime_sma=200, min_dollar_vol=2_000_000, min_price=1.0,
    cost_per_side=0.0020, max_weight=0.075, rebalance="M",
    capital=25000.0,         # AUD; used by the per-clip cost model below
    commission_pct=0.0008, commission_min=6.0, slippage=0.0010,  # IBKR AU fixed: 0.08% min AUD 6; +10 bps slippage
    cost_model="per_clip",   # "per_clip" (Brainroom D12) or "flat" (cost_per_side)
    abs_mom=True,            # dual-momentum leg: a stock must have positive 12-1 return to be eligible
    regime_off_exposure=0.0, # v1 = BINARY (0 = all cash when index regime is off). 0.5 = pre-registered shadow variant.
    trail_stop=0.0,          # per-name trailing stop from high-water mark since entry (0 = off)
    sector_cap=0.0,          # max combined weight in Financials + Materials (0 = off); needs universe.csv Sector column
)

# ------------------------------------------------------------------ data
def load_sectors():
    try:
        u = pd.read_csv(os.path.join(HERE, "universe.csv"))
        return {str(c).strip().upper(): str(sec) for c, sec in zip(u["Code"], u.get("Sector", [""] * len(u)))}
    except Exception:
        return {}

def load_universe():
    u = pd.read_csv(os.path.join(HERE, "universe.csv"))
    return [str(c).strip().upper() for c in u["Code"] if isinstance(c, str) or not math.isnan(c)]

def fetch(tickers, start, cache="prices.parquet", refresh=False):
    import yfinance as yf
    path = os.path.join(DATA, cache)
    if os.path.exists(path) and not refresh:
        px = pd.read_parquet(path)
        if px.index.max() >= pd.Timestamp.today().normalize() - pd.Timedelta(days=4):
            return px
    syms = [t + ".AX" for t in tickers] + ["^AXJO"]
    raw = yf.download(syms, start=start, auto_adjust=True, progress=False, group_by="column", threads=True)
    close = raw["Close"].copy(); vol = raw["Volume"].copy()
    close.columns = [c.replace(".AX", "") for c in close.columns]
    vol.columns = [c.replace(".AX", "") for c in vol.columns]
    px = pd.concat({"close": close, "volume": vol}, axis=1)
    px.to_parquet(path)
    return px

# ------------------------------------------------------------------ signals
def compute_signals(close, volume, p=P):
    idx = close["^AXJO"].ffill()          # index has stray holiday NaNs; never let one NaN poison a 200-day SMA
    stocks = close.drop(columns=["^AXJO"]).ffill(limit=5)
    stock_vol = volume.drop(columns=["^AXJO"], errors="ignore")

    mom = stocks.shift(p["skip"]) / stocks.shift(p["lookback"]) - 1.0
    rets = stocks.pct_change()
    vol = rets.rolling(p["vol_window"]).std() * math.sqrt(252)
    score = mom / vol.replace(0, np.nan)

    dollar_vol = (stocks * stock_vol).rolling(20).median()
    liquid = (dollar_vol >= p["min_dollar_vol"]) & (stocks >= p["min_price"])
    score = score.where(liquid)
    if p.get("abs_mom", False):
        score = score.where(mom > 0)

    if p["regime_sma"] and p["regime_sma"] > 1:
        regime_on = idx > idx.rolling(p["regime_sma"]).mean()
    else:  # research switch: no regime filter
        regime_on = pd.Series(True, index=idx.index)
    return score, regime_on

# ------------------------------------------------------------------ backtest
def backtest(px, start, p=P, verbose=True):
    close = px["close"].loc[start:].dropna(how="all")
    volume = px["volume"].reindex(close.index)
    score, regime_on = compute_signals(px["close"], px["volume"], p)
    score = score.reindex(close.index); regime_on = regime_on.reindex(close.index).fillna(False)
    rets = close.drop(columns=["^AXJO"]).pct_change().fillna(0.0)
    idx_ret = close["^AXJO"].pct_change().fillna(0.0)

    rebal_days = close.groupby(close.index.to_period(p["rebalance"])).head(1).index
    weights = pd.Series(0.0, index=rets.columns)
    equity = [1.0]; bench = [1.0]; exposure = []; turnover = []; holdings_log = {}
    pending = None  # target weights decided at close of rebalance day, executed next close
    sectors = load_sectors(); hwm = {}; stopped = set()

    for i, d in enumerate(close.index[1:], start=1):
        cost = 0.0
        if pending is not None:
            delta = (pending - weights).abs()
            traded = float(delta.sum())
            if p.get("cost_model") == "per_clip":
                eqv = equity[-1] * p["capital"]
                comm = sum(max(p["commission_pct"] * w * eqv, p["commission_min"]) for w in delta[delta > 1e-9]) / eqv if eqv > 0 else 0.0
                cost = comm + traded * p["slippage"]
            else:
                cost = traded * p["cost_per_side"]
            turnover.append(traded / 2.0)
            weights = pending; pending = None
        day_ret = float((weights * rets.loc[d]).sum()) - cost
        equity.append(equity[-1] * (1 + day_ret)); bench.append(bench[-1] * (1 + idx_ret.loc[d]))
        if weights.sum() > 0:
            grown = weights * (1 + rets.loc[d]); weights = grown / grown.sum() * min(1.0, grown.sum() / (1 + day_ret + cost))
        exposure.append(float(weights.sum()))

        # optional per-name trailing stop: exit at next close, stay out until next rebalance
        if p.get("trail_stop", 0) and pending is None:
            drop = []
            for t in weights[weights > 0].index:
                hwm[t] = max(hwm.get(t, close.loc[d, t]), close.loc[d, t])
                if close.loc[d, t] < hwm[t] * (1 - p["trail_stop"]):
                    drop.append(t)
            for t in list(hwm):
                if weights.get(t, 0) == 0: hwm.pop(t, None)
            if drop:
                tgt = weights.copy(); tgt[drop] = 0.0; pending = tgt; stopped.update(drop)

        if d in rebal_days:
            stopped = set()
            s = score.loc[d].dropna().sort_values(ascending=False)
            scale = 1.0 if regime_on.loc[d] else float(p.get("regime_off_exposure", 0.0))
            if scale == 0.0 or len(s) < p["top_n"]:
                target = pd.Series(0.0, index=rets.columns)
            else:
                top = list(s.index[: p["top_n"]]); wide = set(s.index[: p["top_n"] * p["buffer_mult"]])
                keep = [t for t in weights[weights > 0].index if t in wide]
                book = keep + [t for t in top if t not in keep]
                book = book[: p["top_n"]]
                if p.get("sector_cap", 0) and sectors:
                    heavy = [t for t in book if sectors.get(t, "") in ("Financials", "Materials")]
                    max_heavy = int(p["sector_cap"] * p["top_n"])
                    if len(heavy) > max_heavy:
                        light = [t for t in s.index if t not in book and sectors.get(t, "") not in ("Financials", "Materials")]
                        for t in heavy[max_heavy:]:
                            book.remove(t)
                            if light: book.append(light.pop(0))
                w = min(1.0 / len(book), p["max_weight"]) * scale
                target = pd.Series(0.0, index=rets.columns); target[book] = w
                # NO DRIFT TRADES: a name already held and still in the book keeps its drifted weight
                # unless it breached the cap or the regime scale changed (avoids paying the AUD 6 minimum
                # on tiny top-ups every month). Executor applies the same rule (min trade = max(AUD 500, 1% NAV)).
                for t in book:
                    cur = float(weights.get(t, 0.0))
                    if cur > 0 and abs(cur - w) < max(0.01, 0.35 * w) and cur <= p["max_weight"] * scale + 1e-9:
                        target[t] = cur
            pending = target
            holdings_log[str(d.date())] = [t for t in target[target > 0].index]

    eq = pd.Series(equity, index=close.index); bm = pd.Series(bench, index=close.index)
    stats = summarise(eq, bm, exposure, turnover)
    if verbose:
        print(json.dumps(stats, indent=2))
    return eq, bm, stats, holdings_log

def summarise(eq, bm, exposure, turnover):
    def cagr(s): yrs = (s.index[-1] - s.index[0]).days / 365.25; return s.iloc[-1] ** (1 / yrs) - 1
    def mdd(s): return float((s / s.cummax() - 1).min())
    r = eq.pct_change().dropna(); rb = bm.pct_change().dropna()
    sharpe = float(r.mean() / r.std() * math.sqrt(252)) if r.std() > 0 else 0.0
    yearly = pd.DataFrame({"strategy": eq.resample("YE").last().pct_change(), "asx200": bm.resample("YE").last().pct_change()}).dropna()
    yearly.index = yearly.index.year
    yearly = yearly.round(4).to_dict(orient="index")
    return {
        "period": [str(eq.index[0].date()), str(eq.index[-1].date())],
        "strategy_cagr": round(float(cagr(eq)), 4), "asx200_cagr": round(float(cagr(bm)), 4),
        "strategy_vol": round(float(r.std() * math.sqrt(252)), 4), "asx200_vol": round(float(rb.std() * math.sqrt(252)), 4),
        "strategy_sharpe_rf0": round(sharpe, 3), "asx200_sharpe_rf0": round(float(rb.mean() / rb.std() * math.sqrt(252)), 3),
        "strategy_maxdd": round(mdd(eq), 4), "asx200_maxdd": round(mdd(bm), 4),
        "avg_exposure": round(float(np.mean(exposure)), 3),
        "avg_monthly_one_way_turnover": round(float(np.mean(turnover)) if turnover else 0.0, 4),
        "yearly": yearly,
        "caveat": "PROTOTYPE EVIDENCE. Current-constituent universe => survivorship bias (upper bound); benchmark is the PRICE index (no dividends); costs per Brainroom D12 (max(0.08%, AUD 6) + 10 bps slippage per clip at the configured capital).",
    }

# ------------------------------------------------------------------ live signals
def signals_today(px, p=P):
    close = px["close"]; volume = px["volume"]
    score, regime_on = compute_signals(close, volume, p)
    d = close.index[-1]
    s = score.loc[d].dropna().sort_values(ascending=False)
    on = bool(regime_on.loc[d])
    idx = close["^AXJO"].ffill()
    out = {
        "asof": str(d.date()), "regime_on": on,
        "asx200_close": round(float(idx.loc[d]), 1), "asx200_sma200": round(float(idx.rolling(p["regime_sma"]).mean().loc[d]), 1),
        "top_n": p["top_n"], "buffer_rank": p["top_n"] * p["buffer_mult"],
        "ranked": [{"code": t, "rank": i + 1, "score": round(float(v), 3), "close": round(float(close[t].loc[d]), 3)}
                    for i, (t, v) in enumerate(s.head(p["top_n"] * p["buffer_mult"]).items())],
        "target_weight_each": round(min(1.0 / p["top_n"], p["max_weight"]) * (1.0 if on else float(p.get("regime_off_exposure", 0.0))), 4),
    }
    return out

# ------------------------------------------------------------------ cli
if __name__ == "__main__":
    ap = argparse.ArgumentParser(); ap.add_argument("mode", choices=["backtest", "signals", "fetch"])
    ap.add_argument("--start", default="2010-01-01"); ap.add_argument("--refresh", action="store_true")
    ap.add_argument("--top", type=int, default=P["top_n"]); ap.add_argument("--capital", type=float, default=P["capital"])
    ap.add_argument("--graded", action="store_true", help="shadow variant: 50%% exposure when regime off")
    a = ap.parse_args()
    P["top_n"] = a.top; P["capital"] = a.capital
    if a.graded: P["regime_off_exposure"] = 0.5
    tick = load_universe()
    px = fetch(tick, "2008-01-01", refresh=a.refresh or a.mode == "fetch")
    if a.mode == "backtest":
        eq, bm, stats, hl = backtest(px, a.start)
        pd.DataFrame({"strategy": eq, "asx200": bm}).to_csv(os.path.join(DATA, "equity_curve.csv"))
        json.dump(stats, open(os.path.join(DATA, "backtest_stats.json"), "w"), indent=2)
        json.dump(hl, open(os.path.join(DATA, "holdings_log.json"), "w"))
    elif a.mode == "signals":
        print(json.dumps(signals_today(px), indent=2))