"""
The Turtle trading rules on US multi-asset ETFs and on Taiwan large caps (FinLab research, 2026-11).

The rules follow Curtis Faith, "Way of the Turtle" (2007), long side only:
  N        the 20-day average true range (Wilder's smoothing).
  entry    the close is above the highest high of the previous E days.
  exit     the close is below the lowest low of the previous X days.
  stop     the close is more than 2N below the price of the last unit bought.
  units    one unit is 1% of the account divided by N; another unit is added each time the close
           rises 0.5N above the last unit's price (at most one a day), up to 4 per market,
           6 per group of related markets (US only) and 12 in total.
  System 1 skips a breakout when the previous one, taken or not, would have made money as a
           one-unit trade with the same stop and exit. A skipped market still enters on a
           55-day breakout.

A unit's weight in the account is 1% x close / N. Each time a market is bought, its weight is reset
to units x 1% x close / N and stays there until its next trade. When the weights add up to more
than 100%, finlab's `sim` scales them all down to 100%, so there is no leverage; the rest is cash.
Signals are read at a close and traded at the next day's open, on adjusted prices.

Designs (each with a 3 x 3 grid of E x X around the original rules)
  us_system1   18 US ETFs, System 1         E 15/20/25 x X 8/10/12 (original 20/10)
  us_system2   18 US ETFs, System 2         E 45/55/65 x X 15/20/25 (original 55/20)
  tw_system1   0050 and the 30 largest TWSE-listed common stocks by market value at each month
               end, System 1 (a stock enters only while it is in the 30; exits follow the rules)
  tw_system2   the same, System 2

"Beats" means a higher daily Sharpe ratio (no risk-free rate) than buying and holding QQQ (US) or
0050 (Taiwan), dividends included, over exactly the same dates. A design passes only if
  1. at least 70% of its settings beat the benchmark,
  2. its original setting beats it in both halves of the backtest,
  3. its original setting beats it at finlab's default cost (US: zero, also run at 0.1% per side;
     Taiwan: 0.1425% fee and 0.3% tax),
  4. its original setting still beats it without the calendar year of its largest lead,
  5. its original setting beats the same exposure held fixed (US: the 18 ETFs at their average
     weights; Taiwan: 0050 and the 30 stocks equally weighted, scaled to the design's average
     invested share), rebalanced monthly,
  6. US only: its original setting beats QQQ on returns above T-bills (BIL).

Run:  python research.py      writes metrics.json and data.csv next to this file
"""
import json
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 = {"us": "2016-09-01", "tw": "2013-05-02"}
END = {"us": "2026-11-05", "tw": "2026-11-06"}
BENCHMARK = {"us": "QQQ", "tw": "0050"}
COST_PER_SIDE = 0.001
RISK = 0.01          # account share risked per N
ATR_DAYS = 20
ADD_STEP = 0.5       # in N
STOP = 2.0           # in N
MAX_UNITS, GROUP_UNITS, TOTAL_UNITS = 4, 6, 12
FAILSAFE = 55
TW_TOP = 30
OUT = Path(__file__).resolve().parent

ETFS = {
    "stocks": ["SPY", "QQQ", "IWM", "EFA", "EEM"],
    "bonds": ["TLT", "IEF"],
    "metals": ["GLD", "SLV", "CPER"],
    "energy": ["USO", "UGA"],
    "agriculture": ["DBA"],
    "currencies": ["FXE", "FXB", "FXF", "FXY", "FXC"],
}
TICKERS = [t for group in ETFS.values() for t in group]

DESIGNS = {  # name: (market, system, grid, original)
    "us_system1": ("us", 1, ((15, 20, 25), (8, 10, 12)), (20, 10)),
    "us_system2": ("us", 2, ((45, 55, 65), (15, 20, 25)), (55, 20)),
    "tw_system1": ("tw", 1, ((15, 20, 25), (8, 10, 12)), (20, 10)),
    "tw_system2": ("tw", 2, ((45, 55, 65), (15, 20, 25)), (55, 20)),
}


@cache
def get(key, market):
    data.set_market(market)
    data.truncate_end = END["tw" if market == "tw" else "us"]
    return data.get(key)


@cache
def trading_days(market):
    key = "us_price:adj_close" if market == "us" else "etl:adj_close"
    return get(key, "us" if market == "us" else "tw").loc[:END[market]].index


@cache
def tw_members():
    """Trading day x stock: True while a TWSE common stock is among the 30 largest by market value
    at the previous month end. 0050 is always in."""
    value = get("etl:market_value", "tw")
    listing = get("security_categories", "tw")
    listed = listing.loc[(listing["market"] == "sii") & listing["stock_id"].str.fullmatch(r"[1-9]\d{3}")
                         & ~listing["category"].astype(str).str.contains("存託"), "stock_id"]
    value = value[value.columns.intersection(listed)]
    month_ends = value.groupby(value.index.to_period("M")).tail(1)
    top = month_ends.rank(axis=1, ascending=False, method="first") <= TW_TOP
    top = top.loc[:, top.any()]
    m = top.reindex(trading_days("tw"), method="ffill").shift(1).fillna(False).astype(bool)
    m.insert(0, "0050", True)
    return m


@cache
def ohlc(market):
    """Adjusted open, high, low and close on the market's trading days."""
    if market == "us":
        factor = get("us_fund_price:adj_close", "us_fund")[TICKERS] / get("us_fund_price:close", "us_fund")[TICKERS]
        frames = {k: get(f"us_fund_price:{k}", "us_fund")[TICKERS] * factor for k in ("open", "high", "low")}
        frames["close"] = get("us_fund_price:adj_close", "us_fund")[TICKERS]
        columns = TICKERS
    else:
        frames = {k: get(f"etl:adj_{k}", "tw") for k in ("open", "high", "low", "close")}
        columns = list(tw_members().columns)
    return {k: v.reindex(index=trading_days(market), columns=columns) for k, v in frames.items()}


@cache
def atr(market):
    o = ohlc(market)
    previous = o["close"].shift(1)
    true_range = np.maximum(o["high"], previous) - np.minimum(o["low"], previous)
    return true_range.ewm(alpha=1 / ATR_DAYS, adjust=False, min_periods=ATR_DAYS).mean()


def allowed(market):
    if market == "us":
        return np.ones((len(trading_days("us")), len(TICKERS)), dtype=bool)
    return tw_members().to_numpy()


def groups(market):
    if market == "tw":
        return None
    return np.array([next(i for i, g in enumerate(ETFS.values()) if t in g) for t in TICKERS])


def weights(design, entry, exit_days):
    """Day x market target weights, dated on the signal day (finlab trades them the next day)."""
    return turtle(design, entry, exit_days)[0]


@cache
def turtle(design, entry, exit_days):
    """Target weights and units held, day x market, from the rules in the docstring."""
    market, system = DESIGNS[design][:2]
    o = ohlc(market)
    O, C, N = o["open"].to_numpy(), o["close"].to_numpy(), atr(market).to_numpy()
    high = o["high"].shift(1).rolling(entry).max().to_numpy()
    low = o["low"].shift(1).rolling(exit_days).min().to_numpy()
    safe = o["high"].shift(1).rolling(FAILSAFE).max().to_numpy()
    member, group = allowed(market), groups(market)
    days, count = C.shape
    last_price = o["close"].notna().to_numpy()[::-1].argmax(axis=0)
    last_price = days - 1 - last_price          # last row with a close, per market

    units = np.zeros(count, dtype=int)
    paid = np.full(count, np.nan)               # price of the last unit bought
    paid_n = np.full(count, np.nan)             # N on that unit's signal day
    weight = np.zeros(count)
    won = np.zeros(count, dtype=bool)           # System 1: the previous breakout made money
    ghost = np.full(count, np.nan)              # System 1: the hypothetical one-unit trade
    ghost_stop = np.full(count, np.nan)
    table = np.zeros((days, count))
    held = np.zeros((days, count), dtype=int)
    breakouts = skipped = 0                     # flat members' breakouts after the start, and skips
    first = trading_days(market).get_indexer([pd.Timestamp(START[market])], method="bfill")[0]

    def room(a):
        if units.sum() >= TOTAL_UNITS:
            return False
        return group is None or units[group == group[a]].sum() < GROUP_UNITS

    for t in range(days - 1):
        c, n, fill = C[t], N[t], O[t + 1]
        with np.errstate(invalid="ignore"):
            tradable = ~np.isnan(c) & ~np.isnan(fill) & (n > 0)
            breakout = tradable & (c > high[t])
            leave = tradable & (c < low[t])
            if system == 1:
                closed = tradable & ~np.isnan(ghost) & (leave | (c < ghost_stop))
            gone = (units > 0) & (t + 1 > last_price)
            sell = gone | ((units > 0) & tradable & (leave | (c < paid - STOP * paid_n)))
            units[sell], weight[sell] = 0, 0.0
            add = (units > 0) & (units < MAX_UNITS) & tradable & (c >= paid + ADD_STEP * paid_n)
            signal = breakout & ~won | tradable & (c > safe[t]) if system == 1 else breakout
            if system == 1:
                # The hypothetical trade fills at the next open, so its result counts from tomorrow
                won[closed] = fill[closed] > ghost[closed]
                ghost[closed] = np.nan
        if t >= first:
            flat = (units == 0) & member[t] & breakout
            breakouts += int(flat.sum())
            skipped += int((flat & ~signal).sum())
        for a in [*np.flatnonzero(add), *np.flatnonzero((units == 0) & signal & member[t])]:
            if room(a):
                units[a] += 1
                paid[a], paid_n[a] = fill[a], n[a]
                weight[a] = units[a] * RISK * c[a] / n[a]
        if system == 1:
            start = breakout & np.isnan(ghost)
            ghost[start] = fill[start]
            ghost_stop[start] = fill[start] - STOP * n[start]
        table[t], held[t] = weight, units
    frames = (pd.DataFrame(x, index=trading_days(market), columns=o["close"].columns).loc[START[market]:]
              for x in (table, held))
    return (*frames, {"breakouts": breakouts, "skipped": skipped})


def run_sim(table, market, fee_ratio=None, resample=None):
    data.set_market("us_fund" if market == "us" else "tw")
    data.truncate_end = END[market]
    table = table.loc[:, (table != 0).any()]
    report = sim(table, resample=resample, trade_at_price="open", fee_ratio=fee_ratio, upload=False)
    curve = report.creturn
    curve = curve[curve.index.isin(trading_days(market))].loc[:END[market]]
    moved = (curve != curve.iloc[0]).to_numpy().nonzero()[0]
    return curve.iloc[max(moved[0] - 1, 0):], report


def run(design, setting, fee_ratio=None):
    return run_sim(weights(design, *setting), DESIGNS[design][0], fee_ratio)


@cache
def backtest(design, setting, fee_ratio=None):
    return run(design, setting, fee_ratio)[0]


def monthly(index):
    """Rebalance dates: the first day, then the last trading day of each month."""
    ends = index.to_series().groupby(index.to_period("M")).max()
    return pd.DatetimeIndex([index[0], *ends[ends > index[0]]])


def scaled(design):
    """The original setting's weights after sim scales a row above 100% down to 100%."""
    w = weights(design, *DESIGNS[design][3])
    return w.div(w.sum(axis=1).clip(lower=1), axis=0)


@cache
def fixed_curve(design):
    """The same exposure held fixed, rebalanced monthly (see the docstring)."""
    market = DESIGNS[design][0]
    w = scaled(design)
    if market == "us":
        table = pd.DataFrame(w.mean().to_dict(), index=w.index, columns=w.columns)
    else:
        # Each row holds the list that takes effect the next day, when the rebalance trades
        members = tw_members()
        upcoming = members.shift(-1, fill_value=False)
        upcoming.iloc[-1] = members.iloc[-1]
        pool = upcoming.loc[w.index].astype(float)
        table = pool.div(pool.sum(axis=1), axis=0) * w.sum(axis=1).mean()
    return run_sim(table, market, resample=monthly(table.index))[0]


# ---------------------------------------------------------------------------
# Statistics
# ---------------------------------------------------------------------------
@cache
def funds(market):
    if market == "us":
        return get("us_fund_price:adj_close", "us_fund")[["QQQ", "SPY", "BIL"]].reindex(trading_days("us"))
    return get("etl:adj_close", "tw")[["0050"]].reindex(trading_days("tw"))


def daily_sharpe(returns):
    return returns.mean() / returns.std() * 252 ** 0.5


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


def excess_sharpe(curve):
    returns = curve.pct_change().dropna()
    bills = funds("us")["BIL"].reindex(curve.index).pct_change(fill_method=None).reindex(returns.index)
    return daily_sharpe(returns - bills)


def stats(curve, market):
    out = {"start": str(curve.index[0].date()), "end": str(curve.index[-1].date())}
    others = ("QQQ", "SPY") if market == "us" else ("0050",)
    for label, series in (("", curve), *((f"{t.lower()}_", funds(market)[t].reindex(curve.index)) for t in others)):
        returns = series.pct_change()
        out[f"{label}sharpe"] = round(daily_sharpe(returns.dropna()), 6)
        out[f"{label}cagr"] = round(cagr(series), 6)
        out[f"{label}mdd"] = round((series / series.cummax() - 1).min(), 6)
        out[f"{label}volatility"] = round(returns.std() * 252 ** 0.5, 6)
    out["beats"] = bool(out["sharpe"] > out[f"{BENCHMARK[market].lower()}_sharpe"])
    return out


def yearly(curve):
    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 halves(curve, market):
    middle = curve.index[len(curve) // 2]
    return [stats(part, market) for part in (curve.loc[:middle], curve.loc[middle:])]


def without_best_year(curve, market):
    fund = funds(market)[BENCHMARK[market]].reindex(curve.index)
    ours, theirs = yearly(curve), yearly(fund)
    best = max(ours, key=lambda year: ours[year] - theirs[year])
    returns, theirs_daily = curve.pct_change().dropna(), fund.pct_change().dropna()
    return {"year": best, "lead": round(ours[best] - theirs[best], 6),
            "sharpe": round(daily_sharpe(returns[returns.index.year != int(best)]), 6),
            "benchmark_sharpe": round(daily_sharpe(theirs_daily[theirs_daily.index.year != int(best)]), 6)}


def trades(design):
    market = DESIGNS[design][0]
    table = run(design, DESIGNS[design][3])[1].get_trades().dropna(subset=["exit_date"])
    wins, losses = table.loc[table["return"] > 0, "return"], table.loc[table["return"] <= 0, "return"]
    out = {"count": len(table), "win_rate": round(len(wins) / len(table), 6),
           "mean_win": round(wins.mean(), 6), "mean_loss": round(losses.mean(), 6),
           "best": round(table["return"].max(), 6), "median_days": float(table["period"].median())}
    if market == "us":
        label = {t: g for g, members in ETFS.items() for t in members}
        table = table.assign(group=table["stock_id"].map(label))
        out["by_group"] = {g: {"count": int(len(part)), "win_rate": round((part["return"] > 0).mean(), 6),
                               "mean_return": round(part["return"].mean(), 6)}
                           for g, part in table.groupby("group")}
    return out


def study(design):
    market, _, axes, original = DESIGNS[design]
    grid = {str(s): stats(backtest(design, s), market) for s in product(*axes)}
    curve = backtest(design, original)
    wins = sum(cell["beats"] for cell in grid.values())
    first, second = halves(curve, market)
    fixed = fixed_curve(design).reindex(curve.index).dropna()
    w = scaled(design)
    raw, units, breakouts = turtle(design, *original)
    result = {
        "name": design, "original": str(original), "original_stats": grid[str(original)], "grid": grid,
        "grid_wins": wins, "grid_size": len(grid), "first_half": first, "second_half": second,
        "yearly": yearly(curve), "without_best_year": without_best_year(curve, market),
        "fixed": {"sharpe": round(daily_sharpe(fixed.pct_change().dropna()), 6), "cagr": round(cagr(fixed), 6),
                  "mdd": round((fixed / fixed.cummax() - 1).min(), 6)},
        "invested": round(float(w.sum(axis=1).mean()), 6),
        "days_scaled_down": round(float((raw.sum(axis=1) > 1).mean()), 6),
        "raw_weight_median": round(float(raw.sum(axis=1).median()), 4),
        "rebalance_days": int((raw.diff().abs().sum(axis=1) > 0).sum()),
        "turnover": round(float(w.diff().abs().sum(axis=1).sum() / 2 / ((w.index[-1] - w.index[0]).days / 365.25)), 4),
        "days_flat": round(float((raw.sum(axis=1) == 0).mean()), 6),
        "markets_held": round(float((raw > 0).sum(axis=1).mean()), 2),
        "units_held": round(float(units.sum(axis=1).mean()), 2),
        "breakouts": breakouts,
        "trades": trades(design),
    }
    checks = {"grid": wins >= 0.7 * len(grid), "halves": first["beats"] and second["beats"],
              "cost": result["original_stats"]["beats"],
              "without_best_year": result["without_best_year"]["sharpe"] > result["without_best_year"]["benchmark_sharpe"],
              "fixed": result["original_stats"]["sharpe"] > result["fixed"]["sharpe"]}
    if market == "us":
        result["with_cost"] = stats(backtest(design, original, COST_PER_SIDE), market)
        result["average_weights"] = {t: round(v, 4) for t, v in w.mean().items()}
        result["excess_sharpe"] = round(excess_sharpe(curve), 6)
        result["qqq_excess_sharpe"] = round(excess_sharpe(funds("us")["QQQ"].reindex(curve.index)), 6)
        checks["cost"] = checks["cost"] and result["with_cost"]["beats"]
        checks["above_tbills"] = result["excess_sharpe"] > result["qqq_excess_sharpe"]
    else:
        held = (raw > 0).sum()
        result["most_traded"] = {t: int(v) for t, v in held.sort_values(ascending=False).head(10).items()}
    result["checks"] = {k: bool(v) for k, v in checks.items()}
    result["passes"] = all(checks.values())
    return result


def main():
    results = {design: study(design) for design in DESIGNS}
    metrics = {"start": START, "end": END, "run_date": str(pd.Timestamp.now(tz="Asia/Taipei").date()),
               "finlab_version": finlab.__version__, "designs": results,
               "yearly_funds": {m: {t: yearly(funds(m)[t].loc[START[m]:END[m]]) for t in funds(m) if t != "BIL"}
                                for m in ("us", "tw")}}
    (OUT / "metrics.json").write_text(json.dumps(metrics, indent=1))
    columns = {}
    for design, (market, _, _, original) in DESIGNS.items():
        columns[design] = backtest(design, original)
        columns[f"{design}_fixed"] = fixed_curve(design)
    for market in ("us", "tw"):
        first = min(columns[d].index[0] for d in DESIGNS if DESIGNS[d][0] == market)
        for ticker in funds(market):
            if ticker != "BIL":
                columns[ticker] = funds(market)[ticker].loc[first:END[market]]
    table = pd.DataFrame(columns).sort_index()
    table = table.ffill()
    (table / table.bfill().iloc[0]).round(6).rename_axis("date").to_csv(OUT / "data.csv")
    for design, result in results.items():
        print(design, result["original_stats"]["sharpe"], result["grid_wins"], "/", result["grid_size"],
              result["checks"], result["passes"])


if __name__ == "__main__":
    main()
