#!/usr/bin/env python3
"""实盘数据累积系统（MoFin）：每日收盘后累积重点行业/个股的实盘数据。
目标：从当前开始累积实盘数据（资金流、行业热度、龙虎榜代理等），定时回测调整策略。
数据源（MoFin已有）：stock_daily(日线)、stock_capital_flow(资金流)、stock_news(新闻)、sector_index_daily(行业指数)
增量：每日收盘后计算并累积：
  - 重点行业（近期热点）的资金流/新闻热度/行业趋势
  - 重点个股的龙虎榜代理（大单/主力资金）
  - 市场状态（大盘RSI/ADX/行业强弱分布）
输出：/home/hmo/MoFin/data/live_data/ 目录下按日期累积
"""
import sqlite3, json, os, pandas as pd, numpy as np
from datetime import datetime, timedelta

DB = "/home/hmo/MoFin/data/mofin.db"
LIVE_DIR = "/home/hmo/MoFin/data/live_data"
os.makedirs(LIVE_DIR, exist_ok=True)

def collect_daily_snapshot(date_str=None):
    """累积某日的实盘数据快照"""
    conn = sqlite3.connect(f"file:{DB}?mode=ro", uri=True)
    conn.execute("PRAGMA query_only=ON")

    if date_str is None:
        date_str = datetime.now().strftime('%Y-%m-%d')
    print(f"Collecting live data for {date_str}")

    # 1. 大盘状态（上证）
    mkt = pd.read_sql("SELECT date, close FROM sector_index_daily WHERE sector='上证指数' ORDER BY date DESC LIMIT 60", conn)
    mkt['date'] = mkt['date'].astype(str)
    mkt = mkt.sort_values('date').reset_index(drop=True)
    mkt['ret'] = mkt['close'].pct_change()
    mkt['gain'] = mkt['ret'].clip(lower=0)
    mkt['loss'] = (-mkt['ret']).clip(lower=0)
    mkt['ag'] = mkt['gain'].rolling(14).mean()
    mkt['al'] = mkt['loss'].rolling(14).mean()
    mkt['rsi'] = 100 - 100 / (1 + mkt['ag'] / mkt['al'].replace(0, np.nan))
    mkt['rsi'] = mkt['rsi'].fillna(50)
    mkt['ma20'] = mkt['close'].rolling(20).mean()
    latest = mkt.iloc[-1] if len(mkt) else None

    # 2. 行业热度（当日新闻数）
    news = pd.read_sql("SELECT code, date FROM stock_news WHERE date >= ? AND date <= ?", conn,
                       params=[(pd.Timestamp(date_str) - pd.Timedelta(days=5)).strftime('%Y-%m-%d'), date_str])
    news['date'] = news['date'].astype(str).str[:10]
    secmap = dict(conn.execute("SELECT code, sector FROM stock_sectors_em").fetchall())
    news['sector'] = news['code'].map(secmap)
    sector_news = news.groupby('sector').size().reset_index(name='news_count').sort_values('news_count', ascending=False)

    # 3. 行业资金流（当日主力资金净流入）
    flow = pd.read_sql("SELECT code, date, main_net FROM stock_capital_flow WHERE date = ?", conn, params=[date_str])
    flow['sector'] = flow['code'].map(secmap)
    sector_flow = flow.groupby('sector')['main_net'].sum().reset_index().sort_values('main_net', ascending=False)

    # 4. 行业指数（当日涨跌）
    sec_idx = pd.read_sql("SELECT sector, date, close, change_pct FROM sector_index_daily WHERE date = ?", conn, params=[date_str])
    sec_idx['sector'] = sec_idx['sector'].astype(str)

    # 5. 重点行业识别（新闻热度+资金流+行业指数强势）
    merged = sector_news.merge(sector_flow, on='sector', how='outer').merge(
        sec_idx[['sector', 'change_pct']], on='sector', how='outer').fillna(0)
    merged['heat_score'] = merged['news_count'] * 0.5 + merged['main_net'].clip(lower=0) / 1e8 * 0.3 + merged['change_pct'] * 10
    hot_sectors = merged.sort_values('heat_score', ascending=False).head(20)

    # 6. 重点个股（热点行业内的强势股）
    # 用当日涨幅 + 资金流筛选
    stocks = pd.read_sql(f"""
        SELECT sd.code, sd.date, sd.close, sd.open, sd.high, sd.low, sd.volume,
               scf.main_net, scf.main_pct
        FROM stock_daily sd
        LEFT JOIN stock_capital_flow scf ON sd.code = scf.code AND sd.date = scf.date
        WHERE sd.date = ?
    """, conn, params=[date_str])
    stocks['sector'] = stocks['code'].map(secmap)
    # 只保留热点行业的股票
    hot_sector_names = set(hot_sectors['sector'].head(10))
    hot_stocks = stocks[stocks['sector'].isin(hot_sector_names)].copy()
    hot_stocks['ret'] = hot_stocks.groupby('code')['close'].transform(lambda x: x.pct_change())
    hot_stocks = hot_stocks.sort_values(['sector', 'ret'], ascending=[True, False])
    top_stocks = hot_stocks.groupby('sector').head(5)

    snapshot = {
        "date": date_str,
        "market": {
            "close": float(latest['close']) if latest is not None else None,
            "rsi": float(latest['rsi']) if latest is not None and not np.isnan(latest['rsi']) else None,
            "ma20": float(latest['ma20']) if latest is not None and not np.isnan(latest['ma20']) else None,
        },
        "hot_sectors": hot_sectors.to_dict('records'),
        "top_stocks": top_stocks.to_dict('records'),
        "sector_news": sector_news.head(20).to_dict('records'),
        "sector_flow": sector_flow.head(20).to_dict('records'),
    }

    # 保存
    out_path = os.path.join(LIVE_DIR, f"{date_str}.json")
    with open(out_path, 'w', encoding='utf-8') as f:
        json.dump(snapshot, f, ensure_ascii=False, indent=1, default=str)
    print(f"  Saved: {out_path} ({os.path.getsize(out_path)//1024}KB)")
    print(f"  Hot sectors: {hot_sectors['sector'].head(5).tolist()}")
    conn.close()
    return snapshot

if __name__ == '__main__':
    import sys
    date_str = sys.argv[1] if len(sys.argv) > 1 else None
    collect_daily_snapshot(date_str)