"""Leveraged ETF decay: how big it is, where it comes from, and whether a volatility switch helps (FinLab).

Part 1, decay. For each leveraged ETF below, over every calendar year and every quarter:

  gap         ln(1 + fund return) - k * ln(1 + underlying return)
  volatility  ln(1 + ideal return) - k * ln(1 + underlying return), where "ideal" is exactly
              k times the underlying every day, reset daily, with no costs. It is close to
              -k(k-1)/2 times the sum of squared daily underlying returns.
  costs       gap - volatility: expense ratio, financing, swap spreads and tracking error.

k is the daily leverage: 3 for TQQQ, 2 for QLD, -3 for SQQQ. Underlying returns include
dividends (dividend-adjusted closes).

Part 2, the volatility switch. At every close, measure the underlying's realised volatility
over the last N trading days (standard deviation of daily returns x sqrt(252)). Above the
threshold, hold the plain 1x ETF; at or below it, hold the leveraged ETF. A change in the
signal trades at the next day's close.

"Beats" means a higher daily Sharpe ratio (no risk-free rate) than QQQ with dividends
(0050 with dividends for the Taiwan pair) over exactly the same dates. A design passes
only if all five hold:
  1. at least 7 of its 9 neighbouring parameter sets beat the benchmark,
  2. it beats the benchmark in both halves of the backtest,
  3. it beats it after trading costs (US: 0.1% per buy and per sell on top of finlab's
     zero-commission default; Taiwan: finlab's default fee and tax are in every run),
  4. it beats a fixed mix that holds the same average weight in the leveraged ETF,
     rebalanced every month, so the timing has to add something,
  5. it still beats the benchmark once its best calendar year is removed.

Complete the AI-assisted setup flow at https://finlab.finance/en/setup before running this file.

    python research.py               # decay summary and the five checks for every design
    python research.py decay         # decay summary only
    python research.py tqqq_switch   # backtest one design and open its report
    python research.py --export      # also write data.csv and metrics.json next to this file
"""
import json
import math
import re
import sys
from pathlib import Path

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

DATA_END = "2026-11-05"
US_START = "2016-04-01"   # US fund prices start 2016-01-04; every 40-day window is ready by April
TW_START = "2015-01-05"   # same start as FinLab's 00631L article
COST_PER_SIDE = 0.001     # check 3 for the US designs
YEAR = 252

# leveraged ETF: (underlying ETF, daily leverage)
US_FUNDS = {
    "TQQQ": ("QQQ", 3), "QLD": ("QQQ", 2), "SQQQ": ("QQQ", -3),
    "UPRO": ("SPY", 3), "SSO": ("SPY", 2),
    "SOXL": ("SOXX", 3), "TECL": ("XLK", 3), "TNA": ("IWM", 3), "UDOW": ("DIA", 3),
    "CURE": ("XLV", 3), "LABU": ("XBI", 3), "YINN": ("FXI", 3), "EDC": ("EEM", 3),
    "TMF": ("TLT", 3), "TYD": ("IEF", 3),
}
TW_FUNDS = {"00631L": ("0050", 2)}
FUNDS = US_FUNDS | TW_FUNDS

SWITCHES = {
    "tqqq_switch": {"fund": "TQQQ", "base": "QQQ", "params": (20, 0.25),
                    "grid": ((10, 20, 40), (0.20, 0.25, 0.30))},
    "qld_switch": {"fund": "QLD", "base": "QQQ", "params": (20, 0.25),
                   "grid": ((10, 20, 40), (0.20, 0.25, 0.30))},
    "upro_switch": {"fund": "UPRO", "base": "SPY", "params": (20, 0.20),
                    "grid": ((10, 20, 40), (0.15, 0.20, 0.25))},
    "soxl_switch": {"fund": "SOXL", "base": "SOXX", "params": (20, 0.35),
                    "grid": ((10, 20, 40), (0.30, 0.35, 0.40))},
    "00631l_switch": {"fund": "00631L", "base": "0050", "params": (20, 0.20),
                      "grid": ((10, 20, 40), (0.15, 0.20, 0.25)), "market": "tw"},
}

_closes = {}


def closes(market):
    """Dividend-adjusted closes for every ETF in this file."""
    if market not in _closes:
        if market == "us":
            data.set_market("us_fund")
            frame = data.get("us_fund_price:adj_close")
            tickers = {t for fund, (base, _) in US_FUNDS.items() for t in (fund, base)} | {"BIL"}
        else:
            data.set_market("tw")
            frame = data.get("etl:adj_close")
            tickers = {t for fund, (base, _) in TW_FUNDS.items() for t in (fund, base)}
        frame = pd.DataFrame(frame[sorted(tickers)])
        frame.index = pd.to_datetime(frame.index)
        _closes[market] = frame.loc[:DATA_END]
    return _closes[market]


def price(ticker, market="us"):
    # us_fund_price keeps market holidays as blank rows: drop them before any rolling window
    return closes(market)[ticker].dropna()


def market_of(fund):
    return "tw" if fund in TW_FUNDS else "us"


# ---------------------------------------------------------------------------
# Part 1: decay
# ---------------------------------------------------------------------------
def split(fund_return, base_return, daily_base, k):
    """Split one period's gap into the volatility term and the cost term."""
    ideal = float((1 + k * daily_base).prod() - 1)
    fund_log = math.log1p(fund_return)
    ideal_log = math.log1p(ideal)
    k_log = k * math.log1p(base_return)
    return {
        "fund": fund_return,
        "underlying": base_return,
        "k_x_underlying": k * base_return,
        "ideal": ideal,
        "simple_gap": fund_return - k * base_return,
        "gap": fund_log - k_log,
        "volatility": ideal_log - k_log,
        "volatility_approx": -k * (k - 1) / 2 * float((daily_base ** 2).sum()),
        "costs": fund_log - ideal_log,
        "underlying_vol": float(daily_base.std() * YEAR ** 0.5),
    }


def periods(fund, freq):
    """One row per calendar year ("Y") or quarter ("Q"); the first and last can be partial."""
    base, k = FUNDS[fund]
    market = market_of(fund)
    f, u = price(fund, market), price(base, market)
    days = f.index.intersection(u.index)
    if freq == "Y":
        labels = days.year.astype(str)
    else:
        labels = days.year.astype(str) + "Q" + days.quarter.astype(str)
    ends = pd.Series(days, index=days).groupby(labels).max()
    starts = [days[0], *ends.iloc[:-1]]
    bil = price("BIL") if market == "us" else None
    rows = []
    for label, start, end in zip(ends.index, starts, ends):
        # underlying returns on its own trading days, so a halt in the fund still counts every day
        daily_base = u.loc[start:end].pct_change().dropna()
        row = {"period": label, "start": str(start.date()), "end": str(end.date()),
               "days": len(daily_base),
               **split(f[end] / f[start] - 1, u[end] / u[start] - 1, daily_base, k)}
        if bil is not None:
            row["bil"] = float(bil.asof(end) / bil.asof(start) - 1)
        rows.append(row)
    return rows


def calibration(fund):
    """Daily returns of the fund against k times the underlying on the days both trade."""
    base, k = FUNDS[fund]
    market = market_of(fund)
    both = pd.concat([price(fund, market), price(base, market)], axis=1, join="inner")
    r = both.pct_change().dropna()
    rf, ru = r.iloc[:, 0], r.iloc[:, 1]
    slope = float(np.polyfit(ru, rf, 1)[0])
    return {"correlation": round(float(rf.corr(k * ru)), 4), "slope": round(slope, 3),
            "days_off_by_2pct": int(((rf - k * ru).abs() > 0.02).sum()), "days": len(r)}


def whole_period(fund):
    base, k = FUNDS[fund]
    market = market_of(fund)
    f, u = price(fund, market), price(base, market)
    days = f.index.intersection(u.index)
    start, end = days[0], days[-1]
    daily_base = u.loc[start:end].pct_change().dropna()
    years = len(daily_base) / YEAR
    row = split(f[end] / f[start] - 1, u[end] / u[start] - 1, daily_base, k)
    per_year = {f"{key}_per_year": row[key] / years for key in ("gap", "volatility", "volatility_approx", "costs")}
    return {"start": str(start.date()), "end": str(end.date()), "years": round(years, 3),
            "leverage": k, "underlying_ticker": base, **row, **per_year,
            "fund_cagr": (1 + row["fund"]) ** (1 / years) - 1,
            "underlying_cagr": (1 + row["underlying"]) ** (1 / years) - 1}


def cost_vs_bil(year_rows, k):
    """Regress the yearly cost term on the T-bill return; theory says the slope is about 1 - k."""
    full = [r for r in year_rows if r["period"] not in ("2016", DATA_END[:4])]
    x = np.log1p([r["bil"] for r in full])
    y = np.array([r["costs"] for r in full])
    slope, intercept = np.polyfit(x, y, 1)
    return {"years": [r["period"] for r in full], "slope": round(float(slope), 3),
            "intercept": round(float(intercept), 5), "theory_slope": 1 - k}


def table(rows):
    """Rows as one compact list each, under a single header."""
    columns = list(rows[0])
    return {"columns": columns,
            "rows": [[round(v, 6) if isinstance(v, float) else v for v in row.values()] for row in rows]}


def decay(fund):
    years = periods(fund, "Y")
    out = {"calibration": calibration(fund), "whole_period": whole_period(fund),
           "years": table(years), "quarters": table(periods(fund, "Q"))}
    if market_of(fund) == "us":
        out["cost_vs_bil"] = cost_vs_bil(years, FUNDS[fund][1])
    return out


def print_decay(fund, result):
    w, c = result["whole_period"], result["calibration"]
    print(f"{fund:<7}{w['leverage']:>3}x {w['underlying_ticker']:<5} corr {c['correlation']:.4f}"
          f" | per year: gap {w['gap_per_year']:+.1%} = volatility {w['volatility_per_year']:+.1%}"
          f" + costs {w['costs_per_year']:+.1%}"
          f" | CAGR {w['fund_cagr']:.1%} vs underlying {w['underlying_cagr']:.1%}")


# ---------------------------------------------------------------------------
# Part 2: the volatility switch
# ---------------------------------------------------------------------------
def switch_position(name, window=None, threshold=None):
    spec = SWITCHES[name]
    if window is None:
        window, threshold = spec["params"]
    market = spec.get("market", "us")
    base = price(spec["base"], market)
    volatility = base.pct_change().rolling(window).std() * YEAR ** 0.5
    calm = (volatility <= threshold).astype(float)
    calm = calm.loc[TW_START if market == "tw" else US_START:]
    return pd.DataFrame({spec["fund"]: calm, spec["base"]: 1 - calm})


def run(position, market, fee_ratio=None, resample=None, tax_ratio=None):
    if market == "us":
        data.set_market("us_fund")
        return sim(position, resample=resample, trade_at_price="close",
                   fee_ratio=fee_ratio, upload=False)
    data.set_market("tw")
    # finlab's default fee (0.1425% per side) and tax (0.3% on sells) unless tax_ratio is given
    return sim(position, resample=resample, tax_ratio=tax_ratio, upload=False)


def equity_curve(report):
    """From the close before the first trade fills to the last day of price data used here."""
    curve = report.creturn.loc[:DATA_END]
    moved = (curve != curve.iloc[0]).to_numpy().nonzero()[0]
    return curve.iloc[max(moved[0] - 1, 0):]


def sharpe(returns):
    return float(returns.mean() / returns.std() * YEAR ** 0.5)


def benchmark_for(name, curve):
    spec = SWITCHES[name]
    market = spec.get("market", "us")
    ticker = "0050" if market == "tw" else "QQQ"
    return price(ticker, market).reindex(curve.index).ffill()


def stats(curve):
    curve = curve / curve.iloc[0]
    years = (curve.index[-1] - curve.index[0]).days / 365.25
    returns = curve.pct_change().dropna()
    return {"cagr": round(float(curve.iloc[-1] ** (1 / years) - 1), 6),
            "sharpe": round(sharpe(returns), 6),
            "vol": round(float(returns.std() * YEAR ** 0.5), 6),
            "mdd": round(float((curve / curve.cummax() - 1).min()), 6)}


def versus(curve, bench):
    s, b = sharpe(curve.pct_change().dropna()), sharpe(bench.pct_change().dropna())
    return {"start": str(curve.index[0].date()), "end": str(curve.index[-1].date()),
            "sharpe": round(s, 4), "bench_sharpe": round(b, 4), "beats": bool(s > b)}


def drop_best_year(curve, bench):
    rs, rb = curve.pct_change().dropna(), bench.pct_change().dropna()
    yearly = (1 + rs).groupby(rs.index.year).prod() - 1
    best = int(yearly.idxmax())
    keep = rs.index.year != best
    s, b = sharpe(rs[keep]), sharpe(rb[keep])
    return {"best_year": best, "best_year_return": round(float(yearly.max()), 4),
            "sharpe": round(s, 4), "bench_sharpe": round(b, 4), "beats": bool(s > b)}


def excess_sharpe(curve):
    bil = price("BIL").reindex(curve.index).ffill().pct_change()
    return round(sharpe((curve.pct_change() - bil).dropna()), 4)


def yearly_returns(curve):
    """Calendar-year returns; the first year runs from the curve's first day."""
    ends = curve.groupby(curve.index.year).last()
    starts = ends.shift(1)
    starts.iloc[0] = curve.iloc[0]
    return {str(year): round(float(r), 4) for year, r in (ends / starts - 1).items()}


def volatile_days(name="tqqq_switch"):
    """Where the decay falls: the switch's days split by the underlying's volatility at the prior close."""
    spec = SWITCHES[name]
    window, threshold = spec["params"]
    k = US_FUNDS[spec["fund"]][1]
    both = pd.concat([price(spec["base"]), price(spec["fund"])], axis=1, join="inner")
    returns = both.pct_change()
    vol = returns.iloc[:, 0].rolling(window).std() * YEAR ** 0.5
    volatile = vol.shift(1) > threshold
    days = returns.index[(returns.index > pd.Timestamp(US_START)) & (returns.index <= pd.Timestamp(DATA_END))]
    base, fund, volatile = returns.iloc[:, 0].loc[days], returns.iloc[:, 1].loc[days], volatile.loc[days]
    term = np.log1p(k * base) - k * np.log1p(base)   # each day's volatility term
    out = {"design": name, "start": str(days[0].date()), "end": str(days[-1].date()),
           "rule": f"{spec['base']} {window}-day volatility above {threshold:.0%} at the prior close"}
    for label, mask in (("calm", ~volatile), ("volatile", volatile)):
        years = mask.sum() / YEAR
        out[label] = {
            "days": int(mask.sum()), "share_of_days": round(float(mask.mean()), 4),
            "share_of_volatility_term": round(float(term[mask].sum() / term.sum()), 4),
            "base_log_per_year": round(float(np.log1p(base[mask]).sum() / years), 4),
            "fund_log_per_year": round(float(np.log1p(fund[mask]).sum() / years), 4),
            "base_sharpe": round(sharpe(base[mask]), 4), "fund_sharpe": round(sharpe(fund[mask]), 4),
        }
    out["volatile_days_by_year"] = {str(y): int(n) for y, n in volatile.groupby(volatile.index.year).sum().items()}
    return out


def validate(name):
    spec = SWITCHES[name]
    market = spec.get("market", "us")
    fund = spec["fund"]

    # 1. Neighbouring parameters: the original value and one step either side.
    grid = []
    for window in spec["grid"][0]:
        for threshold in spec["grid"][1]:
            curve = equity_curve(run(switch_position(name, window, threshold), market))
            grid.append({"params": [window, threshold], **versus(curve, benchmark_for(name, curve))})
    wins = sum(row["beats"] for row in grid)

    position = switch_position(name)
    report = run(position, market)
    curve = equity_curve(report)
    bench = benchmark_for(name, curve)
    full = versus(curve, bench)

    # 2. First half and second half of the backtest, each on its own.
    middle = curve.index[0] + (curve.index[-1] - curve.index[0]) / 2
    halves = [versus(curve.loc[:middle], bench.loc[:middle]),
              versus(curve.loc[middle:], bench.loc[middle:])]

    # 3. Trading costs: the Taiwan runs already pay finlab's default fee and tax.
    if market == "us":
        cost_curve = equity_curve(run(position, market, fee_ratio=COST_PER_SIDE))
        cost = versus(cost_curve, benchmark_for(name, cost_curve))
    else:
        cost = full.copy()

    # 4. The same average weight in the leveraged ETF, held fixed and rebalanced monthly.
    held = position[fund].loc[curve.index[0]:]
    weight = float(held.mean())
    fixed_position = pd.DataFrame({fund: weight, spec["base"]: 1 - weight}, index=position.index)
    fixed_curve = equity_curve(run(fixed_position, market, resample="M")).loc[curve.index[0]:]
    fixed = {"weight": round(weight, 4), "sharpe": round(sharpe(fixed_curve.pct_change().dropna()), 4)}
    fixed["beaten"] = bool(full["sharpe"] > fixed["sharpe"])

    # 5. Remove the best calendar year from both series.
    best_year = drop_best_year(curve, bench)

    passed = {"neighbours": wins >= 7, "halves": all(h["beats"] for h in halves),
              "costs": cost["beats"], "fixed_mix": fixed["beaten"], "best_year": best_year["beats"]}
    years = (curve.index[-1] - curve.index[0]).days / 365.25
    switches = int((held.diff().abs() > 0).sum())
    fund_curve = price(fund, market).reindex(curve.index).ffill()
    result = {
        "design": name, "fund": fund, "base": spec["base"], "params": list(spec["params"]),
        "full": full, "stats": stats(curve), "bench_stats": stats(bench), "fund_stats": stats(fund_curve),
        "yearly": {"strategy": yearly_returns(curve), "fund": yearly_returns(fund_curve),
                   "bench": yearly_returns(bench)},
        "time_in_fund": round(weight, 4), "switches_per_year": round(switches / years, 2),
        "grid": grid, "grid_wins": wins, "halves": halves, "cost": cost,
        "fixed_mix": fixed, "drop_best_year": best_year,
        "passed": {**passed, "all": all(passed.values())},
    }
    if market == "us":
        result["excess_bil_sharpe"] = {"strategy": excess_sharpe(curve), "bench": excess_sharpe(bench)}
    else:
        # reference only: Taiwan ETFs actually pay 0.1% sell tax, not finlab's 0.3% stock default
        etf_tax = equity_curve(run(position, market, tax_ratio=0.001))
        result["etf_tax_reference"] = versus(etf_tax, benchmark_for(name, etf_tax))
    return result, report


def print_validation(result):
    full, h1, h2 = result["full"], *result["halves"]
    verdict = "PASS" if result["passed"]["all"] else "fail"
    print(f"{result['design']:<14} {verdict:<4} Sharpe {full['sharpe']:.2f} vs {full['bench_sharpe']:.2f}"
          f" | neighbours {result['grid_wins']}/9"
          f" | halves {h1['sharpe']:.2f}/{h2['sharpe']:.2f} vs {h1['bench_sharpe']:.2f}/{h2['bench_sharpe']:.2f}"
          f" | costs {result['cost']['sharpe']:.2f}"
          f" | fixed {result['fixed_mix']['weight']:.0%} mix {result['fixed_mix']['sharpe']:.2f}"
          f" | without {result['drop_best_year']['best_year']} {result['drop_best_year']['sharpe']:.2f}"
          f" vs {result['drop_best_year']['bench_sharpe']:.2f}")


# ---------------------------------------------------------------------------
# Output
# ---------------------------------------------------------------------------
def export(decays, validations, reports, volatile):
    here = Path(__file__).resolve().parent
    columns = {}
    for ticker in ("QQQ", "TQQQ", "QLD", "SPY", "UPRO", "SOXX", "SOXL"):
        columns[ticker] = price(ticker)
    qqq = price("QQQ")
    columns["QQQ cubed, no volatility loss"] = (qqq / qqq.iloc[0]) ** 3
    columns["3x QQQ daily, no costs"] = (1 + 3 * qqq.pct_change().fillna(0)).cumprod()
    listed = price("00631L", "tw").index[0]
    for ticker in ("0050", "00631L"):
        columns[ticker] = price(ticker, "tw").loc[listed:]
    for name, report in reports.items():
        columns[name] = equity_curve(report)
    frame = pd.DataFrame({key: series / series.dropna().iloc[0] for key, series in columns.items()})
    frame.round(6).to_csv(here / "data.csv", index_label="date")
    metrics = {"finlab_version": finlab.__version__, "pandas_version": pd.__version__,
               "data_end": DATA_END, "us_start": US_START, "tw_start": TW_START,
               "win_rule": "daily Sharpe (no risk-free rate) above QQQ (0050 for Taiwan) "
                           "dividend-adjusted buy and hold on the same dates",
               "decay": decays, "volatile_days": volatile, "switch": validations}
    text = json.dumps(metrics, indent=1, default=float)
    # keep each innermost list (a table row, a parameter pair) on one line
    text = re.sub(r"\[\s+([^\[\]{}]*?)\s+\]", lambda m: "[" + " ".join(m.group(1).split()) + "]", text)
    (here / "metrics.json").write_text(text)


def main(args):
    names = [a for a in args if a in SWITCHES]
    if names:
        for name in names:
            report = run(switch_position(name), SWITCHES[name].get("market", "us"))
            print(name, report.get_stats())
            report.display()
        return

    decays = {fund: decay(fund) for fund in FUNDS}
    for fund, result in decays.items():
        print_decay(fund, result)
    if "decay" in args:
        return

    volatile = volatile_days()
    print("volatile days", {key: volatile[key] for key in ("calm", "volatile")})

    validations, reports = {}, {}
    for name in SWITCHES:
        validations[name], reports[name] = validate(name)
        print_validation(validations[name])
    if "--export" in args:
        export(decays, validations, reports, volatile)


if __name__ == "__main__":
    finlab.login()  # finlab guides you through login automatically
    main(sys.argv[1:])