"""
Dogs of the Dow, 2017 to 2026 (FinLab research, 2026-11).

Designs, each with a 3 x 3 grid around the original setting
  dogs_of_the_dow   Each year, rank the 30 Dow stocks by dividend yield on the last trading day of
                    the month and hold the N highest, equal weight, for a year (O'Higgins 1991).
                    grid: N 8/10/12 x month November/December/January
  small_dogs        The N lowest-priced of the 10 highest yielders ("Small Dogs").
                    grid: N 4/5/6 x month November/December/January

Yield is the dividends with an ex-date in the 12 months to the signal, restated for later
splits, over that day's close. Dividend data start in January 2016, so December 2016 and
January 2017 are the first complete windows; the November column starts in 2017.
Signals are read at a close and traded at the next close.

"Beats QQQ" means a higher daily Sharpe ratio (no risk-free rate) than buying and holding QQQ,
dividends included, over exactly the same dates. A design passes only if
  1. at least 70% of its settings beat QQQ,
  2. its original setting beats QQQ in both halves of the backtest,
  3. its original setting beats QQQ at finlab's default US cost (zero) and at 0.1% per side,
  4. its original setting still beats QQQ without the calendar year of its largest lead.

Two checks were added after the results and are not designs:
  published_style   ranks on the latest dividend times the payments in the last 12 months, close
                    to the annualised yield in published Dogs lists
  rest_of_dow       holds the Dow stocks each year's list leaves out, equal weight

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 pandas as pd
from finlab import data
from finlab.backtest import sim

END = "2026-10-30"
COST_PER_SIDE = 0.001
FUNDS = ("QQQ", "SPY", "DIA")
OUT = Path(__file__).resolve().parent

# Effective dates, from the S&P Dow Jones Indices press releases:
# https://press.spglobal.com/2017-08-24-DowDuPont-Set-to-Join-Dow-Jones-Industrial-Average
# https://press.spglobal.com/2018-06-19-Walgreens-Boots-Alliance-Set-to-Join-Dow-Jones-Industrial-Average
# https://press.spglobal.com/2019-03-26-Dow-Set-to-Join-Dow-Jones-Industrial-Average
# https://press.spglobal.com/2020-08-24-Salesforce-com-Amgen-and-Honeywell-International-Set-to-Join-Dow-Jones-Industrial-Average
# https://press.spglobal.com/2024-02-20-Amazon-com-Set-to-Join-Dow-Jones-Industrial-Average-Uber-to-Join-Dow-Jones-Transportation-Average
# https://press.spglobal.com/2024-11-01-NVIDIA-and-Sherwin-Williams-Set-to-Join-Dow-Jones-Industrial-Average-Vistra-to-Join-Dow-Jones-Utility-Average
# https://press.spglobal.com/2026-06-23-Alphabet-Set-to-Join-and-Honeywell-International-to-Remain-in-Dow-Jones-Industrial-Average
# United Technologies became Raytheon Technologies on 2020-04-03; FinLab keeps both under RTX.
# DD before 2017-09-01 is not E. I. du Pont's traded price, so DD joins at the DowDuPont merger.
DOW_2016 = ("MMM AXP AAPL BA CAT CVX CSCO KO XOM GE GS HD IBM INTC JNJ JPM MCD MRK MSFT NKE "
            "PFE PG TRV UNH RTX VZ V WMT DIS").split()
DOW_CHANGES = [
    ("2017-09-01", ["DD"], []),
    ("2018-06-26", ["WBA"], ["GE"]),
    ("2019-04-02", ["DOW"], ["DD"]),
    ("2020-08-31", ["AMGN", "HON", "CRM"], ["XOM", "PFE", "RTX"]),
    ("2024-02-26", ["AMZN"], ["WBA"]),
    ("2024-11-08", ["NVDA", "SHW"], ["INTC", "DOW"]),
    ("2026-06-29", ["GOOGL"], ["VZ"]),
]
MONTHS = (11, 12, 1)
FIRST_YEAR = {11: 2017, 12: 2016, 1: 2017}

DESIGNS = {  # name: (grid, original setting)
    "dogs_of_the_dow": (((8, 10, 12), MONTHS), (10, 12)),
    "small_dogs": (((4, 5, 6), MONTHS), (5, 12)),
}

data.truncate_end = END
data.set_market("us")


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


@cache
def prices():
    return get("us_price:adj_close")


@cache
def funds():
    """Fund prices with dividends on stock-market trading days (us_fund_price has holiday rows)."""
    return get("us_fund_price:adj_close")[[*FUNDS, "BIL"]].reindex(prices().index)


# ---------------------------------------------------------------------------
# Dogs of the Dow
# ---------------------------------------------------------------------------
def dow_members(day):
    names = set(DOW_2016)
    for date, joined, left in DOW_CHANGES:
        if pd.Timestamp(date) <= day:
            names = (names | set(joined)) - set(left)
    return sorted(names)


@cache
def actions():
    rows = pd.DataFrame(get("us_corporate_actions")).astype({"stock_id": str})
    return rows[rows["stock_id"].isin({*DOW_2016, *sum((c[1] for c in DOW_CHANGES), [])})]


def dow_yields(day, published_style=False):
    """Yield of each Dow stock on `day`, with dividends restated for later splits."""
    rows = actions()
    rows = rows[(rows["date"] > day - pd.DateOffset(years=1)) & (rows["date"] <= day)]
    paid = pd.Series(0.0, index=dow_members(day))
    for ticker, stock in rows.groupby("stock_id"):
        if ticker in paid.index:
            splits = stock.dropna(subset=["split_factor"])
            cash = stock.dropna(subset=["dividend"]).sort_values("date")
            restated = [amount / splits.loc[splits["date"] > when, "split_factor"].prod()
                        for when, amount in zip(cash["date"], cash["dividend"])]
            if restated:
                paid[ticker] = restated[-1] * len(restated) if published_style else sum(restated)
    return paid / get("us_price:close").loc[day, paid.index]


def rebalance_days(month):
    days = prices().index
    ends = pd.DatetimeIndex(pd.Series(days, index=days).groupby([days.year, days.month]).max().values)
    return ends[(ends.month == month) & (ends.year >= FIRST_YEAR[month]) & (ends < END)]


@cache
def picks(name, n, month, published_style=False):
    chosen = {}
    close = get("us_price:close")
    for day in rebalance_days(month):
        yields = dow_yields(day, published_style)
        if name == "small_dogs":
            chosen[day] = close.loc[day, yields.nlargest(10).index].nsmallest(n).index
        elif name == "rest_of_dow":
            chosen[day] = yields.drop(yields.nlargest(n).index).index
        else:
            chosen[day] = yields.nlargest(n).index
    return chosen


def position(chosen):
    table = pd.DataFrame(False, index=list(chosen), columns=sorted({t for names in chosen.values() for t in names}))
    for day, names in chosen.items():
        table.loc[day, list(names)] = True
    days = prices().index
    return table.reindex(days[days >= table.index[0]], method="ffill").astype(bool)


def run(chosen, fee_ratio=None):
    report = sim(position(chosen), resample=None, trade_at_price="close", fee_ratio=fee_ratio,
                 market="us", upload=False)
    curve = report.creturn
    moved = (curve != curve.iloc[0]).to_numpy().nonzero()[0]
    return curve.iloc[max(moved[0] - 1, 0):]


@cache
def backtest(name, setting, fee_ratio=None):
    return run(picks(name, *setting), fee_ratio)


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


def bench(ticker, index):
    return funds()[ticker].reindex(index)


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


def stats(curve):
    out = {"start": str(curve.index[0].date()), "end": str(curve.index[-1].date())}
    for label, series in (("", curve), *((f"{t.lower()}_", bench(t, curve.index)) for t in FUNDS)):
        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_qqq"] = bool(out["sharpe"] > out["qqq_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):
    middle = curve.index[len(curve) // 2]
    return [stats(part) for part in (curve.loc[:middle], curve.loc[middle:])]


def without_best_year(curve):
    ours, theirs = yearly(curve), yearly(bench("QQQ", curve.index))
    best = max(ours, key=lambda year: ours[year] - theirs[year])
    returns = curve.pct_change().dropna()
    keep = returns.index.year != int(best)
    qqq = bench("QQQ", curve.index).pct_change().dropna()
    return {"year": best, "lead": round(ours[best] - theirs[best], 6),
            "sharpe": round(daily_sharpe(returns[keep]), 6), "qqq_sharpe": round(daily_sharpe(qqq[keep]), 6)}


def study(name):
    axes, original = DESIGNS[name]
    grid = {str(s): stats(backtest(name, s)) for s in product(*axes)}
    curve = backtest(name, original)
    wins = sum(cell["beats_qqq"] for cell in grid.values())
    first, second = halves(curve)
    result = {
        "name": name, "original": str(original), "original_stats": grid[str(original)], "grid": grid,
        "grid_wins": wins, "grid_size": len(grid), "first_half": first, "second_half": second,
        "with_cost": stats(backtest(name, original, COST_PER_SIDE)), "yearly": yearly(curve),
        "without_best_year": without_best_year(curve),
    }
    result["passes"] = bool(
        wins >= 0.7 * len(grid) and first["beats_qqq"] and second["beats_qqq"]
        and result["original_stats"]["beats_qqq"] and result["with_cost"]["beats_qqq"]
        and result["without_best_year"]["sharpe"] > result["without_best_year"]["qqq_sharpe"])
    return result


def holding_returns(chosen):
    """Return of each pick from the fill after its signal to the fill after the next signal."""
    prices_ = prices()
    days = list(chosen)
    out = {}
    for day, end in zip(days, [*days[1:], prices_.index[-2]]):
        fill = prices_.index[prices_.index.searchsorted(day) + 1]
        last = prices_.index[prices_.index.searchsorted(end) + 1]
        out[str(day.date())] = {t: round(float(prices_.loc[last, t] / prices_.loc[fill, t] - 1), 6) for t in chosen[day]}
    return out


def without_2022(curve):
    """Sharpe and CAGR of the second half with 2022 taken out, against QQQ."""
    second = curve.loc[curve.index[len(curve) // 2]:]
    out = {}
    for label, series in (("dogs", second), ("qqq", bench("QQQ", second.index))):
        returns = series.pct_change().dropna()
        returns = returns[returns.index.year != 2022]
        out[label] = {"sharpe": round(daily_sharpe(returns), 6),
                      "cagr": round((1 + returns).prod() ** (252 / len(returns)) - 1, 6)}
    return out


def names_replaced(chosen):
    lists = list(chosen.values())
    return round(sum(len(set(b) - set(a)) for a, b in zip(lists, lists[1:])) / (len(lists) - 1), 2)


def main():
    original = DESIGNS["dogs_of_the_dow"][1]
    results = {name: study(name) for name in DESIGNS}
    for result in results.values():
        for fund in ("DIA", "SPY"):
            result[f"{fund.lower()}_wins"] = int(sum(c["sharpe"] > c[f"{fund.lower()}_sharpe"] for c in result["grid"].values()))
    curves = {name: backtest(name, DESIGNS[name][1]) for name in DESIGNS}
    curves["rest_of_dow"] = backtest("rest_of_dow", original)
    check = {name: run(picks(name, *DESIGNS[name][1], True)) for name in DESIGNS}
    chosen = picks("dogs_of_the_dow", *original)
    pick_yields = {str(day.date()): dow_yields(day).loc[names] for day, names in chosen.items()}
    metrics = {
        "end": END, "run_date": str(pd.Timestamp.now(tz="Asia/Taipei").date()), "finlab_version": finlab.__version__,
        "designs": results,
        "published_style": {name: {**stats(curve), "yearly": yearly(curve)} for name, curve in check.items()},
        "rest_of_dow": {**stats(curves["rest_of_dow"]), "yearly": yearly(curves["rest_of_dow"]),
                        "halves": halves(curves["rest_of_dow"])},
        "second_half_without_2022": without_2022(curves["dogs_of_the_dow"]),
        "names_replaced_a_year": names_replaced(chosen),
        "yearly_funds": {t: yearly(bench(t, curves["dogs_of_the_dow"].index)) for t in FUNDS},
        "picks": {day: {t: round(float(value), 6) for t, value in yields.items()}
                  for day, yields in pick_yields.items()},
        "average_pick_yield": round(float(pd.Series([float(y) for yields in pick_yields.values() for y in yields]).mean()), 6),
        "pick_returns": holding_returns(chosen),
        "picks_published_style": {str(day.date()): list(names)
                                  for day, names in picks("dogs_of_the_dow", *original, True).items()},
        "small_picks": {str(day.date()): list(names) for day, names in picks("small_dogs", 5, 12).items()},
    }
    (OUT / "metrics.json").write_text(json.dumps(metrics, indent=1))
    table = pd.DataFrame(curves)
    for ticker in FUNDS:
        table[ticker] = bench(ticker, table.index)
    (table / table.iloc[0]).round(6).rename_axis("date").to_csv(OUT / "data.csv")
    for name, result in results.items():
        print(name, result["original_stats"]["sharpe"], result["grid_wins"], "/", result["grid_size"], result["passes"])


if __name__ == "__main__":
    main()
