#!/usr/bin/env python3
"""
涨停缩量回调突破 - 网格搜索 v3 (缩量OR定义 + 合并卖出止盈OR回撤 + 最大持有30天)
买入: T0涨停 -> ta∈(T0+1,T0+sw] 缩量(基准zt OR maN, 系数vs, 均量窗n) 且收盘<P0 -> tb∈(ta+1,T0+bu] 收盘突破P0×(1+bs) -> 次日开盘买
卖出: 盘中回撤dd(基准=昨日最高收盘) -> 收盘止盈profit -> 超时MAX_HOLD=30天
阶段1: 卖出网格 止盈{10,15,20,25,30}×回撤{5,7,10,12,15}=25
阶段2a: 缩量网格 基准{zt,maN,or}×n{3,5,8,10}×系数{0.4-0.8}×窗口{3,5,8}=135
阶段2b: 突破网格 上限{6,8,10}×强度{0-3%}×量能{不限,1.5,2,3}=84
"""
import pandas as pd
import numpy as np
import os, sys, json, time
from multiprocessing import Pool

BASE = '/Users/xpresso/zt_app/backtest_zt_full'
START_DATE = '20210101'

LIMIT_UP = {'main': 9.9, 'gem': 19.9, 'star': 19.9, 'bj': 29.9}
MAX_HOLD = 30
INIT_CAP = 1_000_000.0
POS_SHARE = 0.20
MAX_POS = 5
COST_BUY = 0.00025
COST_SELL = 0.00025 + 0.001
SLIP = 0.001

DEF_BUY = {'vol_base': 'zt', 'ma_n': 5, 'vol_shrink': 0.6, 'shrink_win': 5, 'break_up': 8, 'break_strength': 0.0, 'zt_vol_filter': 0}
GRID_PROFIT = [0.10, 0.15, 0.20, 0.25, 0.30, 0.99]  # 0.99≈∞纯回撤(无止盈)
GRID_DD = [0.05, 0.07, 0.10, 0.12, 0.15, 0.99]  # 0.99≈∞纯止盈(无回撤)
GRID_VB = ['zt', 'maN', 'or']
GRID_MA_N = [3, 5, 8, 10]
GRID_VS = [0.4, 0.5, 0.6, 0.7, 0.8]
GRID_SW = [3, 5, 8]
GRID_BU = [6, 8, 10]
GRID_BS = [0.0, 0.005, 0.01, 0.015, 0.02, 0.025, 0.03]
GRID_ZTV = [0, 1.5, 2.0, 3.0]

def board_type(code):
    if code.startswith(('60', '00')): return 'main'
    if code.startswith(('300', '301')): return 'gem'
    if code.startswith('688'): return 'star'
    return 'bj'

def load_pool(pool):
    members = pd.read_pickle(f'{BASE}/{pool}_members.pkl')
    dailies = {}
    ddir = f'{BASE}/daily_{pool}'
    for f in os.listdir(ddir):
        if f.endswith('.csv'):
            code = f[:-4]
            d = pd.read_csv(f'{ddir}/{f}', dtype={'trade_date': str})
            if len(d) > 0:
                _derive_split_prices(d)   # 派生 qfq_/hfq_（仅当文件存 raw+factor）
                dailies[code] = d
    return members, dailies

def _derive_split_prices(d, latest_factor=None):
    """若文件仅存 raw_* + adj_factor，在内存派生 qfq_/hfq_ 列供判定/K线用。
    若无 factor 或 raw（如研究侧旧数据），保持原列不动（fallback 读文件既有列）。
    派生恒等式：hfq = raw×factor（后复权）；qfq = raw×factor/latest（前复权，展示时算）。

    ⚠️ latest 必须是**该股整段末行因子**（切片归一基准会随窗口漂移 → 标注/蜡烛错位）：
    优先序 ①显式传入 latest_factor ②帧内已盖的 `adj_factor_latest` 列(整段常量)
    ③退化为本帧末行因子——**仅当调用方传的是完整序列才正确**。派生后会在帧上盖
    `adj_factor_latest` 常量列，故此后对任意窗口切片再调用都稳定（切片不变性）。
    """
    if 'adj_factor' not in d.columns or 'raw_close' not in d.columns:
        return d
    fac = d['adj_factor'].astype(float)
    if latest_factor is not None:
        lt = float(latest_factor)
    elif 'adj_factor_latest' in d.columns:
        lt = float(d['adj_factor_latest'].iloc[0])
    else:
        lt = float(fac.iloc[-1])       # ⚠️ 切片调用会得到错误的归一基准
    if lt <= 0:
        return d
    for c in ['open', 'high', 'low', 'close']:
        rc = d['raw_' + c].astype(float)
        d['hfq_' + c] = (rc * fac).round(4)
        d['qfq_' + c] = (rc * fac / lt).round(4)
    d['adj_factor_latest'] = lt        # 盖章: 整段常量, 令切片不变性成立
    return d

# === 口径助手(定稿 2026-09-12: 判定=后复权 hfq_*, 撮合/估值=原始价 raw_*) ===
def hfq_col(d, col):
    """判定取价(缩量回调/突破/止盈止损/链特征): 后复权, 免疫除权跳变且历史可复现"""
    return d['hfq_' + col]

def raw_col(d, col):
    """撮合/估值取价(成交价/股数/持仓估值): 原始价, 与实盘下单口径一致"""
    return d['raw_' + col]

def kline_col(r, col):
    """K线/图表展示取价: 优先前复权 qfq_(派生, 除权日平滑无假跳空), 无则该列取原始价。
    定稿(PRICE_SCHEMA): raw=撮合估值, hfq=判定, qfq=图表展示。"""
    c = 'qfq_' + col
    return r[c] if c in r.index else r['raw_' + col]

def active_members(members, date_str):
    ym = date_str[:6]
    best = None
    for m in sorted(members.keys()):
        if m[:6] <= ym:
            best = m
    return members.get(best, set()) if best else set()

def signal_detect(d, code, bp):
    d = d.sort_values('trade_date').reset_index(drop=True)
    n = len(d)
    if n < 12:
        return []
    pct = d['pct_chg'].astype(float).values
    close = d['hfq_close'].astype(float).values     # 判定=后复权
    open_ = d['hfq_open'].astype(float).values      # 判定=后复权
    vol = d['vol'].astype(float).values
    dates = d['trade_date'].values
    bt = LIMIT_UP[board_type(code)]
    vb, vs, bs = bp['vol_base'], bp['vol_shrink'], bp['break_strength']
    sw, bu, mn = int(bp['shrink_win']), int(bp['break_up']), int(bp['ma_n'])
    ztv = bp.get('t0_vol_max', bp.get('zt_vol_filter', 0))  # T0天量上限(原zt_vol_filter)
    t0min = bp.get('t0_vol_min', 0.0)   # T0放量下限: T0量>=前10日均量*X
    taddmax = bp.get('ta_dd_max', 0.0)  # 回调最大深度: ta收盘>=T0收盘*(1-X), 0=不限制
    t0openmax = bp.get('t0_open_max', 0.0)  # T0开盘涨幅/涨停幅度 <= X (排除一字/秒板), 0=不限制
    t0prev = bp.get('t0_prev_limit', 0)     # 连板限制: T0前第L日涨停则排除(L=1仅首板, L=2最多二板), 0=不限制
    talowfloor = bp.get('ta_low_floor', 0.0)  # 铁底: ta收盘 >= T0开盘*X, 0=不限制
    tbt0_min = bp.get('tb_t0_min', 0.0) # tb量>=T0量*X
    tbt0_max = bp.get('tb_t0_max', 0.0) # tb量<=T0量*X
    tbmin = bp.get('tb_vol_min', 0.0)   # tb量>=回调期均量*X
    tbmax = bp.get('tb_vol_max', 0.0)   # tb量<=回调期均量*X
    out = []
    for t0 in np.where(pct >= bt - 0.05)[0]:
        if t0 + 2 >= n:
            continue
        v0 = vol[t0]
        if v0 <= 0:
            continue
        ma10z = vol[max(0, t0 - 10):t0].mean() if t0 >= 10 else vol[:t0].mean()
        if t0min > 0:
            if ma10z <= 0 or v0 < ma10z * t0min:
                continue
        if ztv > 0:
            if ma10z > 0 and v0 >= ma10z * ztv:
                continue
        if t0openmax > 0:
            if t0 >= 1:
                prev_c = close[t0 - 1]
                if prev_c > 0:
                    open_pct = open_[t0] / prev_c - 1
                    if open_pct / bt > t0openmax:
                        continue
        if t0prev > 0 and t0 >= t0prev:
            if pct[t0 - t0prev] >= bt - 0.05:
                continue
        p0 = close[t0]
        p0_floor = p0 * (1 - taddmax) if taddmax > 0 else None
        # ta: t0+1~t0+sw窗口内所有合格缩量日(且回调深度不破位), 选收盘价最低者(洗盘最深)
        ta = None
        ta_min_close = None
        for j in range(t0 + 1, min(t0 + sw + 1, n)):
            shrink_ok = False
            if vb in ('zt', 'or') and vol[j] < v0 * vs:
                shrink_ok = True
            if vb in ('maN', 'or'):
                m = vol[max(0, j - mn):j].mean() if j >= mn else vol[:j].mean()
                if m > 0 and vol[j] < m * vs:
                    shrink_ok = True
            if shrink_ok and close[j] < p0:
                if p0_floor is not None and close[j] < p0_floor:
                    continue  # 回调过深: 破位, 该日不合格
            if shrink_ok and close[j] < p0:
                if talowfloor > 0 and open_[t0] > 0 and close[j] < open_[t0] * talowfloor:
                    continue  # 跌破铁底(T0开盘): 放弃该日
                if ta_min_close is None or close[j] < ta_min_close:
                    ta_min_close = close[j]
                    ta = j
        if ta is None:
            continue
        for k in range(ta + 1, min(t0 + bu + 1, n)):
            if close[k] > p0 * (1 + bs):
                # tb量能检查(基准: 回调期均量=vol[t0+1..tb-1])
                ok = True
                if tbt0_min > 0 and v0 > 0 and vol[k] < v0 * tbt0_min:
                    ok = False
                if ok and tbt0_max > 0 and v0 > 0 and vol[k] > v0 * tbt0_max:
                    ok = False
                if ok and (tbmin > 0 or tbmax > 0):
                    cb = vol[t0 + 1:k].mean() if k > t0 + 1 else vol[k]
                    if cb > 0:
                        if tbmin > 0 and vol[k] < cb * tbmin:
                            ok = False
                        if ok and tbmax > 0 and vol[k] > cb * tbmax:
                            ok = False
                    else:
                        ok = False
                if not ok:
                    continue  # 量能不满足: 继续找下一个突破日
                if k + 1 < n:
                    out.append((dates[k + 1], dates[t0]))
                break
    return out

def build_sig(dailies, bp):
    sig = {}
    for code, d in dailies.items():
        bs = signal_detect(d, code, bp)
        if bs:
            sig[code] = bs
    return sig

def run_backtest(pool, members, dailies, sig, profit, dd, start=START_DATE, mode='merge', stop=None, hold=MAX_HOLD):
    """卖出:
    mode='merge': 盘中回撤dd(trailing) -> 收盘止盈profit -> 超时 (现有)
    mode='A'(止盈策略): 收盘盈利>=profit | 盘中<=买入价*(1-stop或dd) | 超时
    mode='B'(回撤策略): 盘中<=最高价*(1-profit) | 盘中<=买入价*(1-stop或dd) | 超时
    mode='C'(三规则): 盘中回撤dd(trailing) | 盘中<=买入价*(1-stop) | 收盘止盈profit | 超时
    """
    all_dates = sorted(set().union(*[set(d['trade_date']) for d in dailies.values()]))
    all_dates = [x for x in all_dates if x >= start]
    date_set = set(all_dates)
    sig_by_date = {}
    for code, bs in sig.items():
        for (bd, t0d) in bs:
            if bd in date_set:
                sig_by_date.setdefault(bd, []).append((code, t0d))
    cash = INIT_CAP
    positions = []
    trades = []
    for date in all_dates:
        active = active_members(members, date)
        for pos in positions[:]:
            d = dailies[pos['code']]
            rows = d[d['trade_date'] == date]
            if len(rows) == 0:
                continue
            r = rows.iloc[0]
            cur = float(r['hfq_close'])
            pos['hold_n'] += 1
            sell = None
            if mode == 'A':
                if cur / pos['buy_j'] - 1 >= profit:
                    sell = 'profit'
                elif float(r['hfq_low']) <= pos['buy_j'] * (1 - (stop if stop is not None else dd)):
                    sell = 'stop'
                elif pos['hold_n'] >= hold:
                    sell = 'timeout'
            elif mode == 'B':
                if float(r['hfq_low']) <= pos['high_since'] * (1 - profit):
                    sell = 'trail'
                elif float(r['hfq_low']) <= pos['buy_j'] * (1 - (stop if stop is not None else dd)):
                    sell = 'stop'
                elif pos['hold_n'] >= hold:
                    sell = 'timeout'
            elif mode == 'C':
                if float(r['hfq_low']) <= pos['high_since'] * (1 - dd):
                    sell = 'dd'
                elif stop is not None and float(r['hfq_low']) <= pos['buy_j'] * (1 - stop):
                    sell = 'stop'
                elif cur / pos['buy_j'] - 1 >= profit:
                    sell = 'profit'
                elif pos['hold_n'] >= hold:
                    sell = 'timeout'
            else:
                if float(r['hfq_low']) <= pos['high_since'] * (1 - dd):
                    sell = 'dd'
                elif cur / pos['buy_j'] - 1 >= profit:
                    sell = 'profit'
                elif pos['hold_n'] >= hold:
                    sell = 'timeout'
            pos['high_since'] = max(pos['high_since'], cur)
            if sell:
                net_ret = (cur / pos['buy_j'] - 1) - (COST_BUY + COST_SELL + 2 * SLIP)   # 判定口径收益率
                cash += pos['amount'] * (1 + net_ret)
                fill = float(r['raw_close'])          # 成交价=原始价
                trades.append({'code': pos['code'], 'buy_date': pos['buy_date'], 'sell_date': date,
                               'buy_price': pos['fill_px'], 'sell_price': fill,
                               'ret_net': net_ret, 'hold_days': pos['hold_n'], 'reason': sell})
                positions.remove(pos)
        if len(positions) < MAX_POS:
            cands = sorted(sig_by_date.get(date, []), key=lambda x: x[1])
            for code, t0d in cands:
                if len(positions) >= MAX_POS:
                    break
                if code not in active:
                    continue
                if any(p['code'] == code for p in positions):
                    continue
                d = dailies[code]
                rows = d[d['trade_date'] == date]
                if len(rows) == 0:
                    continue
                fill_px = float(rows.iloc[0]['raw_open'])    # 成交价/股数=原始价
                buy_j = float(rows.iloc[0]['hfq_open'])      # 判定入场=后复权
                if fill_px <= 0 or buy_j <= 0:
                    continue
                held_mv = 0.0
                for p in positions:
                    dd_ = dailies[p['code']]
                    rr = dd_[dd_['trade_date'] == date]
                    held_mv += p['qty'] * float(rr.iloc[0]['raw_close']) if len(rr) else p['amount']
                amount = min((cash + held_mv) * POS_SHARE, cash * 0.99)
                if amount <= 0:
                    break
                qty = amount / fill_px
                positions.append({'code': code, 'buy_date': date, 'buy_price': fill_px * (1 + SLIP),
                                  'fill_px': fill_px * (1 + SLIP), 'buy_j': buy_j * (1 + SLIP),
                                  'amount': qty * fill_px,
                                  'qty': qty, 'high_since': buy_j, 'hold_n': 0})
                cash -= amount * (1 + COST_BUY + SLIP)
    return trades

def stats(trades):
    n = len(trades)
    if n == 0:
        return {'trades': 0, 'win_rate': 0, 'avg_ret': 0}
    rets = [t['ret_net'] for t in trades]
    wins = [r for r in rets if r > 0]
    return {'trades': n, 'win_rate': len(wins) / n, 'avg_ret': float(np.mean(rets))}

def run_pool_stage(pool, stage, sell=None):
    members, dailies = load_pool(pool)
    results = []
    t0 = time.time()
    if stage == 1:
        sig = build_sig(dailies, DEF_BUY)
        nsig = sum(len(v) for v in sig.values())
        for pr in GRID_PROFIT:
            for dd_ in GRID_DD:
                results.append({'pool': pool, 'stage': 1, 'profit': pr, 'dd': dd_, 'nsig': nsig,
                                **stats(run_backtest(pool, members, dailies, sig, pr, dd_))})
    elif stage == 2:
        # 2a 缩量网格(固定sell)
        for vb in GRID_VB:
            for mn in GRID_MA_N:
                for vs in GRID_VS:
                    for sw in GRID_SW:
                        bp = {'vol_base': vb, 'ma_n': mn, 'vol_shrink': vs, 'shrink_win': sw,
                              'break_up': 8, 'break_strength': 0.0, 'zt_vol_filter': 0}
                        sig = build_sig(dailies, bp)
                        nsig = sum(len(v) for v in sig.values())
                        s = stats(run_backtest(pool, members, dailies, sig, sell[0], sell[1]))
                        results.append({'pool': pool, 'stage': 2, 'vol_base': vb, 'ma_n': mn, 'vol_shrink': vs,
                                        'shrink_win': sw, 'break_up': 8, 'break_strength': 0.0, 'zt_vol_filter': 0,
                                        'profit': sell[0], 'dd': sell[1], 'nsig': nsig, **s})
    elif stage == 3:
        # 2b 突破网格(固定sell + 最优缩量)
        best2 = json.load(open(f'{BASE}/best_stage2.json'))
        for bu in GRID_BU:
            for bs in GRID_BS:
                for zt in GRID_ZTV:
                    bp = {'vol_base': best2[pool]['vol_base'], 'ma_n': int(best2[pool]['ma_n']),
                          'vol_shrink': best2[pool]['vol_shrink'], 'shrink_win': int(best2[pool]['shrink_win']),
                          'break_up': bu, 'break_strength': bs, 'zt_vol_filter': zt}
                    sig = build_sig(dailies, bp)
                    nsig = sum(len(v) for v in sig.values())
                    s = stats(run_backtest(pool, members, dailies, sig, sell[0], sell[1]))
                    results.append({'pool': pool, 'stage': 3, 'vol_base': best2[pool]['vol_base'], 'ma_n': int(best2[pool]['ma_n']),
                                    'vol_shrink': best2[pool]['vol_shrink'], 'shrink_win': int(best2[pool]['shrink_win']),
                                    'break_up': bu, 'break_strength': bs, 'zt_vol_filter': zt,
                                    'profit': sell[0], 'dd': sell[1], 'nsig': nsig, **s})
    print(f"{pool} stage{stage} 完成: {len(results)}条, 耗时{(time.time()-t0)/60:.1f}分钟", flush=True)
    return results

def main():
    pools = ['hs300', 'zz500', 'zz1000']
    stage = int(sys.argv[1]) if len(sys.argv) > 1 else 1
    t0 = time.time()
    if stage == 1:
        with Pool(3) as p:
            all_results = []
            for r in p.starmap(run_pool_stage, [(pool, 1) for pool in pools]):
                all_results.extend(r)
        json.dump(all_results, open(f'{BASE}/grid_results_stage1.json', 'w'), ensure_ascii=False, indent=1, default=str)
        df = pd.DataFrame(all_results)
        best1 = {}
        for pool in pools:
            sub = df[df['pool'] == pool]
            ok = sub[sub['trades'] >= 20]
            if len(ok) == 0:
                ok = sub
            r = ok.sort_values('avg_ret', ascending=False).iloc[0]
            best1[pool] = {'profit': float(r['profit']), 'dd': float(r['dd']), 'avg_ret': float(r['avg_ret']), 'trades': int(r['trades'])}
            print(f"best1 {pool}: 止盈{r['profit']} 回撤{r['dd']} 均收{r['avg_ret']*100:.2f}% {r['trades']}笔")
        json.dump(best1, open(f'{BASE}/best_stage1.json', 'w'), ensure_ascii=False, indent=1)
    else:
        best1 = json.load(open(f'{BASE}/best_stage1.json'))
        sell = [(best1[pool]['profit'], best1[pool]['dd']) for pool in pools]
        with Pool(3) as p:
            all_results = []
            for r in p.starmap(run_pool_stage, [(pools[i], stage, sell[i]) for i in range(3)]):
                all_results.extend(r)
        fname = f'grid_results_stage{stage}.json'
        json.dump(all_results, open(f'{BASE}/{fname}', 'w'), ensure_ascii=False, indent=1, default=str)
        if stage == 2:
            df = pd.DataFrame(all_results)
            best2 = {}
            for i, pool in enumerate(pools):
                sub = df[df['pool'] == pool]
                ok = sub[sub['trades'] >= 20]
                if len(ok) == 0:
                    ok = sub
                r = ok.sort_values('avg_ret', ascending=False).iloc[0]
                best2[pool] = {'vol_base': r['vol_base'], 'ma_n': float(r['ma_n']), 'vol_shrink': float(r['vol_shrink']),
                               'shrink_win': float(r['shrink_win']), 'avg_ret': float(r['avg_ret']), 'trades': int(r['trades'])}
                print(f"best2 {pool}: 基准{r['vol_base']}/n{int(r['ma_n'])}/系数{r['vol_shrink']}/窗{r['shrink_win']} 均收{r['avg_ret']*100:.2f}% {r['trades']}笔")
            json.dump(best2, open(f'{BASE}/best_stage2.json', 'w'), ensure_ascii=False, indent=1)
        print(f"stage{stage} 完成: {len(all_results)}条, 总耗时{(time.time()-t0)/60:.1f}分钟")

if __name__ == '__main__':
    main()