#!/usr/bin/env python3
"""step40_mkt_atmosphere.py — 市场气氛 vs 信号质量（底部确认依据）
老莫洞察：2024-02 全场下跌，气氛不对——超跌信号可能是"下跌中继"而非"底部"
构建气氛指标（无前视）：
- 大盘: mkt_ret5(短期加速跌) / mkt_dd60(深度) / mkt_adx(趋势强度)
- 消息面: 全市场每日新闻总量(气氛冷热)
- 行业共振: 行业也在深跌 vs 个股独立超跌
扫描这些 vs 信号后续收益，找"底部确认"的数据依据
"""
import numpy as np
import pandas as pd
import sqlite3

print("=== 加载 ===", flush=True)
panel = pd.read_pickle("/tmp/panel_12d.pkl")
panel = panel.sort_values(["code", "date"]).reset_index(drop=True)
panel["fwd_ret60"] = panel.groupby("code")["close"].transform(lambda x: x.shift(-60)/x - 1) * 100
panel["is_big"] = (panel["fwd_ret60"] >= 50).astype(int)
print("面板:", len(panel), flush=True)

# 大盘气氛指标
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["mkt_ret5"] = idx_df["close"].pct_change(5) * 100
idx_df["mkt_ret10"] = idx_df["close"].pct_change(10) * 100
idx_df["mkt_ret20"] = idx_df["close"].pct_change(20) * 100
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 = 0
for c in close_arr:
    if c < close_arr[len(down_days)-1] if len(down_days) > 0 else False:
        streak += 1
    else:
        streak = 0
    down_days.append(streak)
idx_df["mkt_down_days"] = down_days
mkt_map = dict(zip(idx_df["date"], zip(idx_df["mkt_ret5"], idx_df["mkt_ret10"], idx_df["mkt_ret20"], idx_df["mkt_dd60"], idx_df["mkt_down_days"])))
panel["mkt_ret5"] = panel["date"].map(lambda d: mkt_map[d][0] if d in mkt_map else np.nan)
panel["mkt_ret10"] = panel["date"].map(lambda d: mkt_map[d][1] if d in mkt_map else np.nan)
panel["mkt_ret20"] = panel["date"].map(lambda d: mkt_map[d][2] if d in mkt_map else np.nan)
panel["mkt_dd60"] = panel["date"].map(lambda d: mkt_map[d][3] if d in mkt_map else np.nan)
panel["mkt_down_days"] = panel["date"].map(lambda d: mkt_map[d][4] if d in mkt_map else np.nan)
print("大盘指标合并", flush=True)

# 全市场每日新闻量（消息面气氛）
news_cnt = pd.read_sql("SELECT substr(date,1,10) d, COUNT(*) c FROM stock_news GROUP BY d", conn)
news_cnt["d"] = news_cnt["d"].astype(str)
news_cnt["mkt_news"] = news_cnt["c"].rolling(5, min_periods=1).mean()  # 5日均量
news_map = dict(zip(news_cnt["d"], news_cnt["mkt_news"]))
panel["mkt_news5"] = panel["date"].map(lambda d: news_map.get(d, np.nan))
print("新闻气氛合并", flush=True)

# 信号
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)
)
sig = panel[sig_cond].dropna(subset=["fwd_ret60"]).copy()
print("信号:", len(sig), flush=True)
base_rate = sig["is_big"].mean() * 100
base_avg = sig["fwd_ret60"].mean()
print("信号基线: 大涨率={:.2f}% avg60={:.2f}%".format(base_rate, base_avg), flush=True)

def scan(cond, label, min_n=50):
    m = sig[cond]
    if len(m) < min_n:
        print("{}: n={} 不足".format(label, len(m)), flush=True)
        return
    print("{}: n={} 大涨率={:.2f}% avg60={:.2f}% wr={:.1f}% 超额={:+.2f}pp".format(
        label, len(m), m["is_big"].mean()*100, m["fwd_ret60"].mean(),
        (m["fwd_ret60"]>0).mean()*100, m["is_big"].mean()*100 - base_rate), flush=True)

print("\n=== 大盘短期加速下跌（mkt_ret5）===", flush=True)
scan(sig["mkt_ret5"] > -2, "大盘5日跌>-2%(企稳)")
scan((sig["mkt_ret5"] <= -2) & (sig["mkt_ret5"] > -5), "大盘5日跌-2~-5%")
scan(sig["mkt_ret5"] <= -5, "大盘5日跌<-5%(加速)")

print("\n=== 大盘连续下跌天数 ===", flush=True)
scan(sig["mkt_down_days"] <= 2, "大盘连跌<=2天")
scan((sig["mkt_down_days"] > 2) & (sig["mkt_down_days"] <= 5), "连跌3-5天")
scan(sig["mkt_down_days"] > 5, "连跌>5天")

print("\n=== 消息面气氛（mkt_news5 分位）===", flush=True)
sig["news_q"] = sig["mkt_news5"].rank(pct=True)
scan(sig["news_q"] < 0.33, "新闻量低(冷清)")
scan((sig["news_q"] >= 0.33) & (sig["news_q"] < 0.67), "新闻量中")
scan(sig["news_q"] >= 0.67, "新闻量高(热闹)")

print("\n=== 行业共振（行业是否也在深跌）===", flush=True)
scan(sig["sec_ret20"] < -10, "行业深跌<-10%(共振)")
scan((sig["sec_ret20"] >= -10) & (sig["sec_ret20"] < 0), "行业跌-10~0%")
scan(sig["sec_ret20"] < -15, "行业极深跌<-15%")

print("\n=== 组合：底部确认信号 ===", flush=True)
scan((sig["mkt_ret5"] > -3) & (sig["mkt_down_days"] <= 3), "大盘企稳+连跌<=3")
scan((sig["mkt_ret5"] > -3) & (sig["mkt_down_days"] <= 3) & (sig["news_q"] >= 0.33), "企稳+连跌<=3+新闻中高")
scan((sig["mkt_ret5"] > 0), "大盘5日已转正")

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