"""
US sector ETF rotation driven by the fundamentals and prices of each sector's stocks.

FinLab research, 2026-10. Reproduces every number in
https://finlab.finance/blog/sector-rotation-strategy-backtest

Every month, the 1,000 most traded US common stocks (delisted ones included) are
grouped into the 11 sectors that match the SPDR sector ETFs (XLK, XLF, XLV, ...).
Each sector gets the median of a signal across its stocks (for moving-average
breadth, the share of its stocks above the average), and the portfolio holds the
ETFs of the 2 or 3 top-ranked sectors at equal weight until the next month end.

Signals
  revenue_ttm              trailing-12-month revenue, year-over-year growth
  fcf_yield                trailing-12-month free cash flow / market value
  momentum_12_1            12-month price return, skipping the latest month
  above_200d               share of stocks above their 200-day moving average
  etf_momentum_6m          the sector ETF's own 6-month return
  operating_margin_change  trailing operating margin minus a year earlier
  above_50d                share of stocks above their 50-day moving average
  near_52w_high            price / 52-week high
  eps_ttm                  trailing-12-month diluted EPS, year-over-year growth
  gross_margin_change      trailing gross margin minus a year earlier

"Beats QQQ" means a higher daily Sharpe ratio (no risk-free rate) than buying and
holding QQQ, dividends included, over exactly the same dates. A design passes only if
  1. at least 7 of its 9 neighbouring parameter sets beat QQQ,
  2. it beats QQQ in both halves of the backtest, and
  3. it still beats QQQ after paying 0.1% per buy and per sell (finlab's US default
     cost is zero).
Each design is screened at its original settings first; only a design that beats
QQQ there gets the grid, the halves and the cost run.

Run:  python research.py            (screen all 20 designs, a few minutes)
      python research.py --grid     (adds grids, halves, costs, XLK and start-date checks)
Writes metrics.json and data.csv next to this file.

Needs finlab 2.2.21 or later. Earlier versions keep 55 US market-holiday rows in
us_fund_price as flat trading days, which lowers QQQ's Sharpe from 0.900 to 0.893
and lets one revenue growth setting edge past it (see the article).
"""
import json
import sys
from functools import cache
from pathlib import Path

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

START = "2018-04-02"     # first month the 12-month revenue and margin signals cover 800+ stocks
EARLY_START = "2017-08-01"
END = "2026-10-28"
UNIVERSE_SIZE = 1000
MIN_STOCKS = 5           # a sector needs at least this many stocks with a value
COST_PER_SIDE = 0.001
OUT = Path(__file__).resolve().parent

SECTOR_ETF = {
    "Technology": "XLK",
    "Financial Services": "XLF",
    "Healthcare": "XLV",
    "Consumer Cyclical": "XLY",
    "Consumer Defensive": "XLP",
    "Energy": "XLE",
    "Industrials": "XLI",
    "Basic Materials": "XLB",
    "Utilities": "XLU",
    "Real Estate": "XLRE",
    "Communication Services": "XLC",
}
ETFS = list(SECTOR_ETF.values())
SIGNALS = ["revenue_ttm", "fcf_yield", "momentum_12_1", "above_200d", "etf_momentum_6m",
           "operating_margin_change", "above_50d", "near_52w_high", "eps_ttm", "gross_margin_change"]
BREADTH = {"above_200d": 200, "above_50d": 50}   # share of stocks above this average; every other signal is a median
GRIDS = {
    "revenue_ttm": [(top, signal) for signal in ("revenue_q", "revenue_2q", "revenue_ttm") for top in (1, 2, 3)],
    "momentum_12_1": [(top, f"momentum_{months}_1") for months in (9, 12, 15) for top in (1, 2, 3)],
}

data.truncate_end = END


@cache
def get(key, market="us"):
    data.set_market(market)
    return data.get(key)


def prices():
    return get("us_price:adj_close")


def etf_prices():
    return get("us_fund_price:adj_close", "us_fund")[ETFS + ["QQQ", "SPY"]]


def trading_days(start):
    return etf_prices()["QQQ"].dropna().loc[start:END].index


@cache
def rebalance_dates(start=START):
    """The first trading day, then the last trading day of every month."""
    days = trading_days(start)
    ends = pd.Series(days, index=days).groupby(days.year * 12 + days.month).max()
    return pd.DatetimeIndex([days[0], *ends]).unique()


@cache
def all_dates():
    return rebalance_dates(EARLY_START).union(rebalance_dates(START))


def on_dates(frame):
    """Any frame -> rebalance date x ticker; quarterly items carry their latest filed value."""
    if len(frame.index) and isinstance(frame.index[0], str):
        frame = frame.index_str_to_date()
    frame = frame.reindex(columns=prices().columns)
    return frame.reindex(frame.index.union(all_dates())).ffill().reindex(all_dates())


# ---------------------------------------------------------------------------
# Universe and sectors
# ---------------------------------------------------------------------------
@cache
def sector_of():
    """Sector of every common stock (delisted ones included), one share class per company."""
    p = get("us_company_profile")
    keep = (
        (p["security_type"] == "common_stock")
        & ~p["is_etf"].fillna(False).astype(bool)
        & ~p["is_fund"].fillna(False).astype(bool)
        & ~p["is_adr"].fillna(False).astype(bool)
        & p["sector"].isin(list(SECTOR_ETF))
    )
    rows = p.loc[keep, ["stock_id", "company_name", "sector"]].astype(str).sort_values("stock_id")
    named = rows["company_name"].str.len() > 0
    rows = pd.concat([rows[named].drop_duplicates("company_name"), rows[~named]])
    rows = rows[rows["stock_id"].isin(prices().columns)]
    return rows.set_index("stock_id")["sector"]


@cache
def universe():
    """The 1,000 most traded stocks with price > $5, 60-day average dollar volume
    > $5M and no +/-50% day in the past year."""
    close, volume = get("us_price:close"), get("us_price:volume")
    dollar_volume = (close * volume).rolling(60, min_periods=20).mean()
    daily_return = prices().pct_change(fill_method=None)
    recently_extreme = (daily_return.abs() >= 0.5).astype(float).rolling(252, min_periods=1).max() > 0
    ok = (close > 5) & (dollar_volume > 5e6) & ~recently_extreme
    liquid = dollar_volume[ok].reindex(columns=sector_of().index).reindex(all_dates())
    return liquid.is_largest(UNIVERSE_SIZE)


# ---------------------------------------------------------------------------
# Stock-level signals (rebalance date x ticker)
# ---------------------------------------------------------------------------
@cache
def quarterly(key):
    quarters = get("us_income_statement:revenue").index
    return get(key).reindex(index=quarters, columns=prices().columns)


def ttm(key):
    return quarterly(key).rolling(4).sum()


def growth(now, before):
    return now / before[before > 0] - 1


def ttm_margin_change(key):
    revenue = ttm("us_income_statement:revenue")
    margin = ttm(key) / revenue[revenue > 0]
    return margin - margin.shift(4)


@cache
def stock_signal(name):
    if name == "revenue_ttm":
        revenue_12m = ttm("us_income_statement:revenue")
        return on_dates(growth(revenue_12m, revenue_12m.shift(4)))
    if name in ("revenue_q", "revenue_2q"):
        revenue = quarterly("us_income_statement:revenue")
        latest = growth(revenue, revenue.shift(4))   # latest quarter vs the same quarter a year earlier
        if name == "revenue_q":
            return on_dates(latest)
        return on_dates((latest + latest.shift(1)) / 2)   # average of the last two quarters' growth
    if name == "eps_ttm":
        eps = ttm("us_income_statement:eps_diluted")
        return on_dates(growth(eps, eps.shift(4)))
    if name == "operating_margin_change":
        return on_dates(ttm_margin_change("us_income_statement:operating_income"))
    if name == "gross_margin_change":
        return on_dates(ttm_margin_change("us_income_statement:gross_profit"))
    if name == "fcf_yield":
        cap = on_dates(get("etl:us_market_value"))
        return on_dates(ttm("us_cash_flow:free_cash_flow")) / cap[cap > 0]
    p = prices()
    if name.startswith("momentum_"):    # momentum_12_1: 12-month return, skipping the latest month
        months = int(name.split("_")[1])
        return on_dates(p.shift(21) / p.shift(21 * months) - 1)
    if name == "near_52w_high":
        return on_dates(p / p.rolling(252, min_periods=200).max())
    if name in BREADTH:
        average = p.rolling(BREADTH[name]).mean()
        return on_dates((p > average).astype(float).where(average.notna()))
    raise KeyError(name)


# ---------------------------------------------------------------------------
# Sector scores and portfolios
# ---------------------------------------------------------------------------
def listed(dates, exclude=()):
    """Which sector ETFs trade on each date."""
    return etf_prices().reindex(dates)[ETFS].drop(columns=list(exclude)).notna()


@cache
def sector_scores(name):
    """Rebalance date x sector ETF; NaN before the ETF trades or with too few stocks."""
    trading = listed(all_dates())
    if name == "etf_momentum_6m":
        p = etf_prices()[ETFS]
        return (p / p.shift(126) - 1).reindex(all_dates()).where(trading)
    signal = stock_signal(name).reindex(columns=sector_of().index)[universe()]
    by_sector = signal.T.groupby(sector_of())
    stat = by_sector.mean() if name in BREADTH else by_sector.median()
    scores = stat.where(by_sector.count() >= MIN_STOCKS).T.rename(columns=SECTOR_ETF)
    return scores.reindex(columns=ETFS).where(trading)


def picks(name, top, exclude=(), start=START):
    """Top-ranked sectors on each rebalance date; a sector with no score is never picked."""
    scores = sector_scores(name).drop(columns=list(exclude)).loc[rebalance_dates(start)]
    return scores.rank(axis=1, ascending=False, method="first") <= top


def traded(chosen):
    """Drop the pick made on the last day: it never trades inside the window."""
    return chosen.iloc[:-1]


def weights_from(chosen):
    chosen = chosen.astype(float)
    return chosen.div(chosen.sum(axis=1), axis=0).fillna(0)


def run(weights, fee_ratio=None):
    days = trading_days(weights.index[0])
    position = weights.reindex(days, method="ffill").fillna(0)
    data.set_market("us_fund")
    report = sim(position, resample=weights.index, trade_at_price="close",
                 fee_ratio=fee_ratio, upload=False)
    # the first trade fills at the next close; start every curve (and QQQ) there
    return report.creturn.loc[days[1]:END]


@cache
def backtest(name, top, exclude=(), fee_ratio=None, start=START):
    return run(weights_from(picks(name, top, exclude, start)), fee_ratio)


def equal_weight_sectors(exclude=()):
    """Every sector ETF that trades, equal weight, rebalanced monthly."""
    return run(weights_from(listed(rebalance_dates(), exclude)))


# ---------------------------------------------------------------------------
# Statistics
# ---------------------------------------------------------------------------
def daily_sharpe(curve):
    returns = curve.pct_change().dropna()
    return returns.mean() / returns.std() * 252 ** 0.5


def buy_and_hold(ticker, index):
    return etf_prices()[ticker].reindex(index).ffill()


def stats(curve):
    curve = curve.dropna()
    bench = buy_and_hold("QQQ", curve.index)
    years = (curve.index[-1] - curve.index[0]).days / 365.25
    returns = curve.pct_change().dropna()
    return {
        "start": str(curve.index[0].date()), "end": str(curve.index[-1].date()),
        "sharpe": round(daily_sharpe(curve), 6), "qqq_sharpe": round(daily_sharpe(bench), 6),
        "cagr": round((curve.iloc[-1] / curve.iloc[0]) ** (1 / years) - 1, 6),
        "volatility": round(returns.std() * 252 ** 0.5, 6),
        "mdd": round((curve / curve.cummax() - 1).min(), 6),
    }


def beats_qqq(s):
    return bool(s["sharpe"] > s["qqq_sharpe"])


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


def months_held(chosen):
    return {etf: int(n) for etf, n in chosen.sum().items()}


def turnover(chosen):
    """One-way turnover per year: half the summed weight changes at each rebalance."""
    w = weights_from(traded(chosen))
    changes = w.diff().abs().sum(axis=1).iloc[1:].sum() / 2
    years = (w.index[-1] - w.index[0]).days / 365.25
    return round(changes / years, 4)


def rank_ic(name):
    """Mean Spearman correlation between sector scores and the next month's ETF return,
    and the average next-month return of the sector ranked 1st, 2nd, ... over the
    equal-weight average of the ranked sectors."""
    dates = rebalance_dates()
    scores = sector_scores(name).loc[dates]
    following = etf_prices()[ETFS].reindex(dates).pct_change(fill_method=None).shift(-1)
    ics = pd.Series([
        scores.loc[d].corr(following.loc[d], method="spearman")
        for d in dates[:-1]
        if scores.loc[d].notna().sum() >= 5 and following.loc[d].notna().sum() >= 5
    ]).dropna()
    rank = scores.rank(axis=1, ascending=False, method="first")
    excess = following.sub(following.where(scores.notna()).mean(axis=1), axis=0)
    by_rank = {int(k): round(excess[rank == k].stack().mean(), 6) for k in range(1, len(ETFS) + 1)}
    return {"mean": round(ics.mean(), 6), "t": round(ics.mean() / ics.std() * len(ics) ** 0.5, 4),
            "hit_rate": round((ics > 0).mean(), 6), "months": len(ics),
            "next_month_excess_by_rank": by_rank}


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


def screen(start=START):
    result = {}
    for name in SIGNALS:
        for top in (2, 3):
            curve = backtest(name, top, start=start)
            s = stats(curve)
            s["halves"] = halves(curve)
            s["beats_qqq"] = beats_qqq(s)
            s["turnover"] = turnover(picks(name, top, start=start))
            result[f"{name}/top{top}"] = s
            print(f"{name:<24} top{top}  Sharpe {s['sharpe']:.3f} vs QQQ {s['qqq_sharpe']:.3f}"
                  f"  CAGR {s['cagr']:.1%}  MDD {s['mdd']:.1%}", flush=True)
    return result


def validate(name, top=2, start=START):
    curve = backtest(name, top, start=start)
    split = halves(curve)
    grid = []
    for t, variant in GRIDS[name]:
        s = stats(backtest(variant, t, start=start))
        grid.append({"top": t, "signal": variant, **s, "beats_qqq": beats_qqq(s),
                     "months_held": months_held(traded(picks(variant, t, start=start)))})
    grid_wins = sum(g["beats_qqq"] for g in grid)
    cost = stats(backtest(name, top, fee_ratio=COST_PER_SIDE, start=start))
    return {
        "signal": name, "top": top, "full": stats(curve), "halves": split, "grid": grid,
        "grid_wins": grid_wins, "cost_0.1pct": cost,
        "passes": grid_wins >= 7 and all(map(beats_qqq, split)) and beats_qqq(cost),
    }


def half_in_xlk(chosen):
    """50% XLK, the other 50% split equally across the chosen ETFs, rebalanced monthly."""
    weights = weights_from(chosen) / 2
    weights["XLK"] = 0.5
    return stats(run(weights[ETFS]))


def xlk_check(name, top=2):
    """Is the rotation just holding XLK?"""
    chosen = traded(picks(name, top))
    curve = backtest(name, top)
    ex_xlk = backtest(name, top, exclude=("XLK",))
    return {
        "before_2022": stats(curve.loc[:"2021-12-31"]), "year_2022": stats(curve.loc["2021-12-31":"2022-12-30"]),
        "months": len(chosen), "months_with_xlk": int(chosen["XLK"].sum()),
        "months_held": months_held(chosen),
        "xlk_buy_hold": stats(buy_and_hold("XLK", curve.index)),
        "ex_xlk": stats(ex_xlk), "ex_xlk_halves": halves(ex_xlk),
        "equal_weight_ex_xlk": stats(equal_weight_sectors(exclude=("XLK",))),
        "xlk_half_equal_weight_rest": half_in_xlk(listed(rebalance_dates(), exclude=("XLK",))),
        "xlk_half_best_other": half_in_xlk(picks(name, 1, exclude=("XLK",))),
        "holdings": {str(d.date()): [etf for etf in ETFS if chosen.loc[d, etf]] for d in chosen.index},
        "top1": {
            signal: {
                "sharpe": round(daily_sharpe(backtest(signal, 1)), 6),
                "ex_xle": round(daily_sharpe(backtest(signal, 1, exclude=("XLE",))), 6),
                "ex_xlk": round(daily_sharpe(backtest(signal, 1, exclude=("XLK",))), 6),
                "yearly": yearly(backtest(signal, 1)),
            }
            for signal in ("revenue_q", "revenue_ttm")
        },
    }


def coverage():
    """Stocks per sector, and when each statement signal first covers 800 stocks."""
    members = universe().T.groupby(sector_of()).sum().T
    revenue = get("us_income_statement:revenue").reindex(columns=sector_of().index)
    first = {"revenue_filers_per_quarter": {q: int(n) for q, n in revenue.notna().sum(axis=1).iloc[:8].items()}}
    for name in ("revenue_ttm", "revenue_q", "fcf_yield", "eps_ttm", "operating_margin_change"):
        counted = stock_signal(name).reindex(columns=sector_of().index)[universe()].notna().sum(axis=1)
        first[name] = {"at_early_start": int(counted.iloc[0]),
                       "first_800": str(counted[counted >= 800].index[0].date())}
    return {"stocks_per_sector": {str(d.date()): {k: int(v) for k, v in members.loc[d].items()}
                                  for d in (members.index[0], pd.Timestamp(START), members.index[-1])},
            "signal_coverage": first}


def main():
    out = {"window": {"start": START, "end": END, "early_start": EARLY_START},
           "settings": {"universe_size": UNIVERSE_SIZE, "min_stocks": MIN_STOCKS,
                        "cost_per_side": COST_PER_SIDE}}
    out["coverage"] = coverage()
    out["screen"] = screen()
    out["rank_ic"] = {name: rank_ic(name) for name in SIGNALS}
    curves = {
        "revenue_ttm_top2": backtest("revenue_ttm", 2),
        "revenue_ttm_top3": backtest("revenue_ttm", 3),
        "momentum_12_1_top2": backtest("momentum_12_1", 2),
        "revenue_ttm_top2_ex_xlk": backtest("revenue_ttm", 2, exclude=("XLK",)),
        "revenue_ttm_top1": backtest("revenue_ttm", 1),
        "revenue_q_top1": backtest("revenue_q", 1),
        "equal_weight_sectors": equal_weight_sectors(),
    }
    index = curves["revenue_ttm_top2"].index
    benchmarks = {t: buy_and_hold(t, index) for t in ["QQQ", "SPY", *ETFS]}
    out["buy_and_hold"] = {t: stats(c) for t, c in benchmarks.items()}
    out["equal_weight_sectors"] = stats(curves["equal_weight_sectors"])
    out["yearly"] = {name: yearly(c) for name, c in {**curves, **benchmarks}.items()}

    if "--grid" in sys.argv:
        out["validation"] = {name: validate(name) for name in GRIDS}
        out["xlk_check"] = xlk_check("revenue_ttm")
        out["early_start"] = {"screen": screen(EARLY_START),
                              "validation": {name: validate(name, start=EARLY_START) for name in GRIDS}}
        for name, v in out["validation"].items():
            print(name, "PASS" if v["passes"] else "FAIL", f"grid {v['grid_wins']}/9",
                  "halves", [(round(h["sharpe"], 2), round(h["qqq_sharpe"], 2)) for h in v["halves"]],
                  "0.1% cost", round(v["cost_0.1pct"]["sharpe"], 2), flush=True)

    frame = pd.DataFrame({**curves, **{f"{t} (buy and hold)": c for t, c in benchmarks.items()}})
    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()