#!/usr/bin/env python3
"""重建基准池v4: v2定稿四池回测交易 → 11特征CSV → 强度模型(4768笔rho权重+分池百分位)"""
import pandas as pd
import numpy as np
import json, sys
sys.path.insert(0, '/tmp')
import grid_or as g
import ideal_engine as ie
from scan_daily_tb import COMBO, SELL_COMBO
from plot_ideal_top10 import chain_for_buy

BASE = g.BASE
PN = {'hs300': '沪深300', 'zz500': '中证500', 'zz1000': '中证1000', 'zz2000': '中证2000'}

def feats(d, t0i, tai, tbi, t0d):
    """11特征, 与scan_daily_tb.find_tb_today一致"""
    close = d['close'].astype(float).values
    vol = d['vol'].astype(float).values
    pct = d['pct_chg'].astype(float).values
    ma5 = pd.Series(close).rolling(5).mean().values
    ma10 = pd.Series(close).rolling(10).mean().values
    ma20 = pd.Series(close).rolling(20).mean().values
    v0 = vol[t0i]
    m10_t0 = vol[max(0, t0i - 10):t0i].mean() if t0i >= 10 else vol[:t0i].mean()
    m10_tb = vol[max(0, tbi - 10):tbi].mean() if tbi >= 10 else vol[:tbi].mean()
    f = {
        'f1_t0_pct': float(pct[t0i]),
        'f2_t0_vol_ratio': float(v0 / m10_t0) if m10_t0 > 0 else np.nan,
        'f3_ta_shrink': float(vol[tai] / v0) if v0 > 0 else np.nan,
        'f4_ta_drawdown': float(close[tai] / close[t0i] - 1),
        'f5_ta_gap': float(tai - t0i),
        'f6_tb_strength': float(close[tbi] / close[t0i] - 1),
        'f7_tb_gap': float(tbi - t0i),
        'f8_chain_len': float(tbi + 1 - t0i),
        'f9_tb_vol': float(vol[tbi] / m10_tb) if m10_tb > 0 else np.nan,
        'f10_tb_ma5_slope': float((ma5[tbi] - ma5[tbi - 3]) / ma5[tbi - 3]) if tbi >= 3 and ma5[tbi - 3] > 0 else np.nan,
        'f11_tb_ma_align': float(1 if (ma5[tbi] > ma10[tbi] > ma20[tbi]) else (-1 if (ma5[tbi] < ma10[tbi] < ma20[tbi]) else 0)),
    }
    return f

def main():
    rows = []
    cache = {}
    for pool in COMBO:
        bp = COMBO[pool]; sc = SELL_COMBO[pool]
        if pool not in cache:
            cache[pool] = g.load_pool(pool)
        members, dailies = cache[pool]
        sig = g.build_sig(dailies, bp)
        trades = ie.ideal_backtest(pool, members, dailies, sig, sc['profit'], sc['dd'], mode=sc['mode'], stop=sc['stop'], hold=sc['hold'], start='20240201')
        for t in trades:
            d = dailies[t['code']].sort_values('trade_date').reset_index(drop=True)
            t0i, tai, tbi = chain_for_buy(d, t['code'], bp, str(t['buy_date']))
            if t0i is None:
                continue
            f = feats(d, t0i, tai, tbi, None)
            rows.append({'code': t['code'], 'pool': pool, 'buy_date': str(t['buy_date']),
                         'sell_date': str(t['sell_date']), 'ret': t['ret_net'],
                         **f})
        print(f"{PN[pool]}: {len(trades)}笔", flush=True)
    df = pd.DataFrame(rows)
    print(f"总交易: {len(df)}笔")
    df.to_csv(f'{BASE}/ideal_signal_features_v2.csv', index=False)
    # 特征-收益相关性
    fcols = [f'f{i}_' for i in range(1, 12)]
    fcols = [c for c in df.columns if c.startswith('f')]
    from scipy import stats
    print("\n特征 | 整体rho | (池内rho)")
    for c in fcols:
        rhos = []
        for pool in COMBO:
            sub = df[df['pool'] == pool].dropna(subset=[c, 'ret'])
            if len(sub) > 30:
                rhos.append(round(stats.spearmanr(sub[c], sub['ret']).correlation, 3))
        print(f"{c}: 整体{stats.spearmanr(df[c].dropna(), df['ret'][df[c].notna()]).correlation:+.3f} 池内{rhos}")
    # 构建强度模型v4: 权重=整体rho, 分池percentile (格式与scan_daily_tb兼容: percentiles+raw_scores)
    weights = {}
    for c in fcols:
        ok = df[c].notna() & df['ret'].notna()
        weights[c] = stats.spearmanr(df.loc[ok, c], df.loc[ok, 'ret']).correlation
    pools_m = {}
    for pool in COMBO:
        sub = df[df['pool'] == pool]
        per = {}
        for c in fcols:
            qs2 = sub[sub[c].notna()][c]
            per[c] = [float(np.nanpercentile(qs2, q)) for q in range(101)] if len(qs2) > 3 else None
        pools_m[pool] = {'n': len(sub), 'percentiles': per, 'raw_scores': None}
    # raw_score分布 (与scan一致: searchsorted默认left)
    def raw_score(row):
        s = 0.0
        for c in fcols:
            if np.isnan(row[c]):
                continue
            qs = pools_m[row['pool']]['percentiles'][c]
            if qs is None:
                continue
            p = np.searchsorted(qs, row[c]) / 100.0
            s += weights[c] * p
        return s
    df['raw_score'] = df.apply(raw_score, axis=1)
    for pool in COMBO:
        dist = sorted(df[df['pool'] == pool]['raw_score'].tolist())
        pools_m[pool]['raw_scores'] = dist
    # 0-100 分池百分位
    def final_score(row):
        dist = pools_m[row['pool']]['raw_scores']
        return 100 * np.searchsorted(dist, row['raw_score']) / len(dist)
    df['strength'] = df.apply(final_score, axis=1)
    # 分档统计
    tiers = {'强': (70, 100), '中': (30, 70), '弱': (0, 30)}
    print("\n分档统计(4768笔):")
    stats_txt = []
    for name, (lo, hi) in tiers.items():
        sub = df[(df['strength'] >= lo) & (df['strength'] < hi)]
        if len(sub):
            s = f"{name}{len(sub)}笔 均{np.mean(sub['ret'])*100:+.1f}% 高{np.max(sub['ret'])*100:+.1f}% 低{np.min(sub['ret'])*100:+.1f}% 胜{(sub['ret']>0).mean()*100:.0f}%"
            stats_txt.append(s)
            print(s)
    # 保存
    model = {'weights': weights, 'pools': pools_m, 'n_trades': len(df), 'stats': stats_txt}
    json.dump(model, open(f'{BASE}/strength_model.json', 'w'), ensure_ascii=False, default=str)
    df.to_csv(f'{BASE}/strength_baseline.csv', index=False)
    # 池内rho验证
    print("\n池内强度-收益rho:")
    for pool in COMBO:
        sub = df[df['pool'] == pool]
        rho = stats.spearmanr(sub['strength'], sub['ret']).correlation
        strong = sub[sub['strength'] >= 70]['ret'].mean() if (sub['strength'] >= 70).any() else 0
        weak = sub[sub['strength'] < 30]['ret'].mean() if (sub['strength'] < 30).any() else 0
        print(f"{PN[pool]}: rho={rho:+.3f} 强均{strong*100:+.1f}% vs 弱均{weak*100:+.1f}% (n={len(sub)})")

if __name__ == '__main__':
    main()