"""
暫停交易、恢復交易後買進：3 個設計 27 組參數的三項驗證與事件統計（2016 年起至最新資料）

三項驗證（「贏」＝日夏普高於同期 0050 含息買進持有）：
  1. 參數高原：每個設計 2 個主要參數各取 3 個值，9 組裡至少 7 組要贏。
  2. 前後兩段：回測期間對半切，兩段各自要贏。
  3. 交易成本：sim() 預設手續費 0.1425%（買賣各一次）與證交稅 0.3%（賣出）。
時點：訊號日收盤後才知道，sim() 在隔一個交易日收盤成交。流動性＝訊號日 20 日均成交金額 ≥ 2,000 萬。
0050 含息：etl:adj_close 還原價買進持有，不經 sim、不含成本。日夏普＝日報酬平均 ÷ 標準差 × √252，不扣無風險利率。
每個訊號的定義與參數在看結果前就寫定，之後沒有改過。

暫停交易機制 2016-01-15 起實施，個股事件從這天開始。
trading_halt 的 date 是暫停與恢復發生日；恢復日的反應從暫停前最後一個收盤算起，不把停牌天算成 0。

執行：安裝 finlab 後直接跑，資料需要時會自動引導登入，約 30 秒。
輸出：metrics.json（全部結果）、data.csv（各設計原參數、0050 含息與等權流動股的每日淨值）。
"""
import datetime
import json
import math

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

COSTS = {"default": {}, "zero": {"fee_ratio": 0, "tax_ratio": 0}}


def plain(key):
    return pd.DataFrame(data.get(key))


# ---------------------------------------------------------------------------
# 1. 共同資料與績效計算
# ---------------------------------------------------------------------------
close = plain("price:收盤價")
days = close.index
ND = len(days)
adj_close = plain("etl:adj_close").reindex(days)
bench = adj_close["0050"].dropna()

# 上市櫃 4 碼普通股（不含 0 開頭的 ETF、91 開頭的存託憑證），含已下市
STOCKS = [c for c in close.columns if len(c) == 4 and c.isdigit() and c[0] != "0" and not c.startswith("91")]
COL = {c: i for i, c in enumerate(STOCKS)}
amount = plain("price:成交金額").reindex(index=days, columns=STOCKS)
amount_20d = amount.rolling(20, min_periods=20).mean()
adj_frame = adj_close.reindex(columns=STOCKS).ffill()
adj = adj_frame.to_numpy()
adj0050 = adj_close["0050"].ffill().to_numpy()


def liquid_mask(floor=2e7):
    """20 日均成交金額 ≥ floor（預設 2,000 萬）。"""
    return (amount_20d >= floor).fillna(False).to_numpy()


liquid = liquid_mask()


def month_ends():
    s = pd.Series(np.arange(ND), index=days)
    return s.groupby([s.index.year, s.index.month]).max().tolist()


def stats(curve):
    curve = curve.dropna()
    curve = curve / curve.iloc[0]
    ret = curve.pct_change().dropna()
    years = (curve.index[-1] - curve.index[0]).days / 365.25
    sd = ret.std()
    return {"start": str(curve.index[0].date()), "end": str(curve.index[-1].date()),
            "cagr": round(float(curve.iloc[-1] ** (1 / years) - 1), 6) if years > 0 else None,
            "sharpe": round(float(ret.mean() / sd * math.sqrt(252)), 6) if sd > 0 else None,
            "mdd": round(float((curve / curve.cummax() - 1).min()), 6),
            "total_return": round(float(curve.iloc[-1] - 1), 6)}


def active(curve):
    """從第一次淨值變動的前一個交易日起算。"""
    curve = curve.dropna()
    moved = np.flatnonzero(curve.values != curve.values[0])
    return curve if len(moved) == 0 else curve.iloc[max(moved[0] - 1, 0):]


def compare(curve):
    s = stats(active(curve))
    b = stats(bench.loc[s["start"]:s["end"]])
    return {"strategy": s, "bench": b, "win": bool(s["sharpe"] > b["sharpe"])}


def halves(curve):
    curve = active(curve)
    mid = curve.index[0] + (curve.index[-1] - curve.index[0]) / 2
    out = []
    for part in (curve.loc[:mid], curve.loc[mid:]):
        s = stats(part)
        b = stats(bench.loc[s["start"]:s["end"]])
        out.append({"strategy": s, "bench": b, "win": bool(s["sharpe"] > b["sharpe"])})
    return out


def backtest(table, cost, events=None):
    frame = pd.DataFrame(table, index=days, columns=STOCKS)
    report = sim(frame.loc[:, frame.any()], upload=False, **COSTS[cost])
    return report, frame, events


def event_sim(rows, cost):
    """rows：(股票, b, e)＝b 收盤買、e 收盤賣。sim 在持有表為真的隔天收盤成交，所以持有表是 [b−1, e−2]。"""
    table = np.zeros((ND, len(STOCKS)), dtype=bool)
    count = 0
    for s, b, e in rows:
        lo, hi = max(int(b) - 1, 0), min(int(e) - 2, ND - 1)
        if s in COL and hi >= lo and lo < ND:
            table[lo:hi + 1, COL[s]] = True
            count += 1
    return backtest(table, cost, count)


def period_sim(selections, cost):
    """selections：[(t, 股票清單)]，t 收盤決定，持有到下一個 t。"""
    table = np.zeros((ND, len(STOCKS)), dtype=bool)
    for k, (t, stocks) in enumerate(selections):
        end = selections[k + 1][0] if k + 1 < len(selections) else ND
        table[t:end, [COL[s] for s in stocks if s in COL]] = True
    return backtest(table, cost)


def position_stats(report, frame, events):
    curve = active(report.creturn)
    held = frame.sum(axis=1).reindex(curve.index).fillna(0)
    buys = (frame & ~frame.shift(1, fill_value=False)).sum().sum()
    years = max((curve.index[-1] - curve.index[0]).days / 365.25, .01)
    return {"events": events, "buys_per_year": round(float(buys / years), 1),
            "invested": round(float((held > 0).mean()), 4),
            "avg_holdings": round(float(held[held > 0].mean()), 2) if (held > 0).any() else None}


def equal_weight_curve():
    """等權流動股：每月底買進全部流動股、等權、含預設成本，當第二個對照組。"""
    picks = [(m, [STOCKS[j] for j in np.flatnonzero(liquid[m])]) for m in month_ends() if liquid[m].any()]
    return period_sim(picks, "default")[0].creturn


def validate(key, equal_weight):
    design = DESIGNS[key]
    xs, ys, base = design["x"], design["y"], design["base"]
    grid, base_run = [], None
    for x in xs:
        for y in ys:
            run = design["fn"](x, y, "default")
            row = compare(run[0].creturn)
            row.update({"x": x, "y": y, "pos": position_stats(*run)})
            grid.append(row)
            if (x, y) == base:
                base_run = run
    base_row = next(g for g in grid if (g["x"], g["y"]) == base)
    s = base_row["strategy"]
    result = {"title": design["title"], "x_name": design["x_name"], "y_name": design["y_name"],
              "x": xs, "y": ys, "base": list(base), "grid": grid,
              "grid_wins": sum(g["win"] for g in grid),
              "base_default_cost": base_row,
              "halves": halves(base_run[0].creturn),
              "base_zero_cost": compare(design["fn"](*base, "zero")[0].creturn),
              "equal_weight_same_window": stats(equal_weight.loc[s["start"]:s["end"]])}
    h = result["halves"]
    result["pass"] = {"neighbors": result["grid_wins"] >= 7, "halves": h[0]["win"] and h[1]["win"], "cost": base_row["win"]}
    return result, base_run[0]


# ---------------------------------------------------------------------------
# 2. 事件研究：減等權流動股（每日等權、無成本）
# ---------------------------------------------------------------------------
def equal_weight_index():
    r = adj_frame.pct_change(fill_method=None).to_numpy()
    held = np.vstack([np.zeros((1, liquid.shape[1]), dtype=bool), liquid[:-1]])
    v = np.where(held & np.isfinite(r), r, np.nan)
    with np.errstate(invalid="ignore"):
        return np.cumprod(1 + np.nan_to_num(np.nanmean(v, axis=1)))


EW_INDEX = equal_weight_index()


def study(rows, horizons=(5, 10, 20, 40, 60)):
    """b 收盤買進、持有 h 日的平均與中位超額報酬；t 值先按進場月平均，避免同月事件灌大樣本。"""
    out = {}
    for h in horizons:
        rec = []
        for s, b, _ in rows:
            if b + h >= ND:
                continue
            j = COL[s]
            r = adj[b + h, j] / adj[b, j] - 1
            if np.isfinite(r):
                rec.append((b, r - (EW_INDEX[b + h] / EW_INDEX[b] - 1), r - (adj0050[b + h] / adj0050[b] - 1)))
        if not rec:
            continue
        f = pd.DataFrame(rec, columns=["b", "x", "x0050"])
        by_month = f.groupby(pd.Series(days[f.b.values].to_period("M"), index=f.index)).x.mean()
        se = by_month.std(ddof=1) / np.sqrt(len(by_month)) if len(by_month) > 2 else float("nan")
        out[str(h)] = {"n": len(f), "mean_excess": round(float(f.x.mean()), 5), "median_excess": round(float(f.x.median()), 5),
                       "hit_rate": round(float((f.x > 0).mean()), 4),
                       "t_month": round(float(by_month.mean() / se), 2) if np.isfinite(se) and se > 0 else None,
                       "mean_excess_0050": round(float(f.x0050.mean()), 5)}
    return out


def event_curve(rows, lo=-20, hi=60):
    """以 b−1（決策日）為 0 的平均累積超額報酬。"""
    offsets = list(range(lo, hi + 1))
    total, count = np.zeros(len(offsets)), np.zeros(len(offsets))
    for s, b, _ in rows:
        t0 = b - 1
        if t0 + lo < 0 or not np.isfinite(adj[t0, COL[s]]):
            continue
        for k, o in enumerate(offsets):
            t = t0 + o
            if t >= ND:
                break
            v = adj[t, COL[s]] / adj[t0, COL[s]] - EW_INDEX[t] / EW_INDEX[t0]
            if np.isfinite(v):
                total[k] += v
                count[k] += 1
    return {"offset": offsets, "mean": [round(float(x), 5) for x in total / np.maximum(count, 1)], "n": [int(x) for x in count]}


# ---------------------------------------------------------------------------
# 3. 暫停交易事件表
# ---------------------------------------------------------------------------
def halt_events():
    """個股暫停與恢復。p＝暫停前最後收盤，R＝恢復後第一個有收盤價的日子；反應從 p 算到 R，減 0050，不把停牌天算成 0。"""
    halted, resumed, kind = (plain(f"trading_halt:{k}") for k in ("暫停交易", "恢復交易", "有價證券類別"))
    px = close.reindex(columns=STOCKS)
    rows = []
    for s in [c for c in halted.columns if c in COL]:
        starts = halted[s].dropna()
        starts = starts[starts == 1].index
        ends = resumed[s].dropna()
        ends = ends[ends == 1].index
        kinds = kind[s].dropna()
        p_all = px[s].to_numpy()
        for x in starts:
            k = kinds.get(x, "")
            if k not in ("上市證券", "上櫃股票"):
                continue
            later = ends[ends >= x]
            if len(later) == 0:
                continue
            R = days.searchsorted(later[0])
            while R < ND and not np.isfinite(p_all[R]):
                R += 1
            h = days.searchsorted(x)
            traded = np.flatnonzero(np.isfinite(p_all[:h]))
            if R >= ND or len(traded) == 0:
                continue
            p = int(traded[-1])
            j = COL[s]
            rows.append((s, h, p, int(R), float(adj[R, j] / adj[p, j] - adj0050[R] / adj0050[p]), (later[0] - x).days, k))
    frame = pd.DataFrame(rows, columns=["stock_id", "h", "p", "R", "react", "calendar_days", "kind"])
    return frame.sort_values(["R", "stock_id"]).reset_index(drop=True)


HALTS = halt_events()


# ---------------------------------------------------------------------------
# 4. 設計：恢復交易後買進，持有 H 日
# ---------------------------------------------------------------------------
def halt_rows(theta, H, sign, surge):
    a, a20 = amount.to_numpy(), amount_20d.to_numpy()
    rows = []
    for s, p, R, v in zip(HALTS.stock_id, HALTS.p, HALTS.R, HALTS.react):
        j = COL[s]
        if v * sign < theta or not liquid[p, j]:   # 流動性看暫停前最後一天，恢復後 20 日均額還算不出來
            continue
        if surge and not a[R, j] >= 2 * a20[p, j]:
            continue
        rows.append((s, R + 1, R + 1 + H))
    return rows


def up(theta, H, cost):
    return event_sim(halt_rows(theta / 100, H, 1, False), cost)


def down(theta, H, cost):
    return event_sim(halt_rows(theta / 100, H, -1, False), cost)


def up_volume(theta, H, cost):
    return event_sim(halt_rows(theta / 100, H, 1, True), cost)


GRID = dict(x_name="反應門檻 θ（%）", x=[0, 5, 10], y_name="持有 H（交易日）", y=[5, 10, 20], base=(5, 10))
DESIGNS = {
    "S4a": dict(fn=up, title="恢復後上漲買進", **GRID),
    "S4b": dict(fn=down, title="恢復後下跌買進", **GRID),
    "C4": dict(fn=up_volume, title="恢復後上漲＋放量", **GRID),
}


# ---------------------------------------------------------------------------
# 5. 描述統計
# ---------------------------------------------------------------------------
def descriptive():
    high, low = (plain(f"price:{k}").reindex(columns=STOCKS).to_numpy() for k in ("最高價", "最低價"))
    px = close.reindex(columns=STOCKS).to_numpy()
    j = HALTS.stock_id.map(COL).to_numpy()
    p, R = HALTS.p.to_numpy(), HALTS.R.to_numpy()
    raw = px[R, j] / px[p, j] - 1
    next_locked = (R + 1 < ND) & (high[np.minimum(R + 1, ND - 1), j] == low[np.minimum(R + 1, ND - 1), j])
    rows = {k: halt_rows(.05, 10, s, v) for k, s, v in (("up", 1, False), ("down", -1, False), ("up_volume", 1, True))}
    rows["all"] = [(s, R_ + 1, R_ + 11) for s, p_, R_ in zip(HALTS.stock_id, HALTS.p, HALTS.R) if liquid[p_, COL[s]]]
    edges = np.arange(-0.105, 0.1051, 0.01)
    return {"events": len(HALTS), "stocks": int(HALTS.stock_id.nunique()),
            "events_by_year": pd.Series(days[R].year).value_counts().sort_index().to_dict(),
            "up_by_year": pd.Series(days[R[HALTS.react.to_numpy() > 0]].year).value_counts().sort_index().to_dict(),
            "calendar_days": HALTS.calendar_days.value_counts().sort_index().to_dict(),
            "one_trading_day": int((R - p == 2).sum()), "by_market": HALTS.kind.value_counts().to_dict(),
            "liquid_share": round(float(liquid[p, j].mean()), 4),
            "resume_day": {"mean": round(float(raw.mean()), 5), "median": round(float(np.median(raw)), 5),
                           "excess_mean": round(float(HALTS.react.mean()), 5), "max": round(float(raw.max()), 5), "min": round(float(raw.min()), 5),
                           "limit_up": int((raw >= .094).sum()), "limit_down": int((raw <= -.094).sum()),
                           "one_price": int((high[R, j] == low[R, j]).sum()), "next_day_one_price": int(next_locked.sum()),
                           "up": int((raw > 0).sum()), "down": int((raw < 0).sum()),
                           "hist_edges": [round(float(x), 3) for x in edges], "hist": np.histogram(raw, edges)[0].tolist()},
            "reaction_counts": {f"{k}{t}": int((HALTS.react * s >= t / 100).sum()) for k, s in (("up", 1), ("down", -1)) for t in (0, 5, 10)},
            "study": {k: study(v) for k, v in rows.items()},
            "curves": {k: event_curve(v) for k, v in rows.items()}}


# ---------------------------------------------------------------------------
# 9. 執行
# ---------------------------------------------------------------------------
def main():
    equal_weight = equal_weight_curve()
    results, reports = {}, {}
    for key in DESIGNS:
        results[key], reports[key] = validate(key, equal_weight)
    metrics = {"run_date": str(datetime.date.today()), "finlab_version": finlab.__version__,
               "data_end": str(bench.index[-1].date()), "designs": results, "descriptive": descriptive()}
    with open("metrics.json", "w", encoding="utf-8") as f:
        json.dump(metrics, f, ensure_ascii=False, indent=1)
    frame = {}
    for key, report in reports.items():
        curve = active(report.creturn)
        frame[f"{key} {DESIGNS[key]['title']}"] = curve / curve.iloc[0]
    frame = pd.DataFrame(frame)
    frame["0050含息"] = bench.reindex(frame.index)
    frame["等權流動股"] = equal_weight.reindex(frame.index)
    frame.index.name = "date"
    frame.round(6).to_csv("data.csv")
    print("完成：metrics.json、data.csv")


if __name__ == "__main__":
    main()
