"""
Dollar-cost averaging vs lump sum: SPY, QQQ and Taiwan's 0050 (FinLab research, 2026-11).

The same sum of money, starting on the first trading day of every month:
  lump sum   buy the fund with everything at that day's close
  DCA        buy it in equal parts over N months; money not yet invested sits in BIL
             (T-bill ETF) for SPY / QQQ, or in Taiwan dollar cash at 0% for 0050.
             Every purchase uses (remaining balance) / (remaining purchases).
Both are compared on the first trading day of month N after the start: by then both are
fully invested, so the gap between them stays the same percentage afterwards.
Gap = DCA wealth / lump-sum wealth - 1.  Start months overlap.

Each fund gets a 3 x 3 grid around the original setting (12 months, monthly purchases):
  N months   6 / 12 / 18
  purchases  every 5 trading days / first trading day of each month / of each quarter

"DCA is better" only if all hold:
  1. in at least 7 of the 9 settings, DCA ends with more money in more than half the starts,
  2. at the original setting, in more than half the starts of both the first and second half,
  3. at the original setting with 0.1% per side (BIL legs included; 0050: the 0.1425%
     Taiwan default fee), still more than half,
  4. at the original setting, DCA's daily Sharpe above BIL (0050: cash at 0%) over the
     comparison window is higher than lump sum's in more than half the starts,
  5. at the original setting, DCA beats a fixed mix holding its average weight in the fund,
     (N + 1) / (2N), rebalanced on the first trading day of each month,
  6. at the original setting, still more than half after removing the start year in which
     DCA led by the most on average.
"Lump sum is better" if 1-3 hold the other way round.

Run:  python research.py      writes metrics.json and data.csv next to this file
"""
import json
from functools import cache
from pathlib import Path

import numpy as np
import pandas as pd
from finlab import data

END = "2026-11-05"
OUT = Path(__file__).resolve().parent
MONTHS = (6, 12, 18)
FREQS = ("5d", "monthly", "quarterly")
ORIGINAL = (12, "monthly")
US_STRESS_FEE = 0.001          # per side, every trade including BIL
TW_FEE = 0.001425              # Taiwan broker fee, finlab sim default
TW_ETF_TAX = 0.001             # Taiwan transaction tax on selling an ETF
TW_CASH_ALT = 0.01             # annual cash interest, sensitivity only
US_FUNDS = ("SPY", "QQQ")
COMMON_START = "2016-01-01"    # first US fund price in FinLab data
TW_START = "2007-05-01"        # 0050 prices start 2007-04-23, so the first full month is May

data.truncate_end = END


# ---------------------------------------------------------------------------
# Prices: one frame per fund with columns fund / cash, on trading days only
# ---------------------------------------------------------------------------
@cache
def us_prices():
    data.set_market("us")
    prices = data.get("us_fund_price:adj_close")[[*US_FUNDS, "BIL"]]
    return prices.loc[prices["QQQ"].notna()].ffill()


@cache
def tw_prices():
    data.set_market("tw")
    return data.get("etl:adj_close")["0050"].dropna().loc[TW_START:]


def cash_growth(days, annual_rate):
    years = (days - days[0]).days / 365
    return pd.Series((1 + annual_rate) ** years, index=days)


def prices(fund, tw_cash=0.0, start=None):
    if fund == "0050":
        price = tw_prices()
        frame = pd.DataFrame({"fund": price, "cash": cash_growth(price.index, tw_cash)})
    else:
        frame = us_prices()[[fund, "BIL"]].set_axis(["fund", "cash"], axis=1)
    return frame.loc[start:] if start else frame


def costs(fund, stress=False):
    """(fund buy fee, cash buy fee, cash sell fee)"""
    if fund == "0050":
        return TW_FEE, 0.0, 0.0
    fee = US_STRESS_FEE if stress else 0.0
    return fee, fee, fee


# ---------------------------------------------------------------------------
# One start: lump-sum, DCA and fixed-mix wealth paths over the comparison window
# ---------------------------------------------------------------------------
def month_starts(days):
    return days.to_series().groupby(days.to_period("M")).min().to_numpy()


def purchases(window, starts_in_window, freq):
    """Window offsets of the purchase days (the first is 0)."""
    if freq == "5d":
        return np.arange(0, len(window) - 1, 5)
    step = 1 if freq == "monthly" else 3
    return np.searchsorted(window, starts_in_window[:-1:step])


def lump_sum(fund_px, fee):
    return (1 - fee) / fund_px[0] * fund_px


def dca(fund_px, cash_px, buys, fees):
    fund_fee, cash_buy_fee, cash_sell_fee = fees
    n = len(buys)
    fund_units = np.zeros(len(fund_px))
    cash_units = np.zeros(len(fund_px))
    units, cash = (1 / n) * (1 - fund_fee) / fund_px[0], (1 - 1 / n) * (1 - cash_buy_fee) / cash_px[0]
    previous = 0
    for i, day in enumerate(buys[1:], start=1):
        fund_units[previous:day], cash_units[previous:day] = units, cash
        sold = cash / (n - i)
        cash -= sold
        units += sold * cash_px[day] * (1 - cash_sell_fee) * (1 - fund_fee) / fund_px[day]
        previous = day
    fund_units[previous:], cash_units[previous:] = units, cash
    return fund_units * fund_px + cash_units * cash_px


def fixed_mix(fund_px, cash_px, rebalance_days, weight, fund, fees):
    """Fund at `weight` of the portfolio, the rest in cash, reset on every rebalance day."""
    fund_fee = fees[0]
    sell_fee = fund_fee + TW_ETF_TAX if fund == "0050" else fund_fee
    wealth = np.zeros(len(fund_px))
    value, units, cash = 1.0, 0.0, 0.0
    marks = list(rebalance_days) + [len(fund_px)]
    for begin, stop in zip(marks[:-1], marks[1:]):
        value = units * fund_px[begin] + cash * cash_px[begin] if begin else value
        trade = weight * value - units * fund_px[begin]
        value -= abs(trade) * (fund_fee if trade > 0 else sell_fee)
        units = weight * value / fund_px[begin]
        cash = (1 - weight) * value / cash_px[begin]
        wealth[begin:stop] = units * fund_px[begin:stop] + cash * cash_px[begin:stop]
    return wealth


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


def max_drawdown(path):
    return (path / np.maximum.accumulate(path) - 1).min()


def windows(frame, months, freq):
    """Yield (start day, window days, purchase offsets, month-start offsets) per start."""
    days = frame.index
    starts = month_starts(days)
    for k in range(len(starts) - months):
        window = days[(days >= starts[k]) & (days <= starts[k + months])]
        firsts = starts[k:k + months + 1]
        yield pd.Timestamp(starts[k]), window, purchases(window.to_numpy(), firsts, freq), np.searchsorted(window, firsts[:-1])


@cache
def run(fund, months, freq, stress=False, tw_cash=0.0, start=None):
    frame = prices(fund, tw_cash, start)
    fees = costs(fund, stress)
    rows = []
    for day, window, buys, firsts in windows(frame, months, freq):
        fund_px = frame.loc[window, "fund"].to_numpy()
        cash_px = frame.loc[window, "cash"].to_numpy()
        ls = lump_sum(fund_px, fees[0])
        averaged = dca(fund_px, cash_px, buys, fees)
        mix = fixed_mix(fund_px, cash_px, firsts, (months + 1) / (2 * months), fund, fees)
        cash_returns = np.diff(cash_px) / cash_px[:-1]
        ls_returns, dca_returns = np.diff(ls) / ls[:-1], np.diff(averaged) / averaged[:-1]
        rows.append({
            "start": day,
            "end": window[-1],
            "purchases": len(buys),
            "gap": averaged[-1] / ls[-1] - 1,
            "gap_vs_fixed_mix": averaged[-1] / mix[-1] - 1,
            "ls_sharpe": sharpe(ls_returns),
            "dca_sharpe": sharpe(dca_returns),
            "ls_excess_sharpe": sharpe(ls_returns - cash_returns),
            "dca_excess_sharpe": sharpe(dca_returns - cash_returns),
            "ls_mdd": max_drawdown(ls),
            "dca_mdd": max_drawdown(averaged),
            "ls_return": ls[-1] - 1,
            "dca_return": averaged[-1] - 1,
        })
    return pd.DataFrame(rows).set_index("start")


# ---------------------------------------------------------------------------
# Statistics and verdicts
# ---------------------------------------------------------------------------
def summary(table):
    gap = table["gap"]
    worst_start = table["ls_mdd"].idxmin()
    return {
        "starts": len(table),
        "first_start": str(table.index[0].date()),
        "last_start": str(table.index[-1].date()),
        "dca_win_rate": round(float((gap > 0).mean()), 4),
        "mean_gap": round(float(gap.mean()), 5),
        "median_gap": round(float(gap.median()), 5),
        "mean_gap_when_dca_wins": round(float(gap[gap > 0].mean()), 5) if (gap > 0).any() else None,
        "mean_gap_when_ls_wins": round(float(gap[gap < 0].mean()), 5) if (gap < 0).any() else None,
        "dca_best": {"start": str(gap.idxmax().date()), "gap": round(float(gap.max()), 5)},
        "dca_worst": {"start": str(gap.idxmin().date()), "gap": round(float(gap.min()), 5)},
        "dca_sharpe_win_rate": round(float((table["dca_sharpe"] > table["ls_sharpe"]).mean()), 4),
        "dca_excess_sharpe_win_rate":
            round(float((table["dca_excess_sharpe"] > table["ls_excess_sharpe"]).mean()), 4),
        "dca_beats_fixed_mix_rate": round(float((table["gap_vs_fixed_mix"] > 0).mean()), 4),
        "mean_ls_return": round(float(table["ls_return"].mean()), 5),
        "mean_dca_return": round(float(table["dca_return"].mean()), 5),
        "mean_ls_mdd": round(float(table["ls_mdd"].mean()), 5),
        "mean_dca_mdd": round(float(table["dca_mdd"].mean()), 5),
        "worst_ls_mdd": {"start": str(worst_start.date()),
                         "ls_mdd": round(float(table["ls_mdd"].min()), 5),
                         "dca_mdd_same_start": round(float(table.loc[worst_start, "dca_mdd"]), 5)},
    }


def halves(table):
    middle = len(table) // 2
    return table.iloc[:middle], table.iloc[middle:]


def without_best_year(table):
    yearly = table["gap"].groupby(table.index.year).mean()
    best = int(yearly.idxmax())
    return best, table.loc[table.index.year != best]


def win_rate(table, side="dca"):
    gap = table["gap"]
    return float((gap > 0).mean() if side == "dca" else (gap < 0).mean())


def evaluate(fund):
    grid = {f"{months}m_{freq}": run(fund, months, freq) for months in MONTHS for freq in FREQS}
    original = run(fund, *ORIGINAL)
    original_summary = summary(original)
    stressed = run(fund, *ORIGINAL, stress=True)
    first, second = halves(original)
    best_year, trimmed = without_best_year(original)

    def checks(side):
        return {
            "grid_wins": sum(win_rate(t, side) > 0.5 for t in grid.values()),
            "first_half": round(win_rate(first, side), 4),
            "second_half": round(win_rate(second, side), 4),
            "with_costs": round(win_rate(stressed, side), 4),
        }

    dca_checks = checks("dca") | {
        "excess_sharpe": original_summary["dca_excess_sharpe_win_rate"],
        "beats_fixed_mix": original_summary["dca_beats_fixed_mix_rate"],
        "best_year_removed": best_year,
        "without_best_year": round(win_rate(trimmed), 4),
    }
    ls_checks = checks("ls")
    dca_better = (dca_checks["grid_wins"] >= 7 and all(dca_checks[k] > 0.5 for k in (
        "first_half", "second_half", "with_costs", "excess_sharpe", "beats_fixed_mix", "without_best_year")))
    ls_better = ls_checks["grid_wins"] >= 7 and all(
        ls_checks[k] > 0.5 for k in ("first_half", "second_half", "with_costs"))
    return {
        "grid": {name: summary(table) for name, table in grid.items()},
        "original": original_summary,
        "first_half": summary(first),
        "second_half": summary(second),
        "with_costs": summary(stressed),
        "dca_checks": dca_checks,
        "ls_checks": ls_checks,
        "verdict": "dca_better" if dca_better else "lump_sum_better" if ls_better else "no_consistent_winner",
        "by_start_year": {int(y): round(float(g), 5) for y, g in
                          original["gap"].groupby(original.index.year).mean().items()},
    }


def example_paths(fund, start):
    """Daily lump-sum and DCA wealth for one start at the original setting (for charts)."""
    frame = prices(fund)
    for day, window, buys, _ in windows(frame, *ORIGINAL):
        if str(day.date()) == start:
            fund_px = frame.loc[window, "fund"].to_numpy()
            cash_px = frame.loc[window, "cash"].to_numpy()
            fees = costs(fund)
            return pd.DataFrame({"lump_sum": lump_sum(fund_px, fees[0]),
                                 "dca": dca(fund_px, cash_px, buys, fees)}, index=window)
    raise ValueError(start)


def main():
    out = {"end": END, "original": {"months": ORIGINAL[0], "purchases": ORIGINAL[1]}, "funds": {}}
    for fund in (*US_FUNDS, "0050"):
        out["funds"][fund] = evaluate(fund)
    out["0050_since_2016"] = {
        "original": summary(run("0050", *ORIGINAL, start=COMMON_START)),
        "grid": {f"{m}m_{f}": summary(run("0050", m, f, start=COMMON_START)) for m in MONTHS for f in FREQS},
    }
    out["0050_cash_1pct"] = summary(run("0050", *ORIGINAL, tw_cash=TW_CASH_ALT))

    examples = {}
    for fund, result in out["funds"].items():
        for label in ("dca_best", "dca_worst"):
            start = result["original"][label]["start"]
            path = example_paths(fund, start)
            examples[f"{fund}_{label}"] = {"start": start, "dates": [str(d.date()) for d in path.index],
                                           "lump_sum": path["lump_sum"].round(5).tolist(),
                                           "dca": path["dca"].round(5).tolist()}
    out["examples"] = examples

    columns = {}
    for fund in out["funds"]:
        for months in MONTHS:
            for freq in FREQS:
                columns[f"{fund}_{months}m_{freq}_gap"] = run(fund, months, freq)["gap"]
        table = run(fund, *ORIGINAL)
        columns[f"{fund}_12m_monthly_ls_mdd"] = table["ls_mdd"]
        columns[f"{fund}_12m_monthly_dca_mdd"] = table["dca_mdd"]
    frame = pd.DataFrame(columns).sort_index()
    frame.index = frame.index.strftime("%Y-%m-%d")
    frame.round(6).to_csv(OUT / "data.csv", index_label="start")
    (OUT / "metrics.json").write_text(json.dumps(out, indent=1, ensure_ascii=False))
    for fund, result in out["funds"].items():
        print(fund, result["verdict"], result["dca_checks"], result["ls_checks"])


if __name__ == "__main__":
    main()
