import pandas as pd
import numpy as np
from sqlalchemy import create_engine, text


def init_db(engine):
    """
    初始化数据库表结构
    """
    create_table_sql = """
    CREATE TABLE IF NOT EXISTS `stock_limit_data` (
        `id` BIGINT AUTO_INCREMENT PRIMARY KEY,
        `trade_date` DATE NOT NULL COMMENT '交易日期',
        `symbol` VARCHAR(10) NOT NULL COMMENT '股票代码',
        `ts_code` VARCHAR(20) NOT NULL COMMENT 'TS股票代码',
        `name` VARCHAR(50) NOT NULL COMMENT '股票名称',
        `status` TINYINT NOT NULL COMMENT '状态：1-涨停，2-炸板',
        `limit_price` DECIMAL(16, 4) COMMENT '计算出的理论涨停价',
        `close` DECIMAL(16, 4) COMMENT '实际收盘价',
        `market` VARCHAR(20) COMMENT '所属板块',
        UNIQUE KEY `uniq_date_code` (`trade_date`, `ts_code`)
    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
    """
    with engine.begin() as conn:
        conn.execute(text(create_table_sql))


def analyze_and_save_limits(start_date, end_date):
    """
    计算并在数据库中落地涨停、炸板详情
    """
    # 数据库连接配置
    engine = create_engine(
        "mysql+pymysql://yangfan@8.159.129.156:10836/level?charset=utf8mb4",
        connect_args={
            "password": "aA%8023321088"
        },
        pool_pre_ping=True
    )

    init_db(engine)

    # 1. 获取行情数据
    query = f"""
    SELECT 
        d.trade_date, d.symbol, d.ts_code, b.name, 
        d.high, d.close, d.pre_close, b.market
    FROM stock_daily d
    JOIN stock_basic b ON d.ts_code = b.ts_code
    WHERE d.trade_date >= '{start_date}' 
      AND d.trade_date <= '{end_date}'
      AND d.pre_close IS NOT NULL
    ORDER BY d.trade_date ASC
    """

    try:
        df = pd.read_sql(query, engine)
        if df.empty:
            print(f"⚠️ 区间 [{start_date} , {end_date}] 内未获取到数据。")
            return

        # 强制类型转换
        df['pre_close'] = pd.to_numeric(df['pre_close'], errors='coerce')
        df['high'] = pd.to_numeric(df['high'], errors='coerce')
        df['close'] = pd.to_numeric(df['close'], errors='coerce')

        # 2. 涨停价计算规则 (已取消 ST 判定)
        def calc_limit(row):
            m = str(row['market'])
            p = row['pre_close']

            if pd.isna(p) or p <= 0: return 0.0

            # 仅保留板块判定
            if '创业板' in m or '科创板' in m:
                ratio = 1.2
            elif '北交所' in m:
                ratio = 1.3
            else:
                ratio = 1.1  # 普通主板（含ST股现在也按10%计算）

            # 采用标准四舍五入
            return round(p * ratio, 2)

        df['limit_price'] = df.apply(calc_limit, axis=1)

        # 3. 筛选逻辑：利用 abs 差值判断，解决浮点数精度问题
        # 只要最高价触及理论涨停价（误差0.01以内）即视为触板
        is_at_limit = (df['high'] - df['limit_price']).abs() < 0.011
        is_close_at_limit = (df['close'] - df['limit_price']).abs() < 0.011

        df['is_limit'] = is_at_limit & is_close_at_limit
        df['is_broken'] = is_at_limit & (~is_close_at_limit)

        # 分离结果
        limit_df = df[df['is_limit']].copy()
        limit_df['status'] = 1

        broken_df = df[df['is_broken']].copy()
        broken_df['status'] = 2

        save_df = pd.concat([limit_df, broken_df])

        if not save_df.empty:
            final_df = save_df[[
                'trade_date', 'symbol', 'ts_code', 'name',
                'status', 'limit_price', 'close', 'market'
            ]]

            # 4. 事务入库：先删后插
            with engine.begin() as conn:
                target_dates = final_df['trade_date'].astype(str).unique()
                date_list = ", ".join([f"'{d}'" for d in target_dates])

                delete_query = f"DELETE FROM stock_limit_data WHERE trade_date IN ({date_list})"
                conn.execute(text(delete_query))

                final_df.to_sql('stock_limit_data', conn, if_exists='append', index=False)

            print(f"✅ 处理成功！涉及交易日: {len(target_dates)} 天")
            print(f"📊 记录详情: 总计 {len(final_df)} 条 (涨停: {len(limit_df)} / 炸板: {len(broken_df)})")
        else:
            print("ℹ️ 该时段内全市场无涨停或炸板记录。")

    except Exception as e:
        print(f"❌ 运行失败: {e}")


if __name__ == "__main__":
    # 配置日期区间
    START_DATE = '2024-09-23'
    END_DATE = '2026-04-27'

    analyze_and_save_limits(START_DATE, END_DATE)