"""隔日沖分點：事前登記的 9 個設計、81 組參數（FinLab，2026-11）。

安裝 finlab>=2.3.8、plotly、kaleido 後執行 python research.py。
分點逐年讀取，需數 GB 磁碟與足夠記憶體。只讀資料，sim(upload=False)。
輸出檔寫在本程式旁，--cache 可指定可丟棄的中間資料目錄。
更新資料或修改研究程式後，請指定新的空快取目錄，避免沿用舊快照。
"""
import argparse
import contextlib
import gc
import io
import json
import runpy
import tempfile
import urllib.request
from itertools import product
from pathlib import Path

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

START, END, RAW_START = "2023-07-03", "2026-11-09", "2022-07-01"
LOOKBACKS, THRESHOLDS, HORIZONS = (60, 120, 240), (0.4, 0.5, 0.6), (5, 10, 20)
CENTER = (120, 0.5)  # pre-registered (lookback, threshold); the others are neighbours
BASES = {
    "monthly-revenue-momentum": "monthly-revenue-momentum/validation/strategy.py",
    "small-cap-effect": "small-cap-effect/validation/strategy.py",
    "operating-margin-three-factor": "ai-agent-quant-trading-tools/operating-margin-three-factor/strategy.py",
}
OUT = Path(__file__).resolve().parent
REPO = next((p for p in OUT.parents if (p / "content/strategies-data").is_dir()), None)


def get(key):
    return pd.DataFrame(data.get(key, progress="silent"))


def sharpe(returns):
    return float(returns.mean() / returns.std() * np.sqrt(252)) if returns.std() else 0.0


def stats(curve):
    curve = curve.dropna()
    returns = curve.pct_change(fill_method=None).dropna()
    return {
        "cagr": float((curve.iloc[-1] / curve.iloc[0]) ** (365.25 / (curve.index[-1] - curve.index[0]).days) - 1),
        "sharpe": sharpe(returns),
        "mdd": float((curve / curve.cummax() - 1).min()),
    }


def historical_rate(counts, successes, lookback):
    # A pair ending today cannot classify today's branch. It first counts tomorrow.
    denominator = counts.rolling(lookback, min_periods=lookback).sum().shift(1)
    numerator = successes.rolling(lookback, min_periods=lookback).sum().shift(1)
    return (numerator / denominator).where(denominator >= 20), denominator


def branch_events(days, stocks, cache):
    """Rank observed net buyers/sellers, then pair *adjacent market sessions*, across years."""
    path = cache / f"events-{RAW_START}-{END}.pickle"
    if path.exists():
        return pd.read_pickle(path)
    buys, sells, rows = [], [], 0
    for year in range(int(RAW_START[:4]), int(END[:4]) + 1):
        raw = data.get("broker_transactions", start=max(RAW_START, f"{year}-01-01"),
                       end=min(END, f"{year}-12-31"), progress="silent")
        rows += len(raw)
        raw = raw.loc[raw.stock_id.isin(stocks), ["date", "stock_id", "broker", "buy", "sell"]].copy()
        raw["broker"] = raw.broker.str.replace("臺", "台", regex=False)
        raw["d"] = days.get_indexer(raw.date)
        raw["net"] = (raw.buy - raw.sell).astype(float)
        raw[["buy", "sell"]] = raw[["buy", "sell"]].astype(float)
        assert not raw.duplicated(["d", "stock_id", "broker"]).any()
        groups = raw.groupby(["d", "stock_id"], observed=True).net
        buys.append(raw.loc[(raw.net > 0) & (groups.rank(ascending=False, method="first") <= 15)])
        sells.append(raw.loc[(raw.net < 0) & (groups.rank(ascending=True, method="first") <= 15)])
        del raw
        gc.collect()
        print("read year", year, "source rows", rows, flush=True)
    buy, sell = pd.concat(buys, ignore_index=True), pd.concat(sells, ignore_index=True)
    branches = pd.Index(sorted(set(buy.broker) | set(sell.broker)))
    for frame in (buy, sell):
        frame["b"] = branches.get_indexer(frame.broker)
        frame.drop(columns=["date", "broker"], inplace=True)
    # The denominator includes purchases with no observed next-day top-15 sale.
    mature = buy.loc[buy.d + 1 < len(days), ["d", "stock_id", "b", "buy"]].copy()
    mature["d"] += 1
    pairs = mature.merge(sell[["d", "stock_id", "b", "sell"]], on=["d", "stock_id", "b"], how="left")
    pairs["success"] = pairs.sell >= pairs.buy * 0.8
    counts = pd.crosstab(pairs.d, pairs.b).reindex(index=range(len(days)), columns=range(len(branches)), fill_value=0)
    success = pd.crosstab(pairs.loc[pairs.success, "d"], pairs.loc[pairs.success, "b"])
    success = success.reindex(index=counts.index, columns=counts.columns, fill_value=0)
    result = buy, sell, branches, counts, success, {"raw_rows": rows, "purchases": len(buy), "pairs": len(pairs)}
    pd.to_pickle(result, path)
    return result


def large_trades(frame, stocks, volume, liquid):
    """Keep (day, branch, stock) positions of large net trades in liquid stocks; parameter-independent."""
    d, b, col = frame.d.to_numpy(), frame.b.to_numpy(), stocks.get_indexer(frame.stock_id)
    size = frame.net.abs().to_numpy()
    keep = (size >= 100) & (size * 1000 >= volume.to_numpy()[d, col] * 0.02) & liquid.to_numpy()[d, col]
    return d[keep], b[keep], col[keep]


def signal(trades, qualified, days, stocks):
    d, b, col = trades
    valid = qualified.to_numpy()[d, b]
    array = np.zeros((len(days), len(stocks)), dtype=bool)
    array[d[valid], col[valid]] = True
    return pd.DataFrame(array, index=days, columns=stocks)


def run(position, **settings):
    return finlab.backtest.sim(position, upload=False, end_date=END, retain_initial_cash=True, **settings)


def capture_base(slug, cache):
    """Read our published strategy source, preserving its selection, stops and schedule."""
    path = cache / f"base-{slug}-{END}.pickle"
    if path.exists():
        return pd.read_pickle(path)
    source = REPO / "static/blog" / BASES[slug] if REPO else None
    if source and source.exists():
        code = source.read_text()
    else:
        request = urllib.request.Request("https://finlab.finance/blog/" + BASES[slug], headers={"User-Agent": "FinLab-research"})
        code = urllib.request.urlopen(request).read().decode()
    calls, real_sim = [], finlab.backtest.sim

    def record(position, **settings):
        settings.update(upload=False, end_date=END)
        report = real_sim(position, **settings)
        calls.append((report, settings))
        return report

    finlab.backtest.sim = record
    try:
        with tempfile.TemporaryDirectory() as tmp, contextlib.chdir(tmp), contextlib.redirect_stdout(io.StringIO()):
            script = Path(tmp) / "strategy.py"
            script.write_text(code)
            runpy.run_path(str(script), run_name="__main__")
    finally:
        finlab.backtest.sim = real_sim
    report, settings = calls[-1]
    original_stats = stats(report.creturn)
    if REPO:
        page = json.loads((REPO / "content/strategies-data" / f"{slug}.json").read_text())["stats"]
        same_window = stats(report.creturn.loc[page["window"]["start"]:page["window"]["end"]])
        original_stats["page_comparison"] = {"rerun": same_window, "page": page}
        assert abs(same_window["cagr"] - page["cagr"]) <= 0.005, slug
        assert abs(same_window["sharpe"] - page["sharpe"]) <= 0.02, slug
    # Re-run the *sampled* monthly weights at next-day open for both controls and filters.
    weights = report.position.astype(float)
    days = get("price:收盤價").loc[:END].index
    weights = weights.loc[weights.index < days[-1] + pd.offsets.BDay()].copy()
    weights.index = days[days.searchsorted(weights.index, side="right") - 1]
    weights = weights.loc[~weights.index.duplicated(keep="last")]
    settings = {k: v for k, v in settings.items() if k not in ("upload", "end_date", "name", "resample", "resample_offset")}
    settings.update(trade_at_price="open", resample=weights.index)
    opened = run(weights.reindex(days[days >= weights.index[0]], method="ffill"), **settings)
    result = opened, settings, original_stats, weights
    pd.to_pickle(result, path)
    return result


def filtered_position(base, buy_signal, days):
    report, settings, _, weights = base
    weights = weights.reindex(columns=report.position.columns)
    daily = weights.reindex(days, method="ffill").fillna(0)
    # Once excluded, do not re-enter before the next scheduled monthly selection.
    groups = pd.Series(np.arange(len(weights)), index=weights.index).reindex(days, method="ffill")
    excluded = buy_signal.reindex(index=days, columns=daily.columns, fill_value=False).groupby(groups).cummax()
    excluded &= daily.ne(0)
    target = daily.to_numpy(copy=True)
    target[excluded.to_numpy()] = 0
    # Extra rebalances may remove a stock but must never resurrect a base stop-out.
    alive = np.zeros(target.shape, dtype=bool)
    for trade in report.trades.itertuples():
        column = daily.columns.get_indexer([str(trade.stock_id).split()[0]])[0]
        if column < 0:
            continue
        begin = days.searchsorted(trade.entry_sig_date)
        end = len(days) if pd.isna(trade.exit_sig_date) else days.searchsorted(trade.exit_sig_date)
        alive[begin:end, column] = True
    alive[days.isin(weights.index), :] = True
    target[~alive] = 0
    target = pd.DataFrame(target, index=days, columns=daily.columns)
    target.attrs["excluded_stock_days"] = int(excluded.loc[START:].sum().sum())
    # Stop-outs are already executed by sim. Adding their dates as rebalances would
    # also rebalance the other holdings, changing even an all-clear control.
    changes = excluded.index[excluded.ne(excluded.shift()).any(axis=1)]
    rebalance = changes.union(weights.index.intersection(days))
    return target, {**settings, "resample": rebalance}


def compare(curve, bench):
    mid = len(curve) // 2
    output = {}
    for label, part in [("full", slice(None)), ("first", slice(None, mid)), ("second", slice(mid, None))]:
        strategy, benchmark = stats(curve.iloc[part]), stats(bench.iloc[part])
        output[label] = {"strategy": strategy, "benchmark": benchmark,
                         "win": strategy["sharpe"] > benchmark["sharpe"] + 1e-8}
    return output


def three_checks(rows, center):
    """At least 7 of the 9 parameter sets win, and the center setting wins in both halves."""
    return sum(r["full"]["win"] for r in rows) >= 7 and center["first"]["win"] and center["second"]["win"]


def additional_checks(curve, benchmark, exposure):
    returns, market = curve.pct_change().dropna(), benchmark.pct_change().dropna()
    annual = (1 + returns).groupby(returns.index.year).prod() - 1
    best = int(annual.idxmax())
    keep = returns.index.year != best
    matched = (1 + market * exposure).cumprod().reindex(curve.index).fillna(1)
    correlations = {}
    if REPO:
        weekly = curve.resample("W-FRI").last().pct_change().dropna()
        for path in sorted((REPO / "content/strategies-data").glob("*.json")):
            page = json.loads(path.read_text())
            points = page.get("equitySeries", [])
            if page.get("market", "TW") == "US" or not page.get("strategyResearch") or not points:
                continue
            reference = pd.Series([p["value"] for p in points], index=pd.to_datetime([p["date"] for p in points]))
            joint = pd.concat([weekly, reference.resample("W-FRI").last().pct_change()], axis=1, sort=True).dropna()
            if len(joint) >= 60:
                value = float(joint.iloc[:, 0].corr(joint.iloc[:, 1]))
                if np.isfinite(value):
                    correlations[path.stem] = value
    return {"matched_exposure": compare(curve, matched), "best_year_removed": best,
            "without_best_year": {"strategy": sharpe(returns.loc[keep]), "benchmark": sharpe(market.reindex(returns.index).loc[keep])},
            "correlations": correlations}


def forward_returns(close, opening):
    returns = {"next_open_close": close.shift(-1) / opening.shift(-1) - 1}
    returns.update({f"close_{h}": close.shift(-h) / close - 1 for h in HORIZONS})
    return returns


def event_stats(mask, forwards):
    output = {}
    for name, forward in forwards.items():
        masked = forward.where(mask)
        values = masked.stack().dropna()
        daily = masked.mean(axis=1).dropna()
        if values.empty:
            output[name] = {"n": 0, "mean": None, "median": None, "positive": None, "equal_day_mean": None}
            continue
        output[name] = {"n": len(values), "mean": float(values.mean()), "median": float(values.median()),
                        "positive": float((values > 0).mean()),
                        "equal_day_mean": float(daily.mean()) if len(daily) else None}
    return output


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--cache", type=Path, default=Path(tempfile.gettempdir()) / "finlab-overnight-440")
    args = parser.parse_args()
    args.cache.mkdir(parents=True, exist_ok=True)
    data.truncate_end = END
    close_all = get("etl:adj_close")
    days = close_all.loc[RAW_START:END].index
    stocks = close_all.columns[close_all.columns.str.fullmatch(r"[1-9]\d{3}")]
    close = close_all.reindex(index=days, columns=stocks)
    opening = get("etl:adj_open").reindex(index=days, columns=stocks)
    volume = get("price:成交股數").reindex(index=days, columns=stocks)
    liquid = get("price:成交金額").rolling(20).mean().reindex(index=days, columns=stocks) >= 20_000_000
    buy, sell, branches, counts, success, coverage = branch_events(days, stocks, args.cache)
    window = days[days >= START]
    benchmark = close_all["0050"].reindex(window).ffill()
    bases = {slug: capture_base(slug, args.cache) for slug in BASES}
    baseline = {slug: base[0].creturn.reindex(window).ffill() for slug, base in bases.items()}
    for slug, base in bases.items():
        target, settings = filtered_position(base, pd.DataFrame(False, index=days, columns=stocks), days)
        control = run(target, **settings).creturn.reindex(window).ffill()
        reference = baseline[slug]
        assert np.allclose(control / control.iloc[0], reference / reference.iloc[0], atol=1e-8), slug
    curves, results, events_grid = {"0050": benchmark / benchmark.iloc[0]}, {}, {}
    # Everything below up to the loop is parameter-independent: compute it once.
    large = {"follow": large_trades(buy, stocks, volume, liquid), "rebound": large_trades(sell, stocks, volume, liquid)}
    del buy, sell  # the full top-15 tables are no longer needed; free their memory before the backtests
    forwards = forward_returns(close.loc[START:], opening.loc[START:])
    rates = {lookback: historical_rate(counts, success, lookback) for lookback in LOOKBACKS}
    months = pd.period_range(START[:7], END[:7], freq="M")
    for lookback, threshold in product(LOOKBACKS, THRESHOLDS):
        qualified = rates[lookback][0] >= threshold
        signals = {kind: signal(trades, qualified, days, stocks) for kind, trades in large.items()}
        signal_days = {kind: int(mask.loc[START:].sum().sum()) for kind, mask in signals.items()}
        events_grid[f"{lookback}-{threshold}"] = {
            kind: event_stats(mask.loc[START:], forwards) for kind, mask in signals.items()}
        is_center = (lookback, threshold) == CENTER
        designs = {}  # name -> (position, sim settings, signal kind, base slug or None)
        for kind, mask in signals.items():
            for horizon in HORIZONS:
                holding = mask.rolling(horizon, min_periods=1).max().astype(bool).loc[START:]
                position = holding.div(holding.sum(axis=1).clip(lower=10), axis=0)
                designs[f"{kind}_{horizon}"] = position, {"trade_at_price": "open"}, kind, None
        for slug, base in bases.items():
            position, settings = filtered_position(base, signals["follow"], days)
            designs[f"exclude_{slug}"] = position, settings, "follow", slug
        for design, (position, settings, kind, slug) in designs.items():
            key = f"{design}-{lookback}-{threshold}"
            cached = args.cache / f"{START}-{END}-{key}.pickle"
            if cached.exists():
                report = pd.read_pickle(cached)
            else:
                report = run(position, **settings)
                pd.to_pickle(report, cached)
            curve = report.creturn.reindex(window).ffill().fillna(1)
            trades = report.trades.loc[report.trades.entry_date >= pd.Timestamp(START)]
            entry_month = pd.to_datetime(trades.entry_date).dt.to_period("M")
            turnover = trades.position.abs().groupby(entry_month).sum().reindex(months, fill_value=0)
            row = {"lookback": lookback, "threshold": threshold, **compare(curve, benchmark),
                   "average_exposure": float(position.loc[START:].sum(axis=1).mean()),
                   "signal_stock_days": signal_days[kind],
                   "trades": len(trades), "monthly_one_way_turnover": float(turnover.mean()),
                   "max_monthly_one_way_turnover": float(turnover.max()),
                   "max_turnover_month": str(turnover.idxmax()) if turnover.max() else None}
            if slug:
                row["base"] = compare(curve, baseline[slug])
                row["excluded_stock_days"] = position.attrs["excluded_stock_days"]
            results.setdefault(design, []).append(row)
            normalized = curve / curve.iloc[0]
            curves[f"{design}@{lookback}-{threshold}"] = normalized
            if is_center:
                curves[design] = normalized
                row["additional_checks"] = additional_checks(curve, benchmark, row["average_exposure"])
            if lookback == 60 and threshold == 0.4 and design in ("follow_10", "rebound_10"):
                report.to_html(str(OUT / f"report_{design}_60_40.html"))
            print(key, row["full"]["strategy"], flush=True)
            del report
            gc.collect()
    rate, denominator = rates[CENTER[0]]
    last_rate, last_count = rate.iloc[-1], denominator.iloc[-1]
    observed = sorted(({"branch": branches[i], "rate": float(last_rate.iloc[i]), "purchases": int(last_count.iloc[i])}
                       for i in np.flatnonzero(last_rate.notna().to_numpy())),
                      key=lambda r: (-r["rate"], -r["purchases"]))
    latest = [r for r in observed if r["rate"] >= CENTER[1]]
    summary = {}
    for design, rows in results.items():
        center = next(r for r in rows if (r["lookback"], r["threshold"]) == CENTER)
        wins = sum(r["full"]["win"] for r in rows)
        passed = three_checks(rows, center)
        if "base" in center:
            passed = passed and three_checks([r["base"] for r in rows], center["base"])
        additional = center["additional_checks"]
        robust = additional["without_best_year"]["strategy"] > additional["without_best_year"]["benchmark"]
        distinct = all(c < 0.95 for c in additional["correlations"].values())
        matched_wins = additional["matched_exposure"]["full"]["win"]
        summary[design] = {"wins": wins, "three_checks_pass": bool(passed),
                           "validated": bool(passed and robust and distinct and matched_wins), "center": center}
    for slug, curve in baseline.items():
        curves["base_" + slug] = curve / curve.iloc[0]
    metrics = {"run_date": str(pd.Timestamp.now(tz="Asia/Taipei").date()), "finlab_version": finlab.__version__,
               "window": f"{START}/{END}", "coverage": coverage, "summary": summary, "grids": results,
               "events": events_grid[f"{CENTER[0]}-{CENTER[1]}"], "events_grid": events_grid, "branches": latest,
               "observed_branches": observed, "base_original_stats": {k: v[2] for k, v in bases.items()}}
    (OUT / "metrics.json").write_text(json.dumps(metrics, ensure_ascii=False, indent=2, allow_nan=False))
    pd.DataFrame(curves).to_csv(OUT / "data.csv", index_label="date")
    print("DONE", {k: v["three_checks_pass"] for k, v in summary.items()}, flush=True)


if __name__ == "__main__":
    main()