#!/usr/bin/env python3
"""
历史全景扫描回溯脚本 — 从指定日期开始，对每个交易日模拟 11:30 和 16:30 两次全市场扫描，
将推荐结果写入 stock_scan_history 表，供回测直接使用。

用法:
  ./venv/bin/python scan_history.py                       # 从 2026-01-01 扫描到今天
  ./venv/bin/python scan_history.py --start 2026-02-01    # 指定起始日
  ./venv/bin/python scan_history.py --end 2026-02-10      # 指定结束日
  ./venv/bin/python scan_history.py --force               # 强制覆盖已扫描日期

说明:
  11:30 扫描: 用 T-1 日 K 线 + T 日 open 模拟中午数据（对应回测 15:00 决策依据）
  16:30 扫描: 用 T 日完整 K 线（对应回测次日 10:00 买入依据）
"""
import sys
import os
import time
import argparse
from datetime import datetime, date, timedelta

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import pandas as pd
import psycopg2
from psycopg2.extras import Json
from config import Config
from services.signal_detector import detect_all_signals
from services.stock_algorithms import compute_recommend

K_DAYS = 120
LOOKBACK = 5
SAVE_BATCH = 500


def get_db_conn():
    return psycopg2.connect(
        host=Config.DB_HOST, port=Config.DB_PORT,
        dbname=Config.DB_NAME, user=Config.DB_USER, password=Config.DB_PASSWORD,
    )


def get_trading_days(conn, start: date, end: date):
    with conn.cursor() as cur:
        cur.execute("""
            SELECT DISTINCT trade_date::date FROM stock_kline_daily
            WHERE trade_date >= %s AND trade_date <= %s
            ORDER BY trade_date
        """, (start, end))
        return [r[0] for r in cur.fetchall()]


def get_scanned_dates(conn, scan_time: str):
    """返回已扫描的日期集合"""
    with conn.cursor() as cur:
        cur.execute("""
            SELECT DISTINCT scan_date FROM stock_scan_history
            WHERE scan_time = %s
        """, (scan_time,))
        return {r[0] for r in cur.fetchall()}


def preload_all_klines(conn):
    """一次性加载全部 K 线到内存: {code: [(date,o,h,l,c,v), ...]}"""
    print("  加载全市场 K 线数据...", flush=True)
    t0 = time.time()
    result = {}
    with conn.cursor() as cur:
        cur.execute("""
            SELECT code, trade_date, open, high, low, close, volume
            FROM stock_kline_daily
            ORDER BY code, trade_date
        """)
        buf_code = None
        buf_rows = []
        for r in cur:
            code = r[0]
            if code != buf_code:
                if buf_code and buf_rows:
                    result[buf_code] = buf_rows
                buf_code = code
                buf_rows = []
            buf_rows.append((str(r[1]), float(r[2]), float(r[3]), float(r[4]), float(r[5]), float(r[6])))
        if buf_code and buf_rows:
            result[buf_code] = buf_rows
    print(f"  加载完成: {len(result)} 只股票, {time.time()-t0:.1f}s", flush=True)
    return result


def build_df(rows, end_date_str: str, days: int = K_DAYS):
    """从预加载行构建 DataFrame（截止到 end_date_str）"""
    filtered = [r for r in rows if r[0] <= end_date_str]
    if len(filtered) < 30:
        return None
    trimmed = filtered[-days:]
    df = pd.DataFrame(trimmed, columns=['date', 'open', 'high', 'low', 'close', 'volume'])
    for col in ('open', 'high', 'low', 'close', 'volume'):
        df[col] = df[col].astype(float)
    return df


def build_noon_df(rows, day_t: date, ohlc_t: dict, code: str, days: int = K_DAYS):
    """构建 11:30 中午 K 线: T-1 前 + T 日 open"""
    prev_str = str(day_t - timedelta(days=1))
    filtered = [r for r in rows if r[0] <= prev_str]
    if len(filtered) < 30:
        return None
    trimmed = filtered[-days:]
    df = pd.DataFrame(trimmed, columns=['date', 'open', 'high', 'low', 'close', 'volume'])
    for col in ('open', 'high', 'low', 'close', 'volume'):
        df[col] = df[col].astype(float)
    if code in ohlc_t:
        open_t = ohlc_t[code][0]
        extra = pd.DataFrame([{
            'date': str(day_t), 'open': open_t, 'high': open_t,
            'low': open_t, 'close': open_t, 'volume': 0.0,
        }])
        return pd.concat([df, extra], ignore_index=True)
    return df


def get_day_ohlc(conn, trade_date: date):
    with conn.cursor() as cur:
        cur.execute("""
            SELECT code, open, close FROM stock_kline_daily WHERE trade_date = %s
        """, (trade_date,))
        return {r[0]: (float(r[1]), float(r[2])) for r in cur.fetchall()}


def scan_one_day(preloaded, codes, day_t, ohlc_t, scan_time, conn):
    """对指定日期的所有股票做一次扫描，返回结果列表。
    scan_time='16:30': 用 T 日完整 K 线
    scan_time='11:30': 用 T-1 + T 日 open 模拟中午
    """
    results = []
    day_str = str(day_t)
    total = len(codes)
    t0 = time.time()

    for j, code in enumerate(codes):
        if scan_time == '16:30':
            df = build_df(preloaded.get(code, []), day_str, K_DAYS)
        else:
            df = build_noon_df(preloaded.get(code, []), day_t, ohlc_t, code, K_DAYS)

        if df is None or len(df) < 30:
            continue

        try:
            res = detect_all_signals(df, lookback=LOOKBACK)
        except Exception:
            continue
        if res.get('error'):
            continue

        signal_status = res.get('signal_status', [])
        indicators = res.get('indicators', {})
        triggered_count = sum(1 for s in signal_status if s.get('triggered'))

        is_holding = False
        st, disp, reason, rate = compute_recommend(
            signal_status, indicators, triggered_count, is_holding=is_holding
        )

        # 同时计算持仓版推荐（回测 15:00 需要两种）
        st_h, disp_h, reason_h, rate_h = compute_recommend(
            signal_status, indicators, triggered_count, is_holding=True
        )

        results.append({
            'code': code,
            'recommend_display': disp,
            'recommend_type': st,
            'recommend_reason': reason,
            'recommend_rate': rate,
            'triggered_count': triggered_count,
            'signal_status': signal_status,
            'indicators': indicators,
            'disp_holding': disp_h,
            'reason_holding': reason_h,
            'rate_holding': rate_h,
        })

        if (j + 1) % 500 == 0:
            elapsed = time.time() - t0
            print(f"      已扫描 {j+1}/{total}  {elapsed:.0f}s", flush=True)

    return results


def save_results(conn, results, scan_date, scan_time):
    """批量写入扫描结果"""
    if not results:
        return
    with conn.cursor() as cur:
        for r in results:
            # 将持仓版推荐也存入 indicators 字段方便回测
            ind = r.get('indicators', {})
            ind['_holding'] = {
                'display': r.get('disp_holding', ''),
                'reason': r.get('reason_holding', ''),
                'rate': r.get('rate_holding', 0),
            }
            cur.execute("""
                INSERT INTO stock_scan_history
                    (scan_date, scan_time, code, recommend_display, recommend_type,
                     recommend_reason, recommend_rate, triggered_count, signal_status, indicators)
                VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                ON CONFLICT (scan_date, scan_time, code) DO UPDATE SET
                    recommend_display = EXCLUDED.recommend_display,
                    recommend_type = EXCLUDED.recommend_type,
                    recommend_reason = EXCLUDED.recommend_reason,
                    recommend_rate = EXCLUDED.recommend_rate,
                    triggered_count = EXCLUDED.triggered_count,
                    signal_status = EXCLUDED.signal_status,
                    indicators = EXCLUDED.indicators,
                    created_at = CURRENT_TIMESTAMP
            """, (
                scan_date, scan_time, r['code'],
                r['recommend_display'], r['recommend_type'],
                r['recommend_reason'], r['recommend_rate'],
                r['triggered_count'],
                Json(r['signal_status']), Json(ind),
            ))
    conn.commit()


def main():
    parser = argparse.ArgumentParser(description='历史全景扫描回溯（11:30 + 16:30）')
    parser.add_argument('--start', type=str, default='2026-01-01', metavar='YYYY-MM-DD')
    parser.add_argument('--end', type=str, default=None, metavar='YYYY-MM-DD')
    parser.add_argument('--force', action='store_true', help='强制覆盖已扫描日期')
    args = parser.parse_args()

    start_date = datetime.strptime(args.start, '%Y-%m-%d').date()
    end_date = datetime.strptime(args.end, '%Y-%m-%d').date() if args.end else date.today()

    conn = get_db_conn()

    print("=" * 70)
    print("  历史全景扫描回溯")
    print("=" * 70)
    print(f"  扫描区间: {start_date} ~ {end_date}")
    print(f"  扫描时段: 11:30（中午）+ 16:30（收盘后）")
    print(f"  强制覆盖: {'是' if args.force else '否（跳过已扫描日期）'}")
    print("-" * 70)

    trading_days = get_trading_days(conn, start_date, end_date)
    if not trading_days:
        print("错误: 无交易日数据")
        conn.close()
        return
    print(f"  交易日数: {len(trading_days)} 天")

    scanned_1130 = get_scanned_dates(conn, '11:30') if not args.force else set()
    scanned_1630 = get_scanned_dates(conn, '16:30') if not args.force else set()

    preloaded = preload_all_klines(conn)
    all_codes = sorted(preloaded.keys())
    print(f"  可扫描股票: {len(all_codes)} 只")
    print("=" * 70)

    total_start = time.time()
    total_scans = 0

    for i, day_t in enumerate(trading_days):
        need_1130 = day_t not in scanned_1130
        need_1630 = day_t not in scanned_1630

        if not need_1130 and not need_1630:
            continue

        ohlc_t = get_day_ohlc(conn, day_t)
        if not ohlc_t:
            continue

        print(f"\n  [{i+1}/{len(trading_days)}] {day_t}", flush=True)

        # 16:30 收盘后扫描（用 T 日完整 K 线）
        if need_1630:
            t0 = time.time()
            print(f"    16:30 扫描中...", flush=True)
            results = scan_one_day(preloaded, all_codes, day_t, ohlc_t, '16:30', conn)
            save_results(conn, results, day_t, '16:30')
            buy_count = sum(1 for r in results if r['recommend_display'] == '买入')
            sell_count = sum(1 for r in results if r['recommend_display'] == '卖出')
            print(f"    16:30 完成: {len(results)} 只  买入推荐 {buy_count}  卖出 {sell_count}  {time.time()-t0:.0f}s",
                  flush=True)
            total_scans += 1

        # 11:30 中午扫描（用 T-1 + T 日 open 模拟）
        if need_1130:
            t0 = time.time()
            print(f"    11:30 扫描中...", flush=True)
            results = scan_one_day(preloaded, all_codes, day_t, ohlc_t, '11:30', conn)
            save_results(conn, results, day_t, '11:30')
            buy_count = sum(1 for r in results if r['recommend_display'] == '买入')
            sell_count = sum(1 for r in results if r['recommend_display'] == '卖出')
            print(f"    11:30 完成: {len(results)} 只  买入推荐 {buy_count}  卖出 {sell_count}  {time.time()-t0:.0f}s",
                  flush=True)
            total_scans += 1

    elapsed = time.time() - total_start
    print(f"\n{'=' * 70}")
    print(f"  全部完成!")
    print(f"  扫描次数: {total_scans} 次（{len(trading_days)} 天 × 2 时段）")
    print(f"  总耗时  : {elapsed/60:.1f} 分钟")
    print(f"{'=' * 70}")
    conn.close()


if __name__ == '__main__':
    main()