"""
Buy the dip, share buybacks and the Altman Z-score as trading rules (FinLab research, 2026-11).

Five designs, each on a parameter grid fixed before any result was seen. Starred values are the
original settings.

  dip               Buy QQQ on the first close that is X% or more below its highest close of the
                    last N sessions. Sell after H sessions, or as soon as the close gets back to the
                    N-session high recorded at entry, whichever comes first. Cash otherwise.
                    N 63 / 126* / 252, X 5% / 10%* / 15%, H 21 / 63* / 126 (27 settings).
  scale_in          Hold a base weight w of QQQ. While the dip rule (N 126, H 63) is on, hold 100%.
                    Rebalance to the target at every month end and whenever the rule switches.
                    w 25% / 50%* / 75%, X 5% / 10%* / 15% (9 settings).
  buyback           Buyback yield: shares repurchased over the last K quarters (gross, or net of
                    shares issued) divided by market value. Hold the N highest yields above zero,
                    equal weight. K 4* / 8, gross* / net, N 20 / 30* / 50 (12 settings).
  exclude_distress  Hold every stock in the pool whose Altman Z-score is at least T, equal weight.
                    T 1.5 / 1.81* / 2.1, pool 200 / 500* / 1000 (9 settings).
  safe_zone         Hold every stock in the pool whose Z-score is above T, equal weight.
                    T 2.5 / 2.99* / 3.5, pool 200 / 500* / 1000 (9 settings).

Dates. Prices from 2016 are warm-up; results run from 2017-04-03 (the first day eight quarters of
buybacks are on file) to 2026-11-10. A signal read at one day's close trades at the next day's
close (finlab sim default). Statement values are used from their filing date (FinLab aligns US
statements on it). The stock designs rebalance on the first day and on the last trading day of
every quarter.

Universe. Common stocks, delisted ones included, one share class per company, no ETFs, funds or
ADRs. On the rebalance day a stock must close above $5, have traded more than $5M a day on average
over 60 days (raw close x volume) and have had no +/-50% day in the past 252 sessions. The buyback
design ranks the 1,000 most traded of these and treats a yield above 50% as a data error. The
Z-score designs leave out financial companies (the formula was built for manufacturers) and keep
the 200 / 500 / 1,000 most traded stocks with a computable score.

Altman (1968) Z = 1.2 working capital / total assets + 1.4 retained earnings / total assets
                + 3.3 EBIT (last four quarters) / total assets + 0.6 market value / total liabilities
                + 1.0 revenue (last four quarters) / total assets.
Market value is FinLab's etl:us_market_value: latest quarterly share count times the close, with
later stock splits taken out.

"Beats" means a higher daily Sharpe ratio (no risk-free rate) than QQQ with dividends over the same
dates. A design passes only if
  1. at least 70% of its settings beat QQQ at sim's default cost (zero for US stocks and ETFs),
  2. its original setting beats QQQ in both halves of the period,
  3. its original setting beats QQQ at 0.1% per side,
  4. its original setting beats QQQ on returns above T-bills (BIL),
  5. its original setting beats QQQ without the calendar year of its largest lead over QQQ,
  6. (dip and scale_in) with idle cash in BIL, it beats holding QQQ all the time at the same
     average weight, the rest in BIL.

Run:  pip install "finlab>=2.3.8"
      python research.py                 writes metrics.json and data.csv into this folder and the
                                         buyback and Z-score article folders next to it
      python research.py --reports DIR   also writes the FinLab report of each original setting
"""
import argparse
import json
from datetime import date
from functools import cache
from itertools import product
from pathlib import Path

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

START, END = "2017-04-03", "2026-11-10"
HERE = Path(__file__).resolve().parent
OUT = {"dip": HERE, "buyback": HERE.parent / "stock-buyback-strategy", "altman": HERE.parent / "altman-z-score-backtest"}
COST_PER_SIDE = 0.001
MIN_PRICE, MIN_DOLLAR_VOLUME = 5, 5e6
BUYBACK_POOL, MAX_YIELD = 1000, 0.5
DISTRESS, SAFE = 1.81, 2.99

GRIDS = {
    "dip": {"lookback": (63, 126, 252), "drop": (0.05, 0.10, 0.15), "hold": (21, 63, 126)},
    "scale_in": {"base": (0.25, 0.50, 0.75), "drop": (0.05, 0.10, 0.15)},
    "buyback": {"quarters": (4, 8), "measure": ("gross", "net"), "top": (20, 30, 50)},
    "exclude_distress": {"threshold": (1.5, DISTRESS, 2.1), "pool": (200, 500, 1000)},
    "safe_zone": {"threshold": (2.5, SAFE, 3.5), "pool": (200, 500, 1000)},
}
ORIGINAL = {
    "dip": {"lookback": 126, "drop": 0.10, "hold": 63},
    "scale_in": {"base": 0.50, "drop": 0.10},
    "buyback": {"quarters": 4, "measure": "gross", "top": 30},
    "exclude_distress": {"threshold": DISTRESS, "pool": 500},
    "safe_zone": {"threshold": SAFE, "pool": 500},
}
QQQ_TIMING = {"dip", "scale_in"}


def use(market):
    data.set_market(market)
    data.truncate_end = END


@cache
def get(key):
    return data.get(key)


def frame(key):
    return pd.DataFrame(get(key)).loc[:END]


@cache
def qqq():
    return frame("us_fund_price:adj_close")["QQQ"].dropna()


@cache
def calendar():
    return qqq().loc[START:].index


def on_calendar(price):
    r = price.reindex(qqq().index).ffill().pct_change().reindex(calendar()).fillna(0)
    r.iloc[0] = 0.0
    return r


@cache
def benchmark():
    return on_calendar(qqq())


@cache
def bil():
    return on_calendar(frame("us_fund_price:adj_close")["BIL"])


def sharpe(r):
    return float(r.mean() / r.std() * np.sqrt(252)) if r.std() > 0 else 0.0


def stats(r):
    equity = (1 + r).cumprod()
    years = (r.index[-1] - r.index[0]).days / 365.25
    return {"cagr": float(equity.iloc[-1] ** (1 / years) - 1), "sharpe": sharpe(r),
            "mdd": float((equity / equity.cummax() - 1).min())}


def halves(r):
    """Daily Sharpe of the first and second half of the trading days."""
    middle = len(r) // 2
    return [sharpe(r.iloc[:middle]), sharpe(r.iloc[middle:])]


def yearly(r):
    return r.groupby(r.index.year).apply(lambda x: (1 + x).prod() - 1)


def backtest(position, resample=None, fee=None):
    # Stocks never held cannot move the result, and sim copies every column it gets.
    held = position.any()
    if held.any():
        position = position.loc[:, held]
    return sim(position, resample=resample, upload=False, end_date=END, fee_ratio=fee)


def daily(report):
    return report.creturn.reindex(calendar()).ffill().fillna(1).pct_change().fillna(0)


# Buy the dip on QQQ

def dip_state(lookback, drop, hold):
    """True on the closes the dip rule wants to hold QQQ."""
    price = qqq()
    high = price.rolling(lookback, min_periods=lookback).max()
    hit = (price / high - 1 <= -drop).to_numpy()
    p, h = price.to_numpy(), high.to_numpy()
    state = np.zeros(len(p), dtype=bool)
    inside, target, age = False, np.nan, 0
    for i in range(1, len(p)):
        if inside:
            age += 1
            inside = p[i] < target and age < hold
        elif hit[i] and not hit[i - 1]:
            inside, target, age = True, h[i], 0
        state[i] = inside
    return pd.Series(state, index=price.index).loc[START:]


def dip_position(setting):
    return pd.DataFrame({"QQQ": dip_state(**setting)}), None


def scale_in_position(setting):
    dip = ORIGINAL["dip"]
    state = dip_state(dip["lookback"], setting["drop"], dip["hold"])
    weight = setting["base"] + (1 - setting["base"]) * state.astype(float)
    month_ends = state.index.to_series().groupby(state.index.to_period("M")).max()
    switches = state.index[state.ne(state.shift())]  # includes the first day
    dates = pd.DatetimeIndex(sorted(set(month_ends) | set(switches)))
    return pd.DataFrame({"QQQ": weight}), dates


def exposure(strategy):
    """QQQ weight held on each day, read off a zero-cost QQQ-plus-cash portfolio."""
    q = benchmark()
    e = (strategy / q).where(q.abs() > 1e-8).ffill().fillna(0)
    return e.clip(0, 1)


def dip_events(setting):
    """Each entry of the dip rule: how it ended, its return and what QQQ did afterwards."""
    state = dip_state(**setting)
    price, days = qqq().loc[START:], state.index
    runs = state.ne(state.shift()).cumsum()
    starts = days[state & ~state.shift(fill_value=False)]
    trades = []
    for start, held in zip(starts, state[state].groupby(runs[state]).size()):
        first = days.get_loc(start)
        trade = {"signal": str(start.date()), "sessions": int(held)}
        if first + held + 1 < len(days):
            trade["exit"] = "time" if held == setting["hold"] else "recovered"
            trade["return"] = float(price.iloc[first + held + 1] / price.iloc[first + 1] - 1)
        else:
            trade["exit"] = "open"
        trades.append(trade)
    forward = {}
    for n in (21, 63, 126):
        change = price.shift(-n) / price - 1
        after = change.reindex(starts).dropna()
        every_day = change.dropna()
        forward[f"qqq_return_{n}d"] = {"mean": float(after.mean()), "positive": int((after > 0).sum()), "count": len(after),
                                       "all_days_mean": float(every_day.mean())}
    return {"trades": trades, **forward}


# US stocks

@cache
def common_stocks():
    """Common stocks (delisted ones included), one share class per company."""
    p = pd.DataFrame(get("us_company_profile")).astype({"stock_id": str}).drop_duplicates("stock_id")
    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)
    )
    rows = p.loc[keep, ["stock_id", "company_name", "sector"]].sort_values("stock_id").drop_duplicates("company_name")
    return rows[rows["stock_id"].isin(get("us_price:adj_close").columns)].set_index("stock_id")


@cache
def rebalance_days():
    days = calendar().to_series()
    quarter_ends = days.groupby(days.index.to_period("Q")).max()
    return pd.DatetimeIndex(sorted({days.iloc[0], *quarter_ends[quarter_ends < days.iloc[-1]]}))


def on_rebalance(df):
    return pd.DataFrame(df).reindex(index=rebalance_days(), columns=common_stocks().index)


@cache
def liquidity():
    """Eligible stocks and 60-day dollar volume on each rebalance day."""
    cols = common_stocks().index
    adj, close, volume = (frame(k).reindex(columns=cols).dropna(how="all")
                          for k in ("us_price:adj_close", "us_price:close", "us_price:volume"))
    dollar = (close * volume).rolling(60, min_periods=20).mean()
    extreme = adj.pct_change(fill_method=None).abs().ge(0.5).rolling(252, min_periods=1).max().fillna(0).astype(bool)
    eligible = close.gt(MIN_PRICE) & dollar.gt(MIN_DOLLAR_VOLUME) & ~extreme
    return on_rebalance(eligible).fillna(False).astype(bool), on_rebalance(dollar)


def most_traded(mask, n):
    _, dollar = liquidity()
    return mask & dollar.where(mask).rank(axis=1, ascending=False, method="first").le(n)


@cache
def buyback_yield(quarters, measure):
    repurchased = -get("us_cash_flow:common_stock_repurchased")
    if measure == "net":
        repurchased = repurchased - get("us_cash_flow:common_stock_issuance")
    total = repurchased.rolling(quarters, min_periods=quarters).sum()
    return on_rebalance(total / get("etl:us_market_value"))


def buyback_pool():
    eligible, _ = liquidity()
    return most_traded(eligible, BUYBACK_POOL)


def buyback_position(setting):
    y = buyback_yield(setting["quarters"], setting["measure"])
    y = y.where(buyback_pool() & y.gt(0) & y.le(MAX_YIELD))
    return y.rank(axis=1, ascending=False, method="first").le(setting["top"]), rebalance_days()


@cache
def z_score():
    bs = lambda f: get("us_balance_sheet:" + f)
    four = lambda f: get("us_income_statement:" + f).rolling(4, min_periods=4).sum()
    assets = bs("total_assets")
    book = (1.2 * (bs("total_current_assets") - bs("total_current_liabilities")) + 1.4 * bs("retained_earnings")
            + 3.3 * four("ebit") + 1.0 * four("revenue")) / assets
    z = on_rebalance(book + 0.6 * get("etl:us_market_value") / bs("total_liabilities"))
    return z.where(np.isfinite(z))


def z_pool(n):
    eligible, _ = liquidity()
    not_financial = common_stocks()["sector"].ne("Financial Services")
    return most_traded(eligible & z_score().notna() & not_financial, n)


def exclude_distress_position(setting):
    return z_pool(setting["pool"]) & z_score().ge(setting["threshold"]), rebalance_days()


def safe_zone_position(setting):
    return z_pool(setting["pool"]) & z_score().gt(setting["threshold"]), rebalance_days()


POSITIONS = {"dip": dip_position, "scale_in": scale_in_position, "buyback": buyback_position,
             "exclude_distress": exclude_distress_position, "safe_zone": safe_zone_position}


# Validation

def validate(design, rows, center, stressed):
    q, cash, days = benchmark(), bil(), calendar()
    center_years, q_years = yearly(center), yearly(q)
    lead = center_years - q_years
    lead_year = int(lead.idxmax())
    keep = days.year != lead_year
    out = {
        **stats(center),
        "benchmark": stats(q),
        "winning_settings": sum(r["sharpe"] > sharpe(q) for r in rows),
        "settings": len(rows),
        "halves": halves(center),
        "benchmark_halves": halves(q),
        "half_split": str(days[len(days) // 2].date()),
        "cost_0_1pct": stats(stressed),
        "excess_sharpe": sharpe(center - cash),
        "benchmark_excess_sharpe": sharpe(q - cash),
        "largest_lead_year": lead_year,
        "largest_lead": float(lead.max()),
        "without_largest_lead_year": [sharpe(center[keep]), sharpe(q[keep])],
        "yearly": {str(y): [float(a), float(b)] for y, a, b in zip(lead.index, center_years, q_years)},
        "grid": rows,
    }
    checks = {
        "settings_70pct": out["winning_settings"] >= 0.7 * len(rows),
        "both_halves": all(a > b for a, b in zip(out["halves"], out["benchmark_halves"])),
        "cost_0_1pct": out["cost_0_1pct"]["sharpe"] > out["benchmark"]["sharpe"],
        "excess_over_bil": out["excess_sharpe"] > out["benchmark_excess_sharpe"],
        "without_largest_lead_year": out["without_largest_lead_year"][0] > out["without_largest_lead_year"][1],
    }
    if design in QQQ_TIMING:
        e = exposure(center)
        out["mean_exposure"] = float(e.mean())
        out["cash_in_bil_sharpe"] = sharpe(center + (1 - e) * cash)
        out["fixed_exposure_sharpe"] = sharpe(e.mean() * q + (1 - e.mean()) * cash)
        checks["beats_fixed_exposure"] = out["cash_in_bil_sharpe"] > out["fixed_exposure_sharpe"]
    out["checks"] = checks
    out["passed"] = all(checks.values())
    return out


def run_design(design, reports):
    use("us_fund" if design in QQQ_TIMING else "us")
    grid, original = GRIDS[design], ORIGINAL[design]
    rows = []
    for values in product(*grid.values()):
        setting = dict(zip(grid, values))
        position, resample = POSITIONS[design](setting)
        report = backtest(position, resample)
        r = daily(report)
        rows.append({**setting, **stats(r)})
        if design in QQQ_TIMING:
            rows[-1]["exposure"] = float(exposure(r).mean())
        if setting == original:
            center, center_report, center_position, center_resample = r, report, position, resample
    stressed = daily(backtest(center_position, center_resample, fee=COST_PER_SIDE))
    out = validate(design, rows, center, stressed)
    if design not in QQQ_TIMING:
        held = center_position.sum(axis=1)
        out["holdings"] = {"mean": float(held.mean()), "min": int(held.min()), "max": int(held.max())}
        weights = center_position.div(held.replace(0, np.nan), axis=0).fillna(0)
        out["one_way_turnover_per_rebalance"] = float((weights.diff().abs().sum(axis=1) / 2).iloc[1:].mean())
        out["trades"] = int(len(center_report.get_trades()))
    if reports:
        center_report.to_html(str(reports / f"{design}_report.html"))
    print(design, "settings", out["winning_settings"], "/", out["settings"], "sharpe", round(out["sharpe"], 3),
          "qqq", round(out["benchmark"]["sharpe"], 3), "passed", out["passed"], flush=True)
    return out, center


def reference(mask):
    """Equal weight across a whole pool, rebalanced quarterly: what a filter is compared with."""
    use("us")
    r = daily(backtest(mask, rebalance_days()))
    return {**stats(r), "halves": halves(r), "excess_sharpe": sharpe(r - bil()),
            "mean_holdings": float(mask.sum(axis=1).mean())}, r


def zone_shares(pool):
    z = z_score().where(pool)
    n = pool.sum(axis=1)
    share = lambda zone: float((zone.sum(axis=1) / n).mean())
    return {"distress": share(z.lt(DISTRESS)), "grey": share(z.ge(DISTRESS) & z.le(SAFE)), "safe": share(z.gt(SAFE))}


def write(study, results, curves, extra):
    out = OUT[study]
    out.mkdir(parents=True, exist_ok=True)
    payload = {"study": study, "start": START, "end": END, "run_date": str(date.today()),
               "finlab_version": finlab.__version__, "results": results, **extra}
    (out / "metrics.json").write_text(json.dumps(payload, indent=2, allow_nan=False))
    curves = {"QQQ": benchmark(), "BIL": bil(), **curves}
    pd.DataFrame({k: (1 + v).cumprod() for k, v in curves.items()}).round(6).to_csv(out / "data.csv", index_label="date")


def main(reports):
    results, curves = {}, {}
    for design in GRIDS:
        results[design], curves[design] = run_design(design, reports)

    dip_extra = {"events": dip_events(ORIGINAL["dip"])}
    write("dip", {k: results[k] for k in ("dip", "scale_in")}, {k: curves[k] for k in ("dip", "scale_in")}, dip_extra)

    use("us")
    pool = buyback_pool()
    pool_stats, pool_curve = reference(pool)
    original = ORIGINAL["buyback"]
    picks = buyback_position(original)[0]
    y = buyback_yield(original["quarters"], original["measure"])
    buying = pool & y.gt(0)
    _, held = np.nonzero(picks.to_numpy())
    buyback_extra = {"pool_equal_weight": pool_stats,
                     "median_yield_held": float(np.nanmedian(y.where(picks).to_numpy())),
                     "median_yield_pool": float(np.nanmedian(y.where(buying).to_numpy())),
                     "pool_share_buying_back": float(buying.sum(axis=1).div(pool.sum(axis=1)).mean()),
                     "sectors_held": common_stocks()["sector"].iloc[held].fillna("Unknown")
                     .value_counts(normalize=True).round(4).to_dict()}
    write("buyback", {"buyback": results["buyback"]}, {"buyback": curves["buyback"], "pool": pool_curve}, buyback_extra)

    z_pool_500 = z_pool(500)
    z_stats, z_curve = reference(z_pool_500)
    distress_stats, distress_curve = reference(z_pool_500 & z_score().lt(DISTRESS))
    altman_extra = {"pool_equal_weight": z_stats, "distress_zone_only": distress_stats,
                    "zone_shares_pool_500": zone_shares(z_pool_500)}
    write("altman", {k: results[k] for k in ("exclude_distress", "safe_zone")},
          {"exclude_distress": curves["exclude_distress"], "safe_zone": curves["safe_zone"],
           "pool": z_curve, "distress_zone": distress_curve}, altman_extra)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--reports", type=Path)
    args = parser.parse_args()
    if args.reports:
        args.reports.mkdir(parents=True, exist_ok=True)
    finlab.login()
    main(args.reports)