"""Day-of-the-week effect on QQQ, SPY and Taiwan's 0050 (FinLab backtest).

Weeks are calendar weeks (Monday to Sunday). F is the first trading day of a week (Monday,
or Tuesday after a Monday holiday) and E the last (usually Friday). Windows count trading
days, so o = -1 moves a window one trading day earlier, across the weekend if needed.

  D1  Skip Monday: cash from F-L+1+o to F+o, QQQ otherwise. L=1, o=0 is "buy at Monday's
      close, sell at Friday's close": the weekend and Monday are in cash.
  D2  Hold only Friday: QQQ from E-L+1+o to E+o, cash otherwise.
  D3  D1 on 0050, judged against 0050 with Taiwan costs.
  D4  D2 on 0050, judged against 0050 with Taiwan costs.
  Grid for all four: L 1/2/3 x o -1/0/1, original setting L=1, o=0.

"Beats" means a higher daily Sharpe ratio (mean daily return / standard deviation x sqrt(252),
no risk-free rate) than buying and holding the benchmark, dividends included, over the same
dates: QQQ for D1 and D2, 0050 (etl:adj_close) for D3 and D4. A design passes only if
  1. at least 7 of its 9 settings beat the benchmark,
  2. its original setting wins in both halves of the backtest (split at the middle date),
  3. its original setting wins at finlab's default cost: zero for US ETFs, plus 0.1% per
     side as a second US check; 0.1425% commission each way and 0.3% sell tax in Taiwan
     (0050's actual 0.1% ETF tax and zero cost are shown for reference only).
Designs in the market less than half the time must also beat the benchmark held at the
same average weight, and the US ones must beat QQQ on returns above T-bills (BIL). The same
US rules on SPY are compared with SPY for reference only.

Statistics (not pass criteria): average daily return by weekday with 95% confidence
intervals for QQQ, SPY (2016 onward), 0050 (2007 onward) and the TAIEX (1999 onward),
weekday averages by year, and buying 1 dollar (0050: 1 TWD) at the close of the same
weekday every week until the end of the data.

Timing: the calendar is known in advance. To earn day x (close of x-1 to close of x) the
buy happens at the close of x-1; sim(resample=None, trade_at_price="close") trades a signal
at the next close, so the signal is set on x-2.

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

    python research.py      # writes metrics.json and data.csv next to this file
"""
import json
import math
from pathlib import Path

import numpy as np
import pandas as pd
from pandas.tseries.holiday import (
    AbstractHolidayCalendar,
    GoodFriday,
    Holiday,
    USLaborDay,
    USMartinLutherKingJr,
    USMemorialDay,
    USPresidentsDay,
    USThanksgivingDay,
    nearest_workday,
    sunday_to_monday,
)

import finlab
from finlab import data
from finlab.backtest import sim

DATA_END = {"us": "2026-11-05", "tw": "2026-11-06"}   # last closes used in the article; None for the latest data
COST_PER_SIDE = 0.001       # second US cost check: 0.1% per buy and per sell
OUT = Path(__file__).resolve().parent
L_GRID, O_GRID, BASE = (1, 2, 3), (-1, 0, 1), (1, 0)
WEEKDAYS = ["Monday", "Tuesday", "Wednesday", "Thursday", "Friday"]
TW_COSTS = {"default": {}, "etf_tax": {"tax_ratio": 0.001}, "zero": {"fee_ratio": 0, "tax_ratio": 0}}


class NYSEHolidays(AbstractHolidayCalendar):
    """Regular NYSE holidays, used only for trading days after the last data row."""

    rules = [
        Holiday("NewYearsDay", month=1, day=1, observance=sunday_to_monday),
        USMartinLutherKingJr,
        USPresidentsDay,
        GoodFriday,
        USMemorialDay,
        Holiday("Juneteenth", month=6, day=19, start_date="2022-01-01", observance=nearest_workday),
        Holiday("IndependenceDay", month=7, day=4, observance=nearest_workday),
        USLaborDay,
        USThanksgivingDay,
        Holiday("Christmas", month=12, day=25, observance=nearest_workday),
    ]


def cut(table, market):
    return table.loc[:DATA_END[market]] if DATA_END else table


# ---------------------------------------------------------------------------
# 1. Prices and trading calendars
# ---------------------------------------------------------------------------
finlab.login()
data.set_market("tw")
TW = pd.DataFrame({
    "0050": cut(pd.DataFrame(data.get("etl:adj_close")), "tw")["0050"].dropna(),
})
TAIEX = cut(pd.DataFrame(data.get("taiex_total_index:收盤指數")), "tw")["TAIEX"].dropna()
data.set_market("us_fund")
US = cut(pd.DataFrame(data.get("us_fund_price:adj_close"))[["QQQ", "SPY", "BIL"]], "us").dropna()


def with_future(days, us):
    """Trading days plus the next few weeks, so the last rows know where the week ends."""
    future = pd.bdate_range(days[-1] + pd.Timedelta(days=1), days[-1] + pd.Timedelta(days=30))
    if us:
        future = future.difference(NYSEHolidays().holidays(future[0], future[-1]))
    return days.append(future)   # Taiwan holidays after the data are unknown; weekdays stand in


MARKETS = {
    "us": {"days": US.index, "calendar": with_future(US.index, True), "prices": US, "bench": "QQQ"},
    "tw": {"days": TW.index, "calendar": with_future(TW.index, False), "prices": TW, "bench": "0050"},
}


def week_edges(calendar):
    """Positions of the first and last trading day of every calendar week."""
    week = pd.Series(np.arange(len(calendar)), index=calendar).groupby(calendar.to_period("W-SUN"))
    return week.min().to_numpy(), week.max().to_numpy()


# ---------------------------------------------------------------------------
# 2. Designs and backtests
# ---------------------------------------------------------------------------
def window(calendar, anchors, L, o):
    """True on the L trading days ending at anchor+o."""
    mask = np.zeros(len(calendar), dtype=bool)
    for p in anchors:
        hi = p + o
        if 0 <= hi < len(calendar):
            mask[max(hi - L + 1, 0):hi + 1] = True
    return mask


def skip_first_day(market, L, o):
    calendar = MARKETS[market]["calendar"]
    return ~window(calendar, week_edges(calendar)[0], L, o)


def last_day_only(market, L, o):
    calendar = MARKETS[market]["calendar"]
    return window(calendar, week_edges(calendar)[1], L, o)


DESIGNS = {
    "D1": {"name": "Skip Monday, hold QQQ", "market": "us", "ticker": "QQQ", "want": skip_first_day},
    "D2": {"name": "Hold QQQ only on Friday", "market": "us", "ticker": "QQQ", "want": last_day_only},
    "D3": {"name": "Skip Monday, hold 0050", "market": "tw", "ticker": "0050", "want": skip_first_day},
    "D4": {"name": "Hold 0050 only on Friday", "market": "tw", "ticker": "0050", "want": last_day_only},
}


def position(market, want, ticker):
    """want[x] = hold over day x; the signal is set two trading days earlier."""
    days = MARKETS[market]["days"]
    return pd.DataFrame({ticker: want[2:len(days) + 2]}, index=days).astype(float)


def backtest(market, pos, **cost):
    data.set_market("us_fund" if market == "us" else "tw")
    report = sim(pos, resample=None, trade_at_price="close", upload=False, **cost)
    days = MARKETS[market]["days"]
    curve = report.creturn.reindex(days).ffill().fillna(1.0).loc[days[1]:]   # first buy at day 1's close
    return curve / curve.iloc[0], report


def returns(curve):
    return curve.pct_change().dropna()


def sharpe(r):
    return float(r.mean() / r.std() * math.sqrt(252))


def stats(curve):
    curve = curve / curve.iloc[0]
    r = returns(curve)
    years = (curve.index[-1] - curve.index[0]).days / 365.25
    return {
        "start": str(curve.index[0].date()),
        "end": str(curve.index[-1].date()),
        "cagr": round(float(curve.iloc[-1] ** (1 / years) - 1), 6),
        "sharpe": round(sharpe(r), 6),
        "mdd": round(float((curve / curve.cummax() - 1).min()), 6),
        "vol": round(float(r.std() * math.sqrt(252)), 6),
        "total_return": round(float(curve.iloc[-1] - 1), 6),
    }


def hold(market, ticker, index):
    return MARKETS[market]["prices"][ticker].reindex(index)


def versus(market, curve, ticker=None):
    ticker = ticker or MARKETS[market]["bench"]
    s, b = stats(curve), stats(hold(market, ticker, curve.index))
    return {"strategy": s, "bench": b, "win": bool(s["sharpe"] > b["sharpe"])}


def halves(market, curve):
    middle = curve.index[0] + (curve.index[-1] - curve.index[0]) / 2
    return [versus(market, curve.loc[:middle]), versus(market, curve.loc[middle:])]


def above_bills(curve):
    bills = returns(hold("us", "BIL", curve.index))
    s = sharpe(returns(curve) - bills)
    q = sharpe(returns(hold("us", "QQQ", curve.index)) - bills)
    return {"strategy": round(s, 6), "qqq": round(q, 6), "win": s > q}


def fixed_weight(market, curve, weight):
    """The benchmark at the design's average weight, the rest in cash, rebalanced daily."""
    r = returns(hold(market, MARKETS[market]["bench"], curve.index))
    fixed = pd.concat([pd.Series([1.0], index=curve.index[:1]), (1 + weight * r).cumprod()])
    s, f = stats(curve), stats(fixed)
    return {"weight": round(weight, 6), "fixed": f, "win": bool(s["sharpe"] > f["sharpe"])}


def best_year_removed(market, curve):
    rs, rb = returns(curve), returns(hold(market, MARKETS[market]["bench"], curve.index))
    excess = (rs + 1).groupby(rs.index.year).prod() - (rb + 1).groupby(rb.index.year).prod()
    best = int(excess.idxmax())
    s, b = sharpe(rs[rs.index.year != best]), sharpe(rb[rb.index.year != best])
    return {"best_year": best, "excess_that_year": round(float(excess.max()), 6),
            "strategy_sharpe": round(s, 6), "bench_sharpe": round(b, 6), "win": s > b}


def trades_per_year(report, curve):
    years = (curve.index[-1] - curve.index[0]).days / 365.25
    return round(len(report.get_trades()) / years, 2)


def validate(key):
    spec = DESIGNS[key]
    market, ticker = spec["market"], spec["ticker"]
    grid = []
    for L in L_GRID:
        for o in O_GRID:
            row = versus(market, backtest(market, position(market, spec["want"](market, L, o), ticker))[0])
            grid.append({"L": L, "o": o, "sharpe": row["strategy"]["sharpe"], "cagr": row["strategy"]["cagr"],
                         "mdd": row["strategy"]["mdd"], "bench_sharpe": row["bench"]["sharpe"], "win": row["win"]})
    pos = position(market, spec["want"](market, *BASE), ticker)
    weight = float(pos[ticker].mean())
    curve, report = backtest(market, pos)
    out = {
        "name": spec["name"],
        "base": {"L": BASE[0], "o": BASE[1]},
        "grid": grid,
        "grid_wins": sum(r["win"] for r in grid),
        "full": versus(market, curve),
        "halves": halves(market, curve),
        "in_market": round(weight, 6),
        "trades_per_year": trades_per_year(report, curve),
        "best_year_removed": best_year_removed(market, curve),
    }
    checks = [out["grid_wins"] >= 7, all(h["win"] for h in out["halves"]), out["full"]["win"]]
    if market == "us":
        out["cost_0_1pct"] = versus(market, backtest(market, pos, fee_ratio=COST_PER_SIDE)[0])
        spy_pos = pos.rename(columns={ticker: "SPY"})
        out["spy_reference"] = versus(market, backtest(market, spy_pos)[0], "SPY")
        checks.append(out["cost_0_1pct"]["win"])
    else:
        out["costs"] = {
            name: out["full"] if name == "default" else versus(market, backtest(market, pos, **cost)[0])
            for name, cost in TW_COSTS.items()
        }
    if weight < 0.5:
        out["fixed_weight"] = fixed_weight(market, curve, weight)
        checks.append(out["fixed_weight"]["win"])
        if market == "us":
            out["above_bills"] = above_bills(curve)
            checks.append(out["above_bills"]["win"])
    out["passes"] = bool(all(checks))
    print(f"{key} {spec['name']}: {out['grid_wins']}/9 beat {MARKETS[market]['bench']}, original Sharpe "
          f"{out['full']['strategy']['sharpe']:.2f} vs {out['full']['bench']['sharpe']:.2f}, "
          f"{'passes' if out['passes'] else 'does not pass'}")
    return out, curve


# ---------------------------------------------------------------------------
# 3. Weekday statistics (not pass criteria)
# ---------------------------------------------------------------------------
def mean_ci(values):
    v = pd.Series(values, dtype=float).dropna()
    se = v.std() / math.sqrt(len(v))
    return {"n": len(v), "mean": round(float(v.mean()), 6), "ci_low": round(float(v.mean() - 1.96 * se), 6),
            "ci_high": round(float(v.mean() + 1.96 * se), 6), "up_ratio": round(float((v > 0).mean()), 4)}


def weekday_stats(close):
    r = close.pct_change().dropna()
    saturdays = int((r.index.day_of_week == 5).sum())   # Taiwan's make-up Saturday sessions
    r = r[r.index.day_of_week < 5]
    by_day = {WEEKDAYS[d]: mean_ci(r[r.index.day_of_week == d]) for d in range(5)}
    monday, rest = r[r.index.day_of_week == 0], r[r.index.day_of_week != 0]
    t = (monday.mean() - rest.mean()) / math.sqrt(monday.var() / len(monday) + rest.var() / len(rest))
    yearly = r.groupby([r.index.year, r.index.day_of_week]).mean().unstack()
    yearly.columns = [WEEKDAYS[d] for d in yearly.columns]
    return {
        "start": str(r.index[0].date()), "end": str(r.index[-1].date()),
        "saturday_sessions_excluded": saturdays, "all_days": mean_ci(r), "weekdays": by_day, "monday_vs_rest_t": round(float(t), 3),
        "worst_day_count_by_year": yearly.idxmin(axis=1).value_counts().to_dict(),
        "best_day_count_by_year": yearly.idxmax(axis=1).value_counts().to_dict(),
        "yearly": {int(y): {d: round(float(v), 6) for d, v in row.items()} for y, row in yearly.iterrows()},
    }


def weekly_buying(close):
    """Buy 1 unit of money at the close of the same weekday every week; value at the last close."""
    out = {}
    for d in range(5):
        buys = close[close.index.day_of_week == d]
        shares = (1 / buys).sum()
        out[WEEKDAYS[d]] = {"buys": len(buys), "value_per_dollar": round(float(shares * close.iloc[-1] / len(buys)), 6)}
    return out


def main():
    results, curves = {}, {}
    for key in DESIGNS:
        results[key], curves[key] = validate(key)
    series = {"QQQ": US["QQQ"], "SPY": US["SPY"], "0050": TW["0050"], "TAIEX": TAIEX}
    metrics = {
        "data_end": {"us": str(US.index[-1].date()), "tw": str(TW.index[-1].date())},
        "cost_per_side": COST_PER_SIDE,
        "benchmarks": {"QQQ": stats(hold("us", "QQQ", US.index[1:])), "SPY": stats(hold("us", "SPY", US.index[1:])),
                       "0050": stats(hold("tw", "0050", TW.index[1:]))},
        "designs": results,
        "weekday_stats": {name: weekday_stats(close) for name, close in series.items()},
        "weekly_buying": {name: weekly_buying(series[name]) for name in ("QQQ", "SPY", "0050")},
    }
    (OUT / "metrics.json").write_text(json.dumps(metrics, indent=1))
    us = pd.DataFrame({"D1_skip_monday_QQQ": curves["D1"], "D2_friday_only_QQQ": curves["D2"]})
    us["QQQ_buy_hold"] = hold("us", "QQQ", us.index) / hold("us", "QQQ", us.index).iloc[0]
    tw = pd.DataFrame({"D3_skip_monday_0050": curves["D3"], "D4_friday_only_0050": curves["D4"]})
    tw["0050_buy_hold"] = hold("tw", "0050", tw.index) / hold("tw", "0050", tw.index).iloc[0]
    us.join(tw, how="outer").round(6).to_csv(OUT / "data.csv", index_label="date")


if __name__ == "__main__":
    main()
