import pandas as pd
import numpy as np
from sqlalchemy import create_engine, text
import logging
from datetime import datetime, timedelta

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class VolatilityCalculator:
    def __init__(self, db_connection_string):
        """
        初始化波动率计算器
        :param db_connection_string: 数据库连接字符串，例如: 'mysql+pymysql://user:password@localhost:3306/db_gp_cj'
        """
        self.engine = create_engine(db_connection_string)

    def create_volatility_table(self):
        """创建存储波动率的表"""
        create_table_sql = """
        CREATE TABLE IF NOT EXISTS gp_volatility (
            id INT NOT NULL AUTO_INCREMENT,
            gp_code VARCHAR(255),
            calculation_date DATE,
            volatility_yearly DECIMAL(10, 4),
            update_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            PRIMARY KEY (id),
            INDEX idx_gp_code (gp_code),
            INDEX idx_calc_date (calculation_date)
        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin;
        """
        with self.engine.connect() as conn:
            conn.execute(text(create_table_sql))
            conn.commit()

    def calculate_volatility(self, end_date=None):
        """
        计算所有股票的年化波动率
        :param end_date: 计算截止日期，默认为最新日期
        """
        if end_date is None:
            end_date = datetime.now().date()
        start_date = end_date - timedelta(days=365)

        # 获取所有股票代码
        with self.engine.connect() as conn:
            stocks = pd.read_sql("SELECT DISTINCT gp_code FROM gp_code_all", conn)

        results = []
        for stock_code in stocks['gp_code']:
            try:
                # 获取股票的日线数据
                query = f"""
                SELECT timestamp, close 
                FROM gp_day_data 
                WHERE symbol = '{stock_code}'
                AND timestamp BETWEEN '{start_date}' AND '{end_date}'
                ORDER BY timestamp
                """
                
                with self.engine.connect() as conn:
                    df = pd.read_sql(query, conn)
                
                if len(df) < 30:  # 数据太少，跳过
                    continue

                # 计算日收益率
                df['close'] = df['close'].astype(float)
                df['daily_return'] = df['close'].pct_change()
                
                # 计算年化波动率 (标准差 * sqrt(交易日数量))
                volatility = df['daily_return'].std() * np.sqrt(252)  # 252为一年的交易日数
                
                results.append({
                    'gp_code': stock_code,
                    'calculation_date': end_date,
                    'volatility_yearly': float(volatility)
                })
                
                logger.info(f"Calculated volatility for {stock_code}: {volatility:.4f}")
                
            except Exception as e:
                logger.error(f"Error calculating volatility for {stock_code}: {str(e)}")
                continue

        # 批量插入结果
        if results:
            insert_sql = """
            INSERT INTO gp_volatility (gp_code, calculation_date, volatility_yearly)
            VALUES (:gp_code, :calculation_date, :volatility_yearly)
            """
            with self.engine.connect() as conn:
                conn.execute(text(insert_sql), results)
                conn.commit()

    def get_top_volatile_stocks(self, top_n=10, calculation_date=None):
        """
        获取波动率最高的股票
        :param top_n: 返回前N个股票
        :param calculation_date: 特定日期的数据，默认为最新日期
        :return: DataFrame包含股票代码和波动率
        """
        query = """
        SELECT v.gp_code, c.gp_name, v.volatility_yearly, v.calculation_date
        FROM gp_volatility v
        LEFT JOIN gp_code_all c ON v.gp_code = c.gp_code
        WHERE v.calculation_date = (
            SELECT MAX(calculation_date) FROM gp_volatility
            WHERE calculation_date = COALESCE(:calc_date, calculation_date)
        )
        ORDER BY v.volatility_yearly DESC
        LIMIT :top_n
        """
        with self.engine.connect() as conn:
            return pd.read_sql(text(query), conn, params={'calc_date': calculation_date, 'top_n': top_n})

if __name__ == "__main__":
    # 使用示例
    # calculator = VolatilityCalculator('mysql+pymysql://user:password@localhost:3306/db_gp_cj')
    # calculator.create_volatility_table()
    # calculator.calculate_volatility()
    # top_volatile = calculator.get_top_volatile_stocks(10)
    # print(top_volatile)
    pass 