#!/usr/bin/env python3
"""
每日TB信号扫描: 找四池中今天(tb日)满足突破条件且明天可买入的股票, 计算11特征组合强度
输出: 按强度排序的候选股票列表
"""
import pandas as pd
import numpy as np
import json, sys
sys.path.insert(0, '/Users/xpresso/zt_app/code')
import grid_or as g

BASE = g.BASE
LIMIT_UP = {'main': 9.9, 'gem': 19.9, 'star': 19.9, 'bj': 29.9}
FCOLS = ['f1_t0_pct', 'f2_t0_vol_ratio', 'f3_ta_shrink', 'f4_ta_drawdown', 'f5_ta_gap',
         'f6_tb_strength', 'f7_tb_gap', 'f8_chain_len', 'f9_tb_vol', 'f10_tb_ma5_slope', 'f11_tb_ma_align',
         'f12_crowd_pool', 'f13_pool_zt', 'f14_pool_avgret']
PN = {'hs300': '沪深300', 'zz500': '中证500', 'zz1000': '中证1000', 'zz2000': '中证2000'}
COMBO = {
    'hs300': {'vol_base': 'or', 'ma_n': 2, 'vol_shrink': 1.0, 'shrink_win': 24, 'break_up': 15, 'break_strength': 0.0,
              't0_vol_min': 0.0, 't0_vol_max': 0.0, 't0_open_max': 0.5, 't0_prev_limit': 0,
              'ta_dd_max': 0.3, 'ta_low_floor': 0.0,
              'tb_t0_min': 0.0, 'tb_t0_max': 0.0, 'tb_vol_min': 0.5, 'tb_vol_max': 4.0},
    'zz500': {'vol_base': 'or', 'ma_n': 2, 'vol_shrink': 1.1, 'shrink_win': 24, 'break_up': 12, 'break_strength': 0.0,
              't0_vol_min': 0.0, 't0_vol_max': 0.0, 't0_open_max': 0.5, 't0_prev_limit': 0,
              'ta_dd_max': 0.2, 'ta_low_floor': 0.0,
              'tb_t0_min': 0.0, 'tb_t0_max': 0.0, 'tb_vol_min': 0.5, 'tb_vol_max': 3.0},
    'zz1000': {'vol_base': 'or', 'ma_n': 2, 'vol_shrink': 1.1, 'shrink_win': 24, 'break_up': 20, 'break_strength': 0.0,
               't0_vol_min': 0.0, 't0_vol_max': 0.0, 't0_open_max': 0.5, 't0_prev_limit': 0,
               'ta_dd_max': 0.4, 'ta_low_floor': 0.0,
               'tb_t0_min': 0.0, 'tb_t0_max': 0.0, 'tb_vol_min': 0.0, 'tb_vol_max': 5.0},
    'zz2000': {'vol_base': 'or', 'ma_n': 2, 'vol_shrink': 1.1, 'shrink_win': 24, 'break_up': 20, 'break_strength': 0.0,
               't0_vol_min': 0.0, 't0_vol_max': 0.0, 't0_open_max': 0.0, 't0_prev_limit': 0,
               'ta_dd_max': 0.4, 'ta_low_floor': 0.0,
               'tb_t0_min': 0.0, 'tb_t0_max': 0.0, 'tb_vol_min': 0.0, 'tb_vol_max': 4.0},
}

def board_type(code):
    if code.startswith(('60', '00')):
        return 'main'
    if code.startswith(('30', '68')):
        return 'gem'
    return 'bj'

def crowd_denom(bdf, pool, bd):
    """池日均信号数分母: 仅用 buy_date < bd 的历史行(2026-09-12)
    防「重跑自称重」漂移——旧式用全库(含当日行), 同一日重跑会得到不同 f12。"""
    h = bdf[(bdf['pool'] == pool) & (bdf['buy_date'].astype(str) < str(bd))]
    if not len(h):
        return 1.0
    return max(1.0, h.groupby('buy_date').size().mean())


def find_tb_today(d, code, bp, target):
    """找tb==target(今天)的信号链, 返回特征dict(含明天buy_date)"""
    d = d.sort_values('trade_date').reset_index(drop=True)
    n = len(d)
    if n < 25:
        return None
    pct = d['pct_chg'].astype(float).values
    close = d['hfq_close'].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))
    t0min = bp.get('t0_vol_min', 0.0)
    taddmax = bp.get('ta_dd_max', 0.0)
    t0openmax = bp.get('t0_open_max', 0.0)
    t0prev = bp.get('t0_prev_limit', 0)
    talowfloor = bp.get('ta_low_floor', 0.0)
    tbt0_min = bp.get('tb_t0_min', 0.0)
    tbt0_max = bp.get('tb_t0_max', 0.0)
    tbmin = bp.get('tb_vol_min', 0.0)
    tbmax = bp.get('tb_vol_max', 0.0)
    ma5 = pd.Series(close).rolling(5).mean().values
    ma10 = pd.Series(close).rolling(10).mean().values
    ma20 = pd.Series(close).rolling(20).mean().values
    res = []
    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
        p0 = close[t0]
        p0_floor = p0 * (1 - taddmax) if taddmax > 0 else None
        # ta: 窗口内所有合格缩量日(且深度不破位), 选收盘价最低(与引擎signal_detect一致)
        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 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):
                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 dates[k] == target:
                    # 突破日=今日: 明日开盘买入(次日数据可能未发布, buy_date=None=待确认)
                    buy_date = dates[k + 1] if k + 1 < n else None
                    m10 = vol[max(0, k - 10):k].mean() if k >= 10 else vol[:k].mean()
                    res.append({
                        'buy_date': buy_date, 'code': code,
                        'f1_t0_pct': float(pct[t0]),
                        'f2_t0_vol_ratio': float(v0 / vol[max(0, t0 - 10):t0].mean()) if t0 >= 10 and vol[max(0, t0 - 10):t0].mean() > 0 else np.nan,
                        'f3_ta_shrink': float(vol[ta] / v0) if v0 > 0 else np.nan,
                        'f4_ta_drawdown': float(close[ta] / p0 - 1),
                        'f5_ta_gap': float(ta - t0),
                        'f6_tb_strength': float(close[k] / p0 - 1),
                        'f7_tb_gap': float(k - t0),
                        'f8_chain_len': float(k + 1 - t0),
                        'f9_tb_vol': float(vol[k] / m10) if m10 > 0 else np.nan,
                        'f10_tb_ma5_slope': float((ma5[k] - ma5[k - 3]) / ma5[k - 3]) if k >= 3 and ma5[k - 3] > 0 else np.nan,
                        'f11_tb_ma_align': float(1 if (ma5[k] > ma10[k] > ma20[k]) else (-1 if (ma5[k] < ma10[k] < ma20[k]) else 0)),
                        'tb_close': float(close[k]), 't0_date': dates[t0], 'ta_date': dates[ta],
                    })
                break
    return res

def ie_active(members, date_str):
    """members月份分期dict -> 当日成分股集合"""
    import ideal_engine as ie
    return ie.active_members(members, date_str)

def simulate_sell(d, t, c):
    """模拟持仓t(未完成信号)在c参数下的卖出: 与ideal_engine逐行一致
    d: 个股dailies; t: {code,buy_date}; c: SELL_COMBO参数(含mode)
    返回None=未触发(继续持仓), 否则{sell_date,sell_price,reason,hold_days,ret_net}"""
    sub = d[d['trade_date'] > t['buy_date']].sort_values('trade_date').reset_index(drop=True)
    if len(sub) == 0:
        return None
    rows0 = d[d['trade_date'] == t['buy_date']]
    if len(rows0) == 0:
        return None
    buy_px = float(rows0.iloc[0]['hfq_open']) * (1 + g.SLIP)      # 判定入场=后复权
    fill_buy = float(rows0.iloc[0]['raw_open']) * (1 + g.SLIP)    # 成交价=原始价
    if buy_px <= 0 or fill_buy <= 0:
        return None
    high_since = buy_px
    hold_n = 0
    mode = c.get('mode', 'A' if c['dd'] >= 0.99 else ('C' if c['stop'] is not None else 'merge'))
    stop = c['stop']
    for _, r in sub.iterrows():
        date = r['trade_date']
        cur = float(r['hfq_close'])
        hold_n += 1
        sell = None
        if mode == 'A':
            if cur / buy_px - 1 >= c['profit']:
                sell = 'profit'
            elif float(r['hfq_low']) <= buy_px * (1 - (stop if stop is not None else c['dd'])):
                sell = 'stop'
            elif hold_n >= c['hold']:
                sell = 'timeout'
        elif mode == 'B':
            if float(r['hfq_low']) <= high_since * (1 - c['profit']):
                sell = 'trail'
            elif float(r['hfq_low']) <= buy_px * (1 - (stop if stop is not None else c['dd'])):
                sell = 'stop'
            elif hold_n >= c['hold']:
                sell = 'timeout'
        elif mode == 'C':
            if float(r['hfq_low']) <= high_since * (1 - c['dd']):
                sell = 'dd'
            elif stop is not None and float(r['hfq_low']) <= buy_px * (1 - stop):
                sell = 'stop'
            elif cur / buy_px - 1 >= c['profit']:
                sell = 'profit'
            elif hold_n >= c['hold']:
                sell = 'timeout'
        else:
            if float(r['hfq_low']) <= high_since * (1 - c['dd']):
                sell = 'dd'
            elif cur / buy_px - 1 >= c['profit']:
                sell = 'profit'
            elif hold_n >= c['hold']:
                sell = 'timeout'
        high_since = max(high_since, cur)
        if sell:
            net_ret = (cur / buy_px - 1) - (g.COST_BUY + g.COST_SELL + 2 * g.SLIP)   # 判定口径收益率
            fill = float(r['raw_close'])                                            # 成交价=原始价
            return {'sell_date': date, 'sell_price': fill, 'reason': sell, 'hold_days': hold_n, 'ret_net': net_ret, 'buy_price': fill_buy}
    return None

# 卖出参数(最终定稿: hs300/zz2000=v3, zz500=v2+防御值, zz1000=v2)
SELL_COMBO = {
    'hs300': {'mode': 'merge', 'profit': 0.25, 'dd': 0.99, 'stop': None, 'hold': 15},
    'zz500': {'mode': 'B', 'profit': 0.45, 'dd': 0.99, 'stop': 0.20, 'hold': 15},
    'zz1000': {'mode': 'C', 'profit': 0.20, 'dd': 0.30, 'stop': 0.25, 'hold': 15},
    'zz2000': {'mode': 'C', 'profit': 0.15, 'dd': 0.25, 'stop': 0.30, 'hold': 15},
}
SELL_CN = {'profit': '止盈', 'dd': '回撤', 'stop': '止损', 'timeout': '超时', 'trail': '跟踪止盈'}

def main():
    import sys as _sys
    target = _sys.argv[1] if len(_sys.argv) > 1 else None
    model = json.load(open(f'{BASE}/strength_model.json'))
    w = model['weights']
    pools_m = model['pools']
    # 守卫(2026-09-14 事故): 模型权重键必须与 FCOLS 一致, 否则 raw_score 求积时 KeyError 只在一只票上炸开、
    # 报错点远离真因(0913 模型重建丢了 f12/f13/f14 → 14→11 键, 直到次日首跑 scan 才炸)。
    _miss = [fc for fc in FCOLS if fc not in w]
    _extra = [fc for fc in w if fc not in FCOLS]
    if _miss or _extra:
        raise SystemExit(
            f"❌ 模型/特征不一致, 拒绝扫描:\n"
            f"   模型 {BASE}/strength_model.json 权重键 {len(w)} 个, 缺 {_miss}, 多 {_extra}\n"
            f"   scan FCOLS {len(FCOLS)} 个 —— 两者必须完全一致。\n"
            f"   处置: 用 build_strength_model_v3.rebuild() 从 baseline 重建 14 特征模型(与对端同构), 或按裁定改 FCOLS。")
    def pct(pool, fc, v):
        qs = pools_m[pool]['percentiles'][fc]
        if qs is None or v is None or (isinstance(v, float) and np.isnan(v)):
            return 0.5
        return float(np.searchsorted(qs, v) / 100.0)
    def std_strength(pool, raw):
        """raw_score -> 该池基准池百分位(0-100)"""
        rs = pools_m[pool]['raw_scores']
        return float(np.searchsorted(rs, raw) / len(rs) * 100)
    sb = pd.read_pickle(f'{BASE}/stock_basic.pkl')
    name_map = dict(zip(sb['ts_code'], sb['name']))
    ind_map = dict(zip(sb['ts_code'], sb['industry']))
    cache = {}
    all_sig = []
    for pool, bp in COMBO.items():
        if pool not in cache:
            cache[pool] = g.load_pool(pool)
        members, dailies = cache[pool]
        dates = sorted(set().union(*[set(d['trade_date']) for d in dailies.values()]))
        if target is None:
            target = dates[-1]
    baseline = f'{BASE}/strength_baseline.csv'
    bdf = pd.read_csv(baseline)
    # === 卖出检查: 未完成持仓(此前推送的买入信号)模拟卖出 ===
    open_pos = bdf[bdf['ret'].isna()]
    if len(open_pos):
        completed = []
        for _, t in open_pos.iterrows():
            pool = str(t['pool'])
            d = cache[pool][1].get(t['code'])
            if d is None:
                continue
            res = simulate_sell(d, {'code': t['code'], 'buy_date': str(t['buy_date'])}, SELL_COMBO[pool])
            if res is not None:
                completed.append((t, pool, res))
        if completed:
            print(f"📤 卖出信号(当日触发 {len(completed)}笔, 交易完成):")
            for t, pool, res in completed:
                name = name_map.get(t['code'], '?')
                ind = ind_map.get(t['code'], '?')
                print(f"【{PN[pool]}】{name}({t['code'][:8]})·{ind} 买{res['buy_price']:.2f}→卖{res['sell_price']:.2f} 持{res['hold_days']}天 {SELL_CN[res['reason']]} {res['ret_net']*100:+.1f}% 卖出于{res['sell_date']}")
                idx = bdf[(bdf['code'] == t['code']) & (bdf['buy_date'].astype(str) == str(t['buy_date']))].index
                if len(idx):
                    # sell_date列在csv中为float64(未完成持仓=NaN, 历史完成=20260828.0):
                    # 字符串日期'20260903'直接写入会TypeError, 须先转数值
                    bdf.loc[idx, ['ret', 'sell_date', 'sell_price', 'reason', 'hold_days']] = [
                        res['ret_net'], int(pd.to_numeric(res['sell_date'])), res['sell_price'], res['reason'], res['hold_days']]
            bdf.to_csv(baseline, index=False)
            print(f"(已按卖出日完成交易, 纳入基准池)")
            sys.path.insert(0, '/Users/xpresso/zt_app/code')
            import build_strength_model_v3 as bsm
            bsm.rebuild(weights=w)
            print(f"(分池基准已更新: {len(bdf)}笔, 含完成的交易)")
        else:
            print(f"📥 持仓检查: {len(open_pos)}笔持仓中, 今日无触发卖出")
    # === 买入信号扫描(TB突破, 明日可买) ===
    active_sets = {}
    pool_day_stat = {}  # pool -> (当日涨停数归一, 当日成分均涨跌幅)
    for pool, bp in COMBO.items():
        members, dailies = cache[pool]
        active_sets[pool] = ie_active(members, target)
        # 当日池统计: target日成分涨停家数/平均涨跌幅(上下文特征f13/f14)
        pz = pc = 0
        ps = 0.0
        for code2, d2 in dailies.items():
            rows2 = d2[d2['trade_date'] == target]
            if len(rows2) == 0:
                continue
            pp = float(rows2.iloc[0]['pct_chg'])
            if np.isnan(pp):
                continue
            pc += 1
            ps += pp
            if pp >= 9.8:
                pz += 1
        pool_day_stat[pool] = (pz / 20.0, ps / pc if pc else np.nan)
        for code, d in dailies.items():
            if code not in active_sets[pool]:
                continue
            if len(d) < 25:
                continue
            if str(d['trade_date'].iloc[-1]) < target:
                continue
            hits = find_tb_today(d, code, bp, target)
            if not hits:
                continue
            p_zt, p_avg = pool_day_stat[pool]
            for h in hits:
                if h['buy_date'] is None:
                    h['buy_pending'] = True  # 突破日=今日且次日数据未出: 明日开盘买入待确认
                all_sig.append({'pool': pool, 'code': code, 'name': name_map.get(code, '?'), 'ind': ind_map.get(code, '?'),
                                **h, 'f13_pool_zt': p_zt, 'f14_pool_avgret': p_avg})
    if not all_sig:
        print(f"📭 四池今日({target})无满足TB突破条件的股票, 明日无信号")
        return
    # 拥挤度f12(当日池信号数/池历史日均信号数) + 统一重算raw/strength(v6: 14特征)
    sig_by_pool = {}
    for x in all_sig:
        sig_by_pool[x['pool']] = sig_by_pool.get(x['pool'], 0) + 1
    for x in all_sig:
        x['f12_crowd_pool'] = sig_by_pool[x['pool']] / crowd_denom(bdf, x['pool'], x['buy_date'])
        raw = sum(w[fc] * pct(x['pool'], fc, x.get(fc, np.nan)) for fc in FCOLS)
        x['raw_score'] = raw
        x['strength'] = std_strength(x['pool'], raw)
    print(f"📊 A股四池涨停突破信号扫描({target})")
    print(f"共{len(all_sig)}只满足TB条件, 明日开盘可买入, 按强度排序(0-100):")
    for pool in ['hs300', 'zz500', 'zz1000', 'zz2000']:
        ps = sorted([x for x in all_sig if x['pool'] == pool], key=lambda x: -x['strength'])
        if not ps:
            continue
        print(f"\n【{PN[pool]}】{len(ps)}只:")
        for x in ps:
            tag = '🟢强' if x['strength'] >= 70 else ('🟡中' if x['strength'] >= 40 else '⚪弱')
            when = '明日开盘买入(待确认)' if x.get('buy_pending') else f"买入日{x['buy_date']}"
            print(f"{tag} {x['name']}({x['code'][:8]})·{x['ind']} 强度{x['strength']:.0f} | {when} | tb涨{x['f6_tb_strength']*100:+.1f}% 缩量{x['f3_ta_shrink']:.2f} 链长{x['f8_chain_len']:.0f}天")
    # 滚动基准池: 新信号纳入
    baseline = f'{BASE}/strength_baseline.csv'
    bdf = pd.read_csv(baseline, dtype={'buy_date': str})  # 必须str: 否则int64与字符串'20260901'混型致drop_duplicates失效
    # 补记: 突破日=昨日、买入日=今日的信号(昨日buy_pending跳过, 今日数据确认后补写)
    # 语义: 昨日17:00扫描时该链突破日=数据末日被标buy_pending不落库, 今日开盘已买入, 今日扫描补记
    from plot_ideal_top10 import chain_for_buy  # 延迟import避免循环依赖
    from build_strength_v4 import feats
    for pool in ['hs300', 'zz500', 'zz1000', 'zz2000']:
        members, dailies = cache[pool]
        bp = COMBO[pool]
        sig = g.build_sig(dailies, bp)
        for code, bs in sig.items():
            if code not in active_sets[pool]:
                continue
            for (bd, t0d) in bs:
                if bd != target:
                    continue  # 只补"买入日=今日"的链
                if code in {x['code'] for x in all_sig}:
                    continue
                d = dailies[code].sort_values('trade_date').reset_index(drop=True)
                t0i, tai, tbi = chain_for_buy(d, code, bp, bd)
                if t0i is None:
                    continue
                f = feats(d, t0i, tai, tbi, None)
                p_zt, p_avg = pool_day_stat.get(pool, (np.nan, np.nan))
                f['f13_pool_zt'] = p_zt
                f['f14_pool_avgret'] = p_avg
                f['f12_crowd_pool'] = sig_by_pool.get(pool, 1) / crowd_denom(bdf, pool, bd)
                raw = sum(w[fc] * pct(pool, fc, f[fc]) for fc in FCOLS)
                all_sig.append({'pool': pool, 'code': code, 'name': name_map.get(code, '?'), 'ind': ind_map.get(code, '?'),
                                'strength': std_strength(pool, raw), 'raw_score': raw, **f, 'buy_date': bd})
    new_rows = []
    for x in all_sig:
        if x.get('buy_pending'):
            continue  # 明日买入待确认信号不写入基准池(次日扫描补记)
        row = {fc: x.get(fc, np.nan) for fc in FCOLS}
        row.update({'pool': x['pool'], 'code': x['code'], 'buy_date': x['buy_date'], 'ret': np.nan,
                    'reason': 'signal', 'hold_days': np.nan, 'raw_score': x['raw_score'], 'strength': x['strength']})
        new_rows.append(row)
    if new_rows:
        bdf = pd.concat([bdf, pd.DataFrame(new_rows)], ignore_index=True)
        bdf = bdf.drop_duplicates(subset=['code', 'buy_date'], keep='last')
        bdf.to_csv(baseline, index=False)
        # 重算模型v3(权重固定, 分池基准)
        sys.path.insert(0, '/Users/xpresso/zt_app/code')
        import build_strength_model_v3 as bsm
        bsm.rebuild(weights=w)
        print(f"\n(分池基准已滚动更新: {len(bdf)}笔)")

if __name__ == '__main__':
    main()