"""
Seven classic asset-allocation portfolios ("lazy portfolios") on US ETFs
versus QQQ (total return).

FinLab research, 2026-10. Reproduces every number in
https://finlab.finance/blog/lazy-portfolio-backtest

Portfolios
  Fixed weights, rebalanced quarterly:
    60/40, Bogleheads three-fund, All Weather, Permanent Portfolio,
    Golden Butterfly
  Rules-based (tactical), rebalanced monthly:
    Faber trend-following allocation (10-month moving average),
    Antonacci dual momentum (Global Equities Momentum)

Three checks; "beats QQQ" means a higher daily Sharpe ratio than holding
QQQ over exactly the same days (daily mean / daily std x sqrt(252), no
risk-free rate deducted).
  1. Neighbouring parameters: the original value and one step either side
     of each of two parameters, 9 portfolios; at least 7 must beat QQQ.
  2. Both halves: split the window at its midpoint; each half must beat QQQ.
  3. Trading costs: the portfolio must beat QQQ under finlab's US default
     costs (no commission, no tax). A 0.1%-per-side run is reported as a
     stress test.

Complete the AI-assisted setup at https://finlab.finance/en/setup first.
Needs finlab 2.2.21 or later (older versions could trade US funds on market holidays).

Run:  python strategy.py    (about 70 backtests; writes data.csv and metrics.json)
"""
import json

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

SIGNAL_START = "2016-12-30"  # first rebalance; month, quarter and year end all fall on this day
WINDOW_START = "2017-02-01"  # dual momentum's first 12-month signal is 2017-01-31; trade the next day
COST_PER_SIDE = 0.001        # stress test: 0.1% per buy and per sell
MIN_GRID_WINS = 7            # of the 9 neighbouring parameter sets, how many must beat QQQ
TRADING_DAYS = 252           # annualises daily statistics
WEIGHT_TOLERANCE = 1e-9      # smaller moves in target weight count as unchanged

FIXED = {
    # name: (weights, tickers counted as stocks)
    "sixty_forty": ({"SPY": 0.60, "AGG": 0.40}, ["SPY"]),
    "all_weather": ({"VTI": 0.30, "TLT": 0.40, "IEF": 0.15, "GLD": 0.075, "DBC": 0.075}, ["VTI"]),
    "permanent": ({"VTI": 0.25, "TLT": 0.25, "GLD": 0.25, "BIL": 0.25}, ["VTI"]),
    "golden_butterfly": ({"VTI": 0.20, "VBR": 0.20, "TLT": 0.20, "SHY": 0.20, "GLD": 0.20}, ["VTI", "VBR"]),
}
FABER_ASSETS = ["SPY", "EFA", "IEF", "VNQ", "DBC"]  # US stocks, foreign stocks, bonds, REITs, commodities
NAMES = ["sixty_forty", "three_fund", "all_weather", "permanent", "golden_butterfly", "faber", "dual_momentum"]
TICKERS = ["SPY", "QQQ", "VTI", "VXUS", "EFA", "VBR", "AGG", "BND", "TLT", "IEF", "SHY", "BIL", "GLD", "DBC", "VNQ"]


# ---------------------------------------------------------------------------
# 1. Prices
# ---------------------------------------------------------------------------
def load_prices():
    """Daily closes of every ETF used here, and the closes on completed month-ends."""
    data.set_market("us_fund")
    close = pd.DataFrame(data.get("us_fund_price:adj_close"))[TICKERS]  # dividend-adjusted closes
    close = close[close["SPY"].notna()]  # the table also lists some US market holidays; drop them
    month_ends = pd.DatetimeIndex(close.index.to_series().groupby(close.index.to_period("M")).max().values)
    if (month_ends[-1] + pd.offsets.BDay(1)).month == month_ends[-1].month:
        month_ends = month_ends[:-1]  # the current month is not over yet
    return close, close.loc[month_ends]


# ---------------------------------------------------------------------------
# 2. Portfolios (each returns target weights; sim() rebalances to them)
# ---------------------------------------------------------------------------
def constant_weights(weights, days):
    position = pd.DataFrame(weights, index=days)
    return position.loc[SIGNAL_START:]


def stock_weight(weights, stock_tickers):
    return sum(weights[t] for t in stock_tickers)


def set_stock_share(weights, stock_tickers, stock_share):
    """Move the stock sleeve to `stock_share`; scale the other assets pro rata."""
    stocks = stock_weight(weights, stock_tickers)
    return {
        t: w / stocks * stock_share if t in stock_tickers else w / (1 - stocks) * (1 - stock_share)
        for t, w in weights.items()
    }


def three_fund(bond_share=0.20, international_share=0.30):
    """Bogleheads: total US market, total international, total bond market."""
    stocks = 1 - bond_share
    return {"VTI": stocks * (1 - international_share), "VXUS": stocks * international_share, "BND": bond_share}


def faber(monthly, months=10, band=0.0):
    """Faber (2007): hold each of five assets (20% each) only while its month-end
    price is above its `months`-month average; otherwise that 20% sits in T-bills."""
    prices = monthly[FABER_ASSETS]
    average = prices.rolling(months).mean()
    in_trend = prices > average * (1 + band)
    weights = in_trend.astype(float) * (1 / len(FABER_ASSETS))
    weights["BIL"] = 1 - weights.sum(axis=1)
    weights = weights[average.notna().all(axis=1)]
    return weights.loc[SIGNAL_START:]


def dual_momentum(monthly, relative_months=12, absolute_months=12):
    """Antonacci (2014): if US stocks beat T-bills over `absolute_months`, hold
    whichever of US and international stocks is stronger over `relative_months`;
    otherwise hold aggregate bonds."""
    prices = monthly[["SPY", "VXUS", "BIL"]]
    relative = prices.pct_change(relative_months, fill_method=None)
    absolute = prices.pct_change(absolute_months, fill_method=None)
    ready = relative[["SPY", "VXUS"]].notna().all(axis=1) & absolute[["SPY", "BIL"]].notna().all(axis=1)

    risk_on = absolute["SPY"] > absolute["BIL"]
    us_leads = relative["SPY"] >= relative["VXUS"]
    weights = pd.DataFrame(
        {
            "SPY": (risk_on & us_leads).astype(float),
            "VXUS": (risk_on & ~us_leads).astype(float),
            "AGG": (~risk_on).astype(float),
        }
    )
    return weights[ready].loc[SIGNAL_START:]


def build(name, close, monthly, **params):
    """Target weights and rebalance frequency for one portfolio and parameter set."""
    if name == "three_fund":
        return constant_weights(three_fund(**params), close.index), "Q"
    if name == "faber":
        return faber(monthly, **params), "M"
    if name == "dual_momentum":
        return dual_momentum(monthly, **params), "M"
    weights, stock_tickers = FIXED[name]
    stock_share = params.get("stock_share", stock_weight(weights, stock_tickers))
    position = constant_weights(set_stock_share(weights, stock_tickers, stock_share), close.index)
    return position, params.get("rebalance", "Q")


def parameter_grid(name):
    """Original value in the middle, one step either side, for two parameters."""
    if name == "three_fund":
        return [{"bond_share": b, "international_share": i} for b in (0.1, 0.2, 0.3) for i in (0.2, 0.3, 0.4)]
    if name == "faber":
        return [{"months": m, "band": b} for m in (8, 10, 12) for b in (-0.01, 0.0, 0.01)]
    if name == "dual_momentum":
        return [{"relative_months": r, "absolute_months": a} for r in (9, 12, 15) for a in (9, 12, 15)]
    weights, stock_tickers = FIXED[name]
    stocks = stock_weight(weights, stock_tickers)
    return [
        {"stock_share": round(stocks + step, 2), "rebalance": freq}
        for step in (-0.10, 0.0, 0.10)
        for freq in ("M", "Q", "Y")
    ]


# ---------------------------------------------------------------------------
# 3. Backtest and statistics
# ---------------------------------------------------------------------------
def backtest(position, rebalance, days, fee_ratio=None):
    report = sim(
        position,
        resample=rebalance,
        trade_at_price="close",  # signal at the month-end close, trade at the next close
        fee_ratio=fee_ratio,     # None = finlab's US default (no commission, no tax)
        upload=False,
    )
    curve = report.creturn
    return curve.reindex(curve.index.intersection(days))  # sim's price table also lists some holidays


def window(curve, start=WINDOW_START, end=None):
    """Clip to [start, end] (never before the portfolio's own first day) and rebase to 1."""
    curve = curve.dropna()
    curve = curve.loc[max(pd.Timestamp(start), curve.index[0]):end]
    return curve / curve.iloc[0]


def split_in_half(days):
    """(label, start, end) for each half of `days`; the midpoint day starts the second half."""
    middle = days[len(days) // 2]
    return (("first_half", days[0], middle), ("second_half", middle, None))


def sharpe(returns):
    return float(returns.mean() / returns.std() * np.sqrt(TRADING_DAYS))


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


def cagr(curve):
    years = (curve.index[-1] - curve.index[0]).days / 365.25
    return float(curve.iloc[-1] ** (1 / years) - 1)


def stats(curve, risk_free=None):
    returns = curve.pct_change().dropna()
    result = {
        "cagr": cagr(curve),
        "total_return": float(curve.iloc[-1] - 1),
        "sharpe": sharpe(returns),
        "volatility": float(returns.std() * np.sqrt(TRADING_DAYS)),
        "max_drawdown": float((curve / curve.cummax() - 1).min()),
        "start": str(curve.index[0].date()),
        "end": str(curve.index[-1].date()),
    }
    if risk_free is not None:
        excess = returns - risk_free.reindex(curve.index).ffill().pct_change().reindex(returns.index)
        result["excess_sharpe"] = sharpe(excess)
    return result


def same_days(prices, days):
    """Buy and hold `prices` over exactly `days`, rebased to 1."""
    return window(prices.reindex(days).ffill(), start=days[0])


def versus_qqq(curve, qqq):
    strategy, benchmark = daily_sharpe(curve), daily_sharpe(same_days(qqq, curve.index))
    return {"sharpe": strategy, "qqq_sharpe": benchmark, "beats_qqq": strategy > benchmark}


def calendar_returns(curve):
    yearly = curve.groupby(curve.index.year).last()
    first = curve.iloc[0]
    return {int(y): float(v) for y, v in (yearly / yearly.shift(1).fillna(first) - 1).items()}


def rebalance_changes(position):
    """Month-ends on which the target weights changed (tactical portfolios)."""
    changed = position.diff().abs().sum(axis=1) > WEIGHT_TOLERANCE
    return int(changed.iloc[1:].sum())


# ---------------------------------------------------------------------------
# 4. The three checks for one portfolio
# ---------------------------------------------------------------------------
def evaluate(name, close, monthly):
    qqq = close["QQQ"]

    grid, neighbours = [], []
    for params in parameter_grid(name):
        neighbour = window(backtest(*build(name, close, monthly, **params), close.index))
        neighbours.append(neighbour)
        grid.append({"params": params, **versus_qqq(neighbour, qqq), "start": str(neighbour.index[0].date())})
    wins = sum(row["beats_qqq"] for row in grid)
    curve = neighbours[len(neighbours) // 2]  # the middle of the grid is the original parameter set

    halves = {}
    for label, start, end in split_in_half(curve.index):
        part = window(curve, start, end)
        halves[label] = {**versus_qqq(part, qqq), "start": str(part.index[0].date()), "end": str(part.index[-1].date())}

    position, rebalance = build(name, close, monthly)
    costly = window(backtest(position, rebalance, close.index, fee_ratio=COST_PER_SIDE))
    full = {**stats(curve, close["BIL"]), **versus_qqq(curve, qqq)}

    result = {
        "rebalance": rebalance,
        "full": full,
        "grid": grid,
        "grid_wins": wins,
        "halves": halves,
        "stress_0_1pct": {**versus_qqq(costly, qqq), "cagr": cagr(costly)},
        "calendar_returns": calendar_returns(curve),
        "passes": wins >= MIN_GRID_WINS and all(h["beats_qqq"] for h in halves.values()) and full["beats_qqq"],
    }
    if rebalance == "M":
        held = position.loc[curve.index[0] - pd.offsets.MonthEnd(1):]
        result["weight_changes"] = rebalance_changes(held)
        result["months"] = len(held)
        result["month_end_weights"] = {
            str(day.date()): {t: round(float(w), 4) for t, w in row.items() if w} for day, row in held.iterrows()
        }
    return curve, result


def context(close, days):
    """Numbers the article uses to explain the results."""
    hedges = ["AGG", "TLT", "GLD", "DBC"]
    daily = close[["SPY", *hedges]].loc[WINDOW_START:].pct_change().dropna()
    correlation, asset_cagr_halves = {}, {}
    for label, start, end in split_in_half(days):
        part = daily.loc[start:end]
        correlation[label] = {t: float(part["SPY"].corr(part[t])) for t in hedges}
        asset_cagr_halves[label] = {t: cagr(window(close[t], start, end)) for t in TICKERS}
    year_2022 = {
        t: float(close[t].loc["2022"].iloc[-1] / close[t].loc[:"2021"].iloc[-1] - 1)
        for t in ("SPY", "QQQ", "AGG", "BND", "TLT", "IEF", "GLD", "DBC", "VXUS", "VNQ", "BIL")
    }
    asset_cagr = {t: cagr(window(close[t])) for t in TICKERS}
    return {"spy_correlation": correlation, "year_2022": year_2022, "asset_cagr": asset_cagr,
            "asset_cagr_halves": asset_cagr_halves}


def main():
    finlab.login()  # finlab guides you through login automatically
    close, monthly = load_prices()

    curves, results = {}, {}
    for name in NAMES:
        curves[name], results[name] = evaluate(name, close, monthly)
        r = results[name]
        print(
            f"{name:17s} Sharpe {r['full']['sharpe']:.2f} vs QQQ {r['full']['qqq_sharpe']:.2f} | "
            f"neighbours {r['grid_wins']}/{len(r['grid'])} | halves "
            + " / ".join(f"{h['sharpe']:.2f} vs {h['qqq_sharpe']:.2f}" for h in r["halves"].values())
            + f" | {COST_PER_SIDE:.1%} per side {r['stress_0_1pct']['sharpe']:.2f}"
            + f" | {'PASS' if r['passes'] else 'fail'}"
        )

    days = curves["sixty_forty"].index
    benchmarks = {t: same_days(close[t], days) for t in ("QQQ", "SPY")}
    metrics = {
        "window": [str(days[0].date()), str(days[-1].date())],
        "portfolios": results,
        "benchmarks": {
            t: {
                **stats(c, close["BIL"]),
                "halves": {label: daily_sharpe(window(c, start, end)) for label, start, end in split_in_half(days)},
                "calendar_returns": calendar_returns(c),
            }
            for t, c in benchmarks.items()
        },
        "context": context(close, days),
    }
    pd.DataFrame({**curves, **benchmarks}).round(6).to_csv("data.csv", index_label="date")
    with open("metrics.json", "w") as f:
        json.dump(metrics, f, indent=1)


if __name__ == "__main__":
    main()
