#!/usr/bin/env python3
"""step49_final_v5.py — 最终模拟 v5（阴跌中段判定 + 限流）
在 step37 基础上：
- 信号: 预测信号 + 大盘回撤<=-5%
- 阴跌中段判定（step48验证）: 连跌>=2 + ADX<=55 + 大盘RSI>=33 → 跳过信号
- 单日限流5 + 10槽 + 15%
对比: 年化/回撤 vs step37(v3) 和 step39(限流)
"""
import numpy as np
import pandas as pd
import sqlite3
import sys
sys.path.insert(0, "/tmp")
from sr_calculator import SRCalculator

print("=== 加载 ===", flush=True)
panel = pd.read_pickle("/tmp/panel_12d.pkl")
panel = panel.sort_values(["code", "date"]).reset_index(drop=True)
panel["_key"] = panel["code"] + "_" + panel["date"]
pos_map = {k: i for i, k in enumerate(panel["_key"])}
dates = sorted(panel["date"].unique())

# 大盘指标
conn = sqlite3.connect("file:/home/hmo/MoFin/data/mofin.db?mode=ro", uri=True)
idx_df = pd.read_sql("SELECT date, close, high FROM stock_daily WHERE code='sh000001' ORDER BY date", conn)
idx_df["date"] = idx_df["date"].astype(str)
idx_df["hi60"] = idx_df["high"].rolling(60).max()
idx_df["mkt_dd60"] = (idx_df["close"] / idx_df["hi60"] - 1) * 100
close_arr = idx_df["close"].values
down_days, streak, prev = [], 0, None
for c in close_arr:
    streak = streak + 1 if (prev is not None and c < prev) else 0
    down_days.append(streak)
    prev = c
idx_df["mkt_down_days"] = down_days
dd_map = dict(zip(idx_df["date"], idx_df["mkt_dd60"]))
down_map = dict(zip(idx_df["date"], idx_df["mkt_down_days"]))
panel["mkt_dd60"] = panel["date"].map(dd_map)
panel["mkt_down_days"] = panel["date"].map(down_map)

# 信号
sig_cond = (
    (panel["mkt_rsi"] < 50) & (panel["mcap_q"] < 0.2) & (panel["pe_q"] < 0.2) &
    (panel["news3"] >= 1) & (panel["sec_ret20"] < 0) & (panel["bias60"] < -20) &
    (panel["mkt_dd60"] <= -5)
)
# 阴跌中段判定：跳过（step48数据验证）
yin_die = (panel["mkt_down_days"] >= 2) & (panel["mkt_adx"] <= 55) & (panel["mkt_rsi"] >= 33)
final_cond = sig_cond & ~yin_die

print("基础信号:", sig_cond.sum(), "| 阴跌跳过:", (sig_cond & yin_die).sum(), "| 最终:", final_cond.sum(), flush=True)

def build_cand(cond):
    cand = panel[cond][["code", "date"]].copy()
    cand = cand.sort_values(["code", "date"])
    cand["prev"] = cand.groupby("code")["date"].shift(1)
    cand["gap"] = (pd.to_datetime(cand["date"]) - pd.to_datetime(cand["prev"])).dt.days
    cand = cand[(cand["prev"].isna()) | (cand["gap"] > 30)]
    return cand

def prep(sigdf):
    codes = sigdf["code"].unique().tolist()
    ph = ",".join("?" * len(codes))
    df = pd.read_sql("SELECT code, date, close, high, low FROM stock_daily WHERE code IN ({}) ORDER BY code, date".format(ph), conn, params=codes)
    df["date"] = df["date"].astype(str)
    df["code"] = df["code"].astype(str).str.zfill(6)
    df = df.sort_values(["code", "date"]).reset_index(drop=True)
    df["_key"] = df["code"] + "_" + df["date"]
    dpos = {k: i for i, k in enumerate(df["_key"])}
    sigdf["kidx"] = (sigdf["code"] + "_" + sigdf["date"]).map(dpos)
    sigdf = sigdf.dropna(subset=["kidx"]).copy()
    sigdf["kidx"] = sigdf["kidx"].astype(int)
    sr = SRCalculator()
    out = []
    for r in sigdf.itertuples():
        bars_code = sr.get_bars(r.code)
        d_idx = bars_code.index[bars_code["date"] == r.date]
        if len(d_idx) == 0:
            continue
        sr_full = sr.sr_full(r.code, d_idx[0])
        pv = sr_full["pivot"]
        chip = sr_full["chip"]
        if not pv:
            continue
        sig_close = df["close"].values[r.kidx]
        support = pv["s2"]
        if chip and chip["chip_ss"] < sig_close:
            support = max(pv["s2"], chip["chip_ss"])
        resist = pv["r2"]
        if chip and chip["chip_sr"] > sig_close:
            resist = min(pv["r2"], chip["chip_sr"])
        if support >= sig_close or resist <= sig_close:
            continue
        out.append({"code": r.code, "date": r.date, "kidx": r.kidx,
                    "sig_close": sig_close, "support": support, "resist": resist})
    return pd.DataFrame(out), df

def run_sim(sigdf, df, max_daily=5, slots=10, hold=40, pos_frac=0.15, stop_buf=0.05):
    dpos = {k: i for i, k in enumerate(df["_key"])}
    closes = df["close"].values
    highs = df["high"].values
    lows = df["low"].values
    sig_by_date = {}
    for r in sigdf.itertuples():
        sig_by_date.setdefault(r.date, []).append(r)
    INIT_CAP = 1_000_000
    positions = {}
    cash = INIT_CAP
    navs = []
    for di, d in enumerate(dates):
        for code in list(positions.keys()):
            pos = positions[code]
            k = dpos.get(code + "_" + d)
            if k is None:
                continue
            hi, lo, cl = highs[k], lows[k], closes[k]
            exited = False
            if hi >= pos["tp"]:
                sell_qty = pos["qty"] // 2
                if sell_qty > 0:
                    cash += sell_qty * pos["tp"]
                    pos["qty"] -= sell_qty
                if pos["qty"] <= 0:
                    del positions[code]
                    exited = True
            if not exited and lo <= pos["stop"]:
                cash += pos["qty"] * lo
                del positions[code]
                exited = True
            if not exited and di - pos["entry_di"] >= hold:
                cash += pos["qty"] * cl
                del positions[code]
        if d in sig_by_date:
            bought = 0
            for r in sig_by_date[d]:
                if bought >= max_daily or len(positions) >= slots:
                    break
                if r.code in positions:
                    continue
                cur_nav = cash
                for c, p in positions.items():
                    k = dpos.get(c + "_" + d)
                    px = closes[k] if k is not None else p["avg_cost"]
                    cur_nav += p["qty"] * px
                pos_val = cur_nav * pos_frac
                price = r.sig_close
                if price <= 0:
                    continue
                qty = int(pos_val / price)
                if qty <= 0 or qty * price > cash:
                    continue
                cash -= qty * price
                positions[r.code] = {
                    "entry_di": di, "qty": qty, "avg_cost": price,
                    "stop": r.support * (1 - stop_buf), "tp": r.resist,
                }
                bought += 1
        nav = cash
        for c, p in positions.items():
            k = dpos.get(c + "_" + d)
            px = closes[k] if k is not None else p["avg_cost"]
            nav += p["qty"] * px
        navs.append(nav)
    nav_arr = np.array(navs)
    final = nav_arr[-1]
    years = len(navs) / 250
    cagr = ((final / INIT_CAP) ** (1 / years) - 1) * 100 if final > 0 else -100
    cummax = np.maximum.accumulate(nav_arr)
    dd = ((nav_arr - cummax) / cummax).min() * 100
    return cagr, dd, nav_arr

print("\n=== 基准（无阴跌判定，限流5）===", flush=True)
c0, d0, _ = run_sim(*prep(build_cand(sig_cond)))
print("基准: 年化={:.2f}% 回撤={:.1f}%".format(c0, d0), flush=True)

print("\n=== 阴跌判定 + 限流5 ===", flush=True)
s1, df1 = prep(build_cand(final_cond))
c1, d1, nav1 = run_sim(s1, df1)
print("阴跌判定: 年化={:.2f}% 回撤={:.1f}% 信号={}".format(c1, d1, len(s1)), flush=True)

print("\n=== 阴跌判定 限流参数 ===", flush=True)
for md in [3, 5, 8, 99]:
    c, d, _ = run_sim(s1, df1, max_daily=md)
    print("限{}: 年化={:.2f}% 回撤={:.1f}%".format(md, c, d), flush=True)

print("\n=== 分年净值（阴跌判定+限流5）===", flush=True)
nav_df = pd.DataFrame({"date": dates, "nav": nav1})
nav_df["year"] = nav_df["date"].str[:4]
yr = nav_df.groupby("year")["nav"].last()
for y, v in yr.items():
    print("  {}: {:.0f} ({}%)".format(y, v, (v/1000000-1)*100))

print("\n=== 完成 ===", flush=True)