"""Four preregistered ETF designs, 36 settings.

Run with finlab >=2.3, pandas and numpy. FinLab handles account login.
Outputs daily equity, all statistics, and an HTML report for every simulation.
Signals at a close fill at the next close; no leverage or short positions.
"""
from functools import cache
from itertools import product
from pathlib import Path
import json

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

END = "2026-10-30"
OUT = Path(__file__).resolve().parent
MOMENTUM = ["SPMO", "MTUM", "XMMO", "PDP", "IMTM", "FDMO"]
FACTORS = ["MTUM", "QUAL", "USMV", "VLUE"]
ASSETS = list(dict.fromkeys(MOMENTUM + FACTORS + ["SPLV", "NOBL", "QQQ", "SPY", "BIL"]))
DESIGNS = {
    "A_rotation": ((9, 12, 15), (1, 2, 3), (12, 2)),
    "B_spmo_trend": ((150, 200, 250), ("W", "M", "Q"), (200, "M")),
    "C_rotation_trend": ((9, 12, 15), (150, 200, 250), (12, 200)),
    "D_factor_rotation": ((9, 12, 15), (1, 2, 3), (12, 2)),
}


@cache
def prices():
    data.truncate_end = END
    data.set_market("us_fund")
    # The stock-price calendar removes holiday rows in the fund-price table.
    calendar = data.get("us_price:adj_close").index
    return pd.DataFrame(data.get("us_fund_price:adj_close"))[ASSETS].reindex(calendar).loc[:END].ffill()


def month_ends(index):
    return pd.DatetimeIndex(pd.Series(index, index=index).groupby(index.to_period("M")).max())


def start(name):
    assets = ["SPMO"] if name.startswith("B_") else FACTORS if name.startswith("D_") else MOMENTUM
    days = prices()[assets].dropna().index
    ready = days[249 if name.startswith("B_") else 21 * 15]
    ends = month_ends(prices().index)
    return ends[ends >= ready][0]


def weights(name, a, b):
    p = prices()
    if name.startswith("B_"):
        above = (p.SPMO > p.SPMO.rolling(a).mean()).astype(float)
        return pd.DataFrame({"SPMO": above, "BIL": 1 - above})
    assets = FACTORS if name.startswith("D_") else MOMENTUM
    score = p[assets] / p[assets].shift(21 * a) - 1
    k = 2 if name.startswith("C_") else b
    chosen = (score.rank(axis=1, ascending=False, method="first") <= k).astype(float)
    w = chosen.div(chosen.sum(axis=1), axis=0).fillna(0)
    if name.startswith("C_"):
        above = p.QQQ > p.QQQ.rolling(b).mean()
        w = w.mul(above, axis=0)
        w["BIL"] = (~above).astype(float)
    return w


def sharpe(r):
    return float(r.mean() / r.std() * np.sqrt(252)) if r.std() > 0 else None


def stats(curve):
    curve = curve.dropna()
    r = curve.pct_change(fill_method=None).dropna()
    bills = prices().BIL.reindex(curve.index).pct_change(fill_method=None).reindex(r.index)
    years = (curve.index[-1] - curve.index[0]).days / 365.25
    return {"start": str(curve.index[0].date()), "end": str(curve.index[-1].date()),
            "total_return": float(curve.iloc[-1] / curve.iloc[0] - 1),
            "cagr": float((curve.iloc[-1] / curve.iloc[0]) ** (1 / years) - 1),
            "sharpe": sharpe(r), "excess_sharpe": sharpe(r - bills),
            "mdd": float((curve / curve.cummax() - 1).min()),
            "volatility": float(r.std() * np.sqrt(252))}


def halves(curve):
    middle = len(curve) // 2
    return [stats(curve.iloc[:middle + 1]), stats(curve.iloc[middle:])]


def comparisons(curve):
    qqq = prices().QQQ.reindex(curve.index)
    return {"stats": stats(curve), "qqq": stats(qqq),
            "halves": halves(curve), "qqq_halves": halves(qqq)}


class QQQMarket(USFundMarket):
    @staticmethod
    def get_benchmark():
        return prices().QQQ


def run(w, frequency, tag, fee=None):
    report = sim(w, resample=frequency, trade_at_price="close", fee_ratio=fee,
                 market=QQQMarket(), upload=False)
    path = OUT / "reports" / (tag + ".html")
    path.parent.mkdir(exist_ok=True)
    report.to_html(str(path), title=tag)
    calendar = prices().loc[w.index[0]:END].index
    curve = report.creturn.reindex(calendar).ffill().fillna(1.0)
    # Include initial trading costs; share a first-fill baseline across settings.
    return curve


def evaluate(name):
    va, vb, original = DESIGNS[name]
    since = start(name)
    grid, curves = [], {}
    for a, b in product(va, vb):
        w = weights(name, a, b).loc[since:]
        frequency = b if name.startswith("B_") else "M"
        curve = run(w, frequency, f"{name}_{a}_{b}")
        curves[a, b] = curve
        row = {"a": a, "b": b, **comparisons(curve)}
        grid.append(row)
        print(name, a, b, round(row["stats"]["sharpe"], 3), flush=True)
    curve = curves[original]
    base = next(r for r in grid if (r["a"], r["b"]) == original)
    w = weights(name, *original).loc[since:]
    frequency = original[1] if name.startswith("B_") else "M"
    cost = comparisons(run(w, frequency, name + "_cost", fee=0.001))
    sampled = w.loc[month_ends(w.index)].iloc[:-1]
    average = sampled.mean()
    fixed = comparisons(run(pd.DataFrame(dict(average), index=w.index), "M", name + "_fixed"))
    standalone = prices().SPMO if name.startswith("B_") else curve
    if name.startswith("C_"):
        standalone = run(weights("A_rotation", original[0], 2).loc[since:], "M", name + "_no_trend")
    standalone = comparisons(standalone.reindex(curve.index))
    wins = sum(r["stats"]["sharpe"] > r["qqq"]["sharpe"] for r in grid)
    checks = {"grid": wins >= 7,
              "halves": all(a["sharpe"] > b["sharpe"] for a, b in zip(base["halves"], base["qqq_halves"])),
              "cost": cost["stats"]["sharpe"] > base["qqq"]["sharpe"],
              "above_bil": base["stats"]["excess_sharpe"] > base["qqq"]["excess_sharpe"],
              "fixed_weights": base["stats"]["sharpe"] > fixed["stats"]["sharpe"]}
    if name.startswith(("B_", "C_")):
        checks["standalone"] = base["stats"]["sharpe"] > standalone["stats"]["sharpe"]
    # Target-weight turnover is a diagnostic, excluding initial allocation.
    changes = w.resample({"W": "W-FRI", "M": "ME", "Q": "QE"}[frequency]).last().diff().abs().sum(axis=1).iloc[1:] / 2
    result = {"grid": grid, "grid_wins": wins, "original": list(original), "original_stats": base,
              "cost": cost, "fixed": fixed,
              "fixed_weights": average.to_dict(), "standalone": standalone,
              "turnover_mean": float(changes.mean()), "turnover_max": float(changes.max()),
              "checks": checks, "passes": all(checks.values())}
    return result, curve


def main():
    p = prices()
    common = p.dropna().index[0]
    out = {"run_date": pd.Timestamp.now(tz="UTC").isoformat(), "end": END,
           "finlab_version": finlab.__version__, "pandas_version": pd.__version__,
           "common_hold_start": str(common.date()), "designs": {},
           "coverage": {t: str(p[t].first_valid_index().date()) for t in ASSETS},
           "hold": {t: {**comparisons(p[t].loc[common:]), "since_first_price": comparisons(p[t].dropna())} for t in ASSETS}}
    curves = {}
    for name in DESIGNS:
        result, curve = evaluate(name)
        out["designs"][name] = result
        curves[name] = curve / curve.iloc[0]
        print(name, result["checks"], flush=True)
    curves.update({t: p[t] / p[t].loc[common] for t in ASSETS})
    pd.DataFrame(curves).to_csv(OUT / "data.csv", index_label="date")
    (OUT / "metrics.json").write_text(json.dumps(out, indent=1, allow_nan=False))


if __name__ == "__main__":
    main()