#!/usr/bin/env python3
"""今日TB突破信号(明日开盘买入)按强度top10/池 交易详情图(40格模板v3)
与plot_holdings.py同布局: 列=池(左→右hs300/zz500/zz1000/zz2000), 行=强度(上→下递减)
只画 bd==最新交易日 的信号(明日买入), 非持仓图
"""
import pandas as pd
import numpy as np
import json, sys
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
plt.rcParams['font.sans-serif'] = ['STHeiti', 'Noto Sans CJK SC', 'WenQuanYi Zen Hei', 'SimHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False
sys.path.insert(0, '/Users/xpresso/zt_app/code')
import grid_or as g
from scan_daily_tb import COMBO, PN
from plot_ideal_top10 import chain_for_buy
from build_strength_v4 import feats

BASE = g.BASE
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']))
model = json.load(open(f'{BASE}/strength_model.json'))
w = model['weights']
pools_m = model['pools']

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 strength_of(pool, fdict):
    raw = sum(w[fc] * pct(pool, fc, fdict[fc]) for fc in w)
    rs = pools_m[pool]['raw_scores']
    return float(np.searchsorted(rs, raw) / len(rs) * 100)

cache = {}
def load(pool):
    if pool not in cache:
        cache[pool] = g.load_pool(pool)
    return cache[pool]

# 1. 收集今日信号 = 突破日(tb)==最新交易日, 明日开盘买入(与scan_daily_tb.find_tb_today口径一致)
all_dates = {}
for pool in COMBO:
    members, dailies = load(pool)
    dates = sorted(set().union(*[set(d['trade_date']) for d in dailies.values()]))
    all_dates[pool] = dates[-1]
target = max(all_dates.values())
print(f"目标交易日(最新): {target}", flush=True)

from scan_daily_tb import find_tb_today

def idx_of(d2, date_str):
    if date_str is None:
        return None
    pos = np.where(d2['trade_date'].values == str(date_str))[0]
    return int(pos[0]) if len(pos) else None

holdings = []
# 池历史日均信号数(拥挤度f12基准, 与scan一致: baseline池行数/去重买入日数)
bdf_ref = pd.read_csv(f'{BASE}/strength_baseline.csv', dtype={'buy_date': str})
pool_avg_sig = {p: bdf_ref[bdf_ref['pool'] == p].groupby('buy_date').size().mean() for p in COMBO}
for pool in COMBO:
    bp = COMBO[pool]
    members, dailies = load(pool)
    # 当期成分过滤(与回测引擎口径一致)
    active = g.active_members(members, target)
    # 当日池统计(上下文特征f13/f14, 与scan同口径)
    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
    p_zt, p_avg = pz / 20.0, (ps / pc if pc else np.nan)
    for code, d0 in dailies.items():
        if code not in active:
            continue
        d2 = d0.sort_values('trade_date').reset_index(drop=True)
        hits = find_tb_today(d2, code, bp, target)
        if not hits:
            continue
        for h in hits:
            bd = h['buy_date']  # 明日买入日(None=数据末日突破, 待确认)
            t0i = idx_of(d2, h['t0_date'])
            tai = idx_of(d2, h['ta_date'])
            tbi = idx_of(d2, target)  # tb=今日
            if t0i is None or tai is None or tbi is None:
                continue
            f = feats(d2, t0i, tai, tbi, None)
            holdings.append({'pool': pool, 'code': code, 'buy_date': bd, 't0i': t0i, 'tai': tai, 'tbi': tbi,
                             'feat': f, 'name': name_map.get(code, '?'), 'ind': ind_map.get(code, '?')})
    # 池内信号收集完毕: 补上下文特征+算强度(v6 14特征)
    n_sig_pool = sum(1 for h in holdings if h['pool'] == pool)
    for h in [h for h in holdings if h['pool'] == pool]:
        f = h['feat']
        f['f12_crowd_pool'] = n_sig_pool / max(1.0, pool_avg_sig[pool])
        f['f13_pool_zt'] = p_zt
        f['f14_pool_avgret'] = p_avg
        h['std'] = strength_of(pool, f)
        rows = dailies[h['code']].sort_values('trade_date').reset_index(drop=True)
        rows = rows[rows['trade_date'] == target]
        h['buy_px'] = float(rows.iloc[0]['open']) * (1 + g.SLIP) if len(rows) else None
        h['cur'] = float(dailies[h['code']].sort_values('trade_date')['close'].iloc[-1])
    print(f"{PN[pool]}: 今日突破信号{n_sig_pool}只", flush=True)

# 同池同股去重: 信号场景保留强度最高(与sim_live口径一致: 资金只能买一次, 取最强链)
dedup = {}
for h in holdings:
    k = (h['pool'], h['code'])
    if k not in dedup or h['std'] > dedup[k]['std']:
        dedup[k] = h
holdings = list(dedup.values())

# 2. 每池强度top10
samples = {}
for pool in COMBO:
    ph = sorted([h for h in holdings if h['pool'] == pool], key=lambda x: x['std'], reverse=True)[:10]
    samples[pool] = ph
    if ph:
        print(f"{PN[pool]} top10: 强度{ph[0]['std']:.0f}~{ph[-1]['std']:.0f}", flush=True)

# 3. 画图(与plot_holdings.py同模板)
def draw_holding(gs, row, col, h, d):
    t0i, tai, tbi = h['t0i'], h['tai'], h['tbi']
    std = h['std']
    tier = '强' if std >= 85 else ('中' if std >= 30 else '弱')
    tier_c = {'强': '#E91E63', '中': '#FF8F00', '弱': '#7f8c8d'}[tier]
    gss = gs[row, col].subgridspec(2, 1, height_ratios=[2.6, 1], hspace=0.06)
    ax = fig.add_subplot(gss[0])
    axv = fig.add_subplot(gss[1], sharex=ax)
    t0d = pd.Timestamp(d['trade_date'].iloc[t0i])
    win_start = t0d - pd.Timedelta(days=4)
    mask = pd.to_datetime(d['trade_date']) >= win_start
    sub = d[mask].reset_index(drop=True)
    x = np.arange(len(sub))
    for xi, r in sub.iterrows():
        o, hi, lo, c_ = r['open'], r['high'], r['low'], r['close']
        up = c_ >= o
        color = '#e74c3c' if up else '#2ecc71'
        ax.plot([xi, xi], [lo, hi], color=color, linewidth=0.7, zorder=2)
        ax.add_patch(plt.Rectangle((xi - 0.35, min(o, c_)), 0.7, max(abs(c_ - o), 0.01), facecolor=color, edgecolor=color, linewidth=0.3, zorder=3))
        axv.bar(xi, r['vol'], color=color, alpha=0.7, width=0.7)
    # 信号链 (日期→sub内序号) + 文字标注
    ymax = ax.get_ylim()[1]
    ymin = ax.get_ylim()[0]
    for idx, color, txt in [(t0i, '#e74c3c', 'T0'), (tai, '#1565C0', 'ta'), (tbi, '#FF8F00', 'tb')]:
        dt_s = str(d['trade_date'].iloc[idx])
        pos = np.where(sub['trade_date'].values == dt_s)[0]
        if len(pos):
            ax.axvline(pos[0], color=color, linestyle='--', linewidth=0.9, alpha=0.8)
            axv.axvline(pos[0], color=color, linestyle='--', linewidth=0.9, alpha=0.8)
            ax.text(pos[0], ymin + (ymax - ymin) * 0.97, txt, fontsize=8, color=color, ha='center', fontweight='bold',
                    bbox=dict(facecolor='white', alpha=0.7, edgecolor='none', pad=0.6))
    # 买入点不再标注(明日买入, 价格未定—信号图只展示T0/ta/tb链+今收参照)
    # 最新收盘参照(信号日尚无持仓, 标今日收盘价)
    xlast = x[-1]
    ax.scatter(xlast, h['cur'], marker='o', color='#2c3e50', s=50, zorder=6, edgecolors='white', linewidths=0.8)
    ax.annotate(f"今收{h['cur']:.2f}", (xlast, h['cur']), textcoords='offset points', xytext=(0, -11), fontsize=8, ha='center', color='#2c3e50', fontweight='bold')
    ax.set_title(f"{h['name']}·{h['ind']}\n{h['code'][:9]} {PN[h['pool']]} 明日开盘买入", fontsize=8.5, fontweight='bold')
    ax.text(0.98, 0.02, f"强度{std:.0f}·{tier}", transform=ax.transAxes, fontsize=7.5, fontweight='bold',
            ha='right', va='bottom', color=tier_c, bbox=dict(facecolor='white', alpha=0.85, edgecolor=tier_c, linewidth=0.8, pad=1.5))
    # 关键节点x轴(年月日) - 统一sub内序号
    key_ticks = []
    for idx in [t0i, tai, tbi]:
        dt_s = str(d['trade_date'].iloc[idx])
        pos = np.where(sub['trade_date'].values == dt_s)[0]
        if len(pos):
            key_ticks.append((pos[0], dt_s))
    key_ticks.append((xlast, d.iloc[-1]['trade_date']))
    ticks = sorted(set(p for p, _ in key_ticks))
    labs = dict(key_ticks)
    ax.set_xticks(ticks)
    axv.set_xticks(ticks)
    axv.set_xticklabels([f"{labs[p][2:4]}-{labs[p][4:6]}-{labs[p][6:]}" for p in ticks], fontsize=6, rotation=45)
    ax.tick_params(labelbottom=False)
    ax.tick_params(labelsize=5)
    axv.tick_params(labelsize=4)
    ax.grid(True, alpha=0.12)
    axv.grid(True, alpha=0.12)

fig = plt.figure(figsize=(26, 31))
gs = fig.add_gridspec(10, 4, hspace=0.34, wspace=0.22)
# 列标题: 左→右 = 四池 (统计框下方, 不与标题/统计框重叠)
for pi, pool in enumerate(COMBO):
    fig.text(0.05 + (pi + 0.5) * (0.935 / 4), 0.954, PN[pool], fontsize=14, fontweight='bold', ha='center', color='#2c3e50')
for pi, pool in enumerate(COMBO):
    for j in range(10):
        row, col = j, pi  # 行=强度排名(上→下递减), 列=池(左→右四池)
        if j < len(samples[pool]):
            h = samples[pool][j]
            d = cache[pool][1][h['code']].sort_values('trade_date').reset_index(drop=True)
            draw_holding(gs, row, col, h, d)
        else:
            ax = fig.add_subplot(gs[row, col])
            ax.text(0.5, 0.5, f"{PN[pool]}\n今日无信号\n(当前{len(samples[pool])}只)",
                    ha='center', va='center', fontsize=13, color='#999', fontweight='bold')
            ax.set_xticks([]); ax.set_yticks([])
            for sp in ax.spines.values():
                sp.set_visible(False)
from matplotlib.lines import Line2D
leg = [Line2D([0], [0], marker='o', color='w', markerfacecolor='#2c3e50', markersize=9, label='今日收盘(参照)'),
       Line2D([0], [0], color='#e74c3c', linestyle='--', label='T0涨停日'),
       Line2D([0], [0], color='#1565C0', linestyle='--', label='ta缩量回调日'),
       Line2D([0], [0], color='#FF8F00', linestyle='--', label='tb突破日')]
fig.legend(handles=leg, loc='lower center', ncol=4, fontsize=12, framealpha=0.9, bbox_to_anchor=(0.5, 0.012))

fig.suptitle(f'四池今日TB突破信号(明日开盘买入)强度Top10共40只: 强度v6分档标注(右上角)', fontsize=15, y=0.997)
n_total = len(holdings)
n_strong = sum(1 for h in holdings if h['std'] >= 85)
n_mid = sum(1 for h in holdings if 30 <= h['std'] < 85)
n_weak = sum(1 for h in holdings if h['std'] < 30)
bs_ = bdf_ref['strength'].dropna()
rt_ = bdf_ref.loc[bs_.index, 'ret']
fig.text(0.5, 0.978, f"今日({target})TB突破信号共{n_total}只: 强{n_strong} 中{n_mid} 弱{n_weak}   |   分档(全{len(bs_)}笔): 强≥85均{rt_[bs_>=85].mean()*100:+.1f}%胜{(rt_[bs_>=85]>0).mean()*100:.0f}% 中30-85均{rt_[(bs_>=30)&(bs_<85)].mean()*100:+.1f}%胜{(rt_[(bs_>=30)&(bs_<85)]>0).mean()*100:.0f}% 弱<30均{rt_[bs_<30].mean()*100:+.1f}%胜{(rt_[bs_<30]>0).mean()*100:.0f}%   |   买入门槛=强档线: ≥85",
         fontsize=11.5, ha='center', bbox=dict(facecolor='#f8f8f8', edgecolor='#999', linewidth=0.8, boxstyle='round,pad=0.4'))
fig.subplots_adjust(top=0.945, bottom=0.05, left=0.05, right=0.985, hspace=0.34, wspace=0.22)
out = f'{BASE}/backtest_zt_signals_today_top10.png'
plt.savefig(out, dpi=200)
print(f"今日信号top10图: {out}", flush=True)