"""
Risk parity on US ETFs: inverse volatility, equal risk contribution and a trend filter.

FinLab research, 2026-11.

Designs (each with a 3 x 3 grid of two parameters that includes the original setting)
  ivol_spy, ivol_qqq            weight proportional to 1 / volatility of daily returns
  erc_spy, erc_qqq              equal risk contribution (Maillard, Roncalli & Teiletche 2010),
                                long only, from the covariance of daily returns
                                grid for all four: window 21/63/126 days x rebalance every
                                1/2/3 months
  erc_trend_spy, erc_trend_qqq  the erc weights; at each month end an asset whose close is not
                                above its own M-month average gives its slice to T-bills
                                (BIL), after Faber (2007)
                                grid: M 8/10/12 months x window 21/63/126 days, monthly

Asset sets: "spy" is SPY, TLT, GLD, DBC and "qqq" is QQQ, TLT, GLD. No leverage: weights
add up to 100%. Signals are read on the last trading day of the month and trade at the next
close (finlab `sim(trade_at_price="close")`).

A design "beats QQQ" only with a higher daily Sharpe ratio than buying and holding QQQ,
dividends included, over the same dates, both with no risk-free rate deducted and on returns
above T-bills (BIL). Bonds and cash flatter the first measure. A design passes only if
  1. at least 7 of its 9 grid settings beat QQQ,
  2. its original setting beats QQQ in both halves of the backtest,
  3. its original setting beats QQQ at finlab's default cost (zero for US funds) and after
     paying 0.1% per buy and per sell,
  4. for the two trend designs, its original setting also beats the same risk parity without
     the filter and the filter on equal weights without risk parity.

Controls, not judged: RPAR and UPAR since listing; 60/40, All Weather and the Permanent
Portfolio as defined in FinLab's lazy portfolio backtest; equal weights on the same assets;
each original design's average weights held fixed.

Run:  python research.py      writes metrics.json and data.csv next to this file
Needs finlab 2.2.21 or later (older versions trade US funds on market holidays).
"""
import json
from functools import cache
from pathlib import Path

import numpy as np
import pandas as pd
from finlab import data
from finlab.backtest import sim

SIGNAL_START = "2016-12-30"   # first month end with 12 month-end closes behind it
FIRST_FILL = "2017-01-03"     # every curve (and QQQ) starts at this close
END = "2026-10-30"
COST_PER_SIDE = 0.001
OUT = Path(__file__).resolve().parent

ASSETS = {"spy": ["SPY", "TLT", "GLD", "DBC"], "qqq": ["QQQ", "TLT", "GLD"]}
SAFE = "BIL"
LISTED_FUNDS = ["RPAR", "UPAR"]
FIXED = {   # same weights as the lazy portfolio backtest, rebalanced quarterly
    "sixty_forty": {"SPY": 0.60, "AGG": 0.40},
    "all_weather": {"VTI": 0.30, "TLT": 0.40, "IEF": 0.15, "GLD": 0.075, "DBC": 0.075},
    "permanent": {"VTI": 0.25, "TLT": 0.25, "GLD": 0.25, "BIL": 0.25},
}

WINDOW_GRID = (("window", "every"), ((21, 63, 126), (1, 2, 3)), (63, 1))
TREND_GRID = (("months", "window"), ((8, 10, 12), (21, 63, 126)), (10, 63))
DESIGNS = {
    "ivol_spy": WINDOW_GRID, "ivol_qqq": WINDOW_GRID,
    "erc_spy": WINDOW_GRID, "erc_qqq": WINDOW_GRID,
    "erc_trend_spy": TREND_GRID, "erc_trend_qqq": TREND_GRID,
}

data.truncate_end = END


@cache
def prices():
    data.set_market("us_fund")
    # us_fund_price keeps market holidays as blank rows
    return data.get("us_fund_price:adj_close").dropna(subset=["SPY"])


@cache
def returns():
    return prices().pct_change(fill_method=None)


@cache
def month_ends():
    days = prices().index
    return pd.DatetimeIndex(pd.Series(days, index=days).groupby(days.to_period("M")).max().values)


def signal_dates(every=1):
    return month_ends()[month_ends() >= SIGNAL_START][::every]


# ---------------------------------------------------------------------------
# Weights (signal date x asset)
# ---------------------------------------------------------------------------
def equal_risk_contribution(cov):
    """Long-only weights with equal risk contributions, by cyclical coordinate descent
    (Griveau-Billion, Richard & Roncalli 2013)."""
    n = len(cov)
    w = np.full(n, 1 / n)
    for _ in range(200):
        for i in range(n):
            others = cov[i] @ w - cov[i, i] * w[i]
            w[i] = (-others + (others ** 2 + 4 * cov[i, i] / n) ** 0.5) / (2 * cov[i, i])
    return w / w.sum()


def inverse_volatility(assets, window):
    inverse = 1 / returns()[assets].rolling(window).std().reindex(signal_dates())
    return inverse.div(inverse.sum(axis=1), axis=0)


@cache
def risk_parity(asset_set, window):
    assets = ASSETS[asset_set]
    rows = {d: equal_risk_contribution(returns()[assets].loc[:d].tail(window).cov().to_numpy())
            for d in signal_dates()}
    return pd.DataFrame(rows, index=assets).T


def equal_weights(asset_set):
    assets = ASSETS[asset_set]
    return pd.DataFrame(1 / len(assets), index=signal_dates(), columns=assets)


def trend_filter(w, months):
    """Each slice whose asset is not above its M-month average goes to T-bills."""
    monthly = prices()[w.columns].reindex(month_ends())
    above = (monthly > monthly.rolling(months).mean()).reindex(w.index)
    filtered = w.where(above, 0.0)
    filtered[SAFE] = 1 - filtered.sum(axis=1)
    return filtered


def weights(name, a, b):
    """Target weights and the signal dates that trade."""
    method, asset_set = name.rsplit("_", 1)
    if method == "ivol":
        return inverse_volatility(ASSETS[asset_set], a), signal_dates(b)
    if method == "erc":
        return risk_parity(asset_set, a), signal_dates(b)
    return trend_filter(risk_parity(asset_set, b), a), signal_dates()


def run(w, dates, fee_ratio=None):
    position = w.reindex(prices().loc[SIGNAL_START:].index, method="ffill").fillna(0)
    data.set_market("us_fund")
    report = sim(position, resample=dates, trade_at_price="close", fee_ratio=fee_ratio, upload=False)
    return report.creturn.loc[FIRST_FILL:END]


@cache
def backtest(name, a, b, fee_ratio=None):
    return run(*weights(name, a, b), fee_ratio)


# ---------------------------------------------------------------------------
# Statistics
# ---------------------------------------------------------------------------
def daily_sharpe(r):
    return r.mean() / r.std() * 252 ** 0.5


def sharpe(curve):
    return daily_sharpe(curve.pct_change().dropna())


def excess_sharpe(curve):
    """Daily Sharpe of the returns above T-bills (BIL)."""
    r = curve.pct_change().dropna()
    return daily_sharpe(r - returns()[SAFE].reindex(r.index))


def qqq(index):
    return prices()["QQQ"].reindex(index)


def stats(curve):
    curve = curve.dropna()
    bench = qqq(curve.index)
    years = (curve.index[-1] - curve.index[0]).days / 365.25
    r = curve.pct_change().dropna()
    curve_sharpe = daily_sharpe(r)
    bench_sharpe = sharpe(bench)
    curve_excess = excess_sharpe(curve)
    bench_excess = excess_sharpe(bench)
    return {
        "start": str(curve.index[0].date()), "end": str(curve.index[-1].date()),
        "sharpe": round(curve_sharpe, 6), "qqq_sharpe": round(bench_sharpe, 6),
        "excess_sharpe": round(curve_excess, 6), "qqq_excess_sharpe": round(bench_excess, 6),
        "cagr": round((curve.iloc[-1] / curve.iloc[0]) ** (1 / years) - 1, 6),
        "qqq_cagr": round((bench.iloc[-1] / bench.iloc[0]) ** (1 / years) - 1, 6),
        "volatility": round(r.std() * 252 ** 0.5, 6),
        "mdd": round((curve / curve.cummax() - 1).min(), 6),
        "beats_qqq": bool(curve_sharpe > bench_sharpe and curve_excess > bench_excess),
        "beats_qqq_raw": bool(curve_sharpe > bench_sharpe),
    }


def halves(curve):
    middle = curve.index[len(curve) // 2]
    return [stats(curve.loc[:middle]), stats(curve.loc[middle:])]


def yearly(curve):
    ends = curve.resample("YE").last()
    r = ends / ends.shift(1, fill_value=curve.iloc[0]) - 1
    return {str(d.year): round(v, 6) for d, v in r.items()}


def turnover_per_year(w, dates):
    """One-way turnover of the target weights, ignoring drift between rebalances."""
    held = w.reindex(dates).fillna(0)
    return round(float(held.diff().abs().sum(axis=1).sum() / 2 / (len(held) / 12)), 4)


def risk_shares(w, window):
    """Average share of portfolio variance from each asset under the trailing covariance."""
    rows = []
    for d, row in w.iterrows():
        cov = returns()[w.columns].loc[:d].tail(window).cov().to_numpy()
        contribution = row.to_numpy() * (cov @ row.to_numpy())
        rows.append(contribution / contribution.sum())
    return {k: round(float(v), 4) for k, v in zip(w.columns, np.mean(rows, axis=0))}


def evaluate(name):
    (pa, pb), (va, vb), (oa, ob) = DESIGNS[name]
    grid = []
    for a in va:
        for b in vb:
            curve = backtest(name, a, b)
            row = {pa: a, pb: b, **stats(curve), "halves": halves(curve)}
            grid.append(row)
            print(f"{name:<14} {pa}={a} {pb}={b}  Sharpe {row['sharpe']:.3f} vs QQQ {row['qqq_sharpe']:.3f}"
                  f"  above BIL {row['excess_sharpe']:.3f} vs {row['qqq_excess_sharpe']:.3f}", flush=True)
    original = next(g for g in grid if (g[pa], g[pb]) == (oa, ob))
    cost = stats(backtest(name, oa, ob, COST_PER_SIDE))
    w, dates = weights(name, oa, ob)
    wins = sum(g["beats_qqq"] for g in grid)
    return {
        "params": [pa, pb], "values": [list(va), list(vb)], "original": [oa, ob],
        "grid": grid, "grid_wins": wins, "grid_wins_raw": sum(g["beats_qqq_raw"] for g in grid),
        "original_stats": {k: original[k] for k in original if k not in (pa, pb)},
        "cost_0.1pct": cost,
        "average_weights": {k: round(float(v), 4) for k, v in w.reindex(dates).iloc[:-1].mean().items()},
        "turnover_per_year": turnover_per_year(w, dates),
        "checks": {"grid": wins >= 7, "halves": all(h["beats_qqq"] for h in original["halves"]),
                   "cost": original["beats_qqq"] and cost["beats_qqq"]},
    }


def fixed_at_average(name):
    w, dates = weights(name, *DESIGNS[name][2])
    average = w.reindex(dates).iloc[:-1].mean()   # the last month end trades after the data ends
    return run(pd.DataFrame(average.to_dict(), index=w.index), signal_dates()), average


def levered_to_qqq_volatility(curve):
    """Scale a curve to QQQ's volatility, borrowing at the T-bill return with no other cost."""
    r = curve.pct_change().dropna()
    leverage = qqq(curve.index).pct_change().std() / r.std()
    bills = returns()[SAFE].reindex(r.index)
    levered = pd.concat([curve.iloc[:1] / curve.iloc[0], 1 + leverage * r - (leverage - 1) * bills]).cumprod()
    return {"leverage": round(leverage, 4), **stats(levered)}


def since_listing(curves):
    """The listed risk parity funds and our curves over the same dates."""
    out = {}
    for fund in LISTED_FUNDS:
        listed = prices()[fund].dropna()
        out[fund] = {name: stats(c.reindex(listed.index).dropna()) for name, c in {fund: listed, **curves}.items()}
    return out


def fixed_weights(mix):
    return pd.DataFrame(mix, index=signal_dates(), columns=list(mix))


def quarter_ends():
    dates = signal_dates()
    return dates[dates.month.isin([3, 6, 9, 12])]


def main():
    out = {"window": {"signal_start": SIGNAL_START, "end": END}, "cost_per_side": COST_PER_SIDE,
           "finlab_version": __import__("finlab").__version__, "pandas_version": pd.__version__}
    out["designs"] = {name: evaluate(name) for name in DESIGNS}

    curves = {name: backtest(name, *DESIGNS[name][2]) for name in DESIGNS}
    controls = {}
    for asset_set in ASSETS:
        controls[f"equal_{asset_set}"] = run(equal_weights(asset_set), signal_dates())
        controls[f"trend_only_{asset_set}"] = run(trend_filter(equal_weights(asset_set), 10), signal_dates())
    for name, mix in FIXED.items():
        controls[name] = run(fixed_weights(mix), quarter_ends())
    averages = {}
    for name in DESIGNS:
        controls[f"{name}_fixed_average"], average = fixed_at_average(name)
        averages[name] = {k: round(float(v), 4) for k, v in average.items()}
    out["controls"] = {name: {**stats(c), "halves": halves(c)} for name, c in controls.items()}
    out["fixed_average_weights"] = averages

    for asset_set in ASSETS:
        trend, base, only = (stats(c) for c in (curves[f"erc_trend_{asset_set}"], curves[f"erc_{asset_set}"],
                                                controls[f"trend_only_{asset_set}"]))
        out["designs"][f"erc_trend_{asset_set}"]["checks"]["beats_parts"] = all(
            trend["sharpe"] > x["sharpe"] and trend["excess_sharpe"] > x["excess_sharpe"] for x in (base, only))
    for name, result in out["designs"].items():
        result["passes"] = all(result["checks"].values())
        print(name, "PASS" if result["passes"] else "FAIL", result["checks"], flush=True)

    index = curves["erc_spy"].index
    held = {t: prices()[t].reindex(index) for t in ["QQQ", "SPY", "TLT", "GLD", "DBC"]}
    out["buy_and_hold"] = {t: {**stats(c), "halves": halves(c)} for t, c in held.items()}
    bills = prices()[SAFE].reindex(index)
    out["bil"] = {"sharpe": round(sharpe(bills), 6),
                  "cagr": round((bills.iloc[-1] / bills.iloc[0]) ** (365.25 / (index[-1] - index[0]).days) - 1, 6)}
    out["since_listing"] = since_listing({
        "QQQ": prices()["QQQ"], "erc_spy": curves["erc_spy"], "erc_qqq": curves["erc_qqq"],
        "sixty_forty": controls["sixty_forty"]})
    out["levered_to_qqq_volatility"] = {name: levered_to_qqq_volatility(c) for name, c in curves.items()}

    out["risk_shares"] = {name: risk_shares(weights(name, 63, 1)[0], 63)
                          for name in ("ivol_spy", "ivol_qqq", "erc_spy", "erc_qqq")}
    out["month_end_weights"] = {name: {str(d.date()): {k: round(float(v), 4) for k, v in row.items()}
                                       for d, row in weights(name, *DESIGNS[name][2])[0].iterrows()}
                                for name in ("erc_spy", "erc_qqq", "erc_trend_spy", "erc_trend_qqq")}
    r = pd.DataFrame(returns()[["SPY", "TLT", "GLD", "DBC"]]).loc[index[0]:]
    middle = index[len(index) // 2]
    out["correlation_with_spy"] = {
        part: {k: round(float(v), 3) for k, v in r.loc[days].corr()["SPY"].items() if k != "SPY"}
        for part, days in (("first_half", slice(None, middle)), ("second_half", slice(middle, None)))
    }

    listed = {t: prices()[t].dropna() for t in LISTED_FUNDS}
    everything = {**curves, **controls, **held, SAFE: bills, **listed}
    out["yearly"] = {name: yearly(c.dropna()) for name, c in everything.items()}
    frame = pd.DataFrame(everything)
    frame = frame / frame.bfill().iloc[0]
    frame.round(6).to_csv(OUT / "data.csv", index_label="date")
    (OUT / "metrics.json").write_text(json.dumps(out, indent=1, ensure_ascii=False))


if __name__ == "__main__":
    main()