# coding:utf-8

import pandas as pd
from sqlalchemy import create_engine
from datetime import datetime

class StockFilter:
    """股票筛选器类，用于根据特定条件筛选股票"""

    def __init__(self, db_url):
        """
        初始化股票筛选器

        Parameters:
        -----------
        db_url : str
            数据库连接URL
        """
        self.engine = create_engine(db_url)

    def get_high_amplitude_stocks(self, date, amplitude_threshold=0.07):
        """
        获取指定日期振幅超过阈值且收盘价低于开盘价的股票

        Parameters:
        -----------
        date : str
            日期字符串，格式为'YYYY-MM-DD'
        amplitude_threshold : float
            振幅阈值，默认为0.07（7%）

        Returns:
        --------
        pandas.DataFrame
            符合条件的股票数据，包含以下列：
            - symbol: 股票代码
            - high: 最高价
            - low: 最低价
            - open: 开盘价
            - close: 收盘价
        """
        query = f"""
        WITH ranked_data AS (
            SELECT symbol, timestamp, high, low, open, close,
                   ROW_NUMBER() OVER (PARTITION BY symbol ORDER BY timestamp) AS row_num,
                   ROW_NUMBER() OVER (PARTITION BY symbol ORDER BY timestamp DESC) AS rev_row_num
            FROM gp_min_data
            WHERE DATE(timestamp) = '{date}'
        ),
        daily_data AS (
            SELECT symbol,
                   MAX(high) AS high,
                   MIN(low) AS low,
                   MAX(CASE WHEN row_num = 1 THEN open ELSE NULL END) AS open,
                   MAX(CASE WHEN rev_row_num = 1 THEN close ELSE NULL END) AS close
            FROM ranked_data
            GROUP BY symbol
        )
        SELECT symbol, high, low, open, close
        FROM daily_data
        WHERE (high - low) / open > {amplitude_threshold}
          AND close < open
        """
        try:
            df = pd.read_sql(query, self.engine)
            return df
        except Exception as e:
            print(f"查询数据时发生错误: {str(e)}")
            return pd.DataFrame()

    # def save_filtered_stocks(self, date, table_name='filtered_stocks'):
    #     """
    #     将筛选结果保存到数据库

    #     Parameters:
    #     -----------
    #     date : str
    #         日期字符串，格式为'YYYY-MM-DD'
    #     table_name : str
    #         保存结果的表名，默认为'filtered_stocks'

    #     Returns:
    #     --------
    #     bool
    #         保存成功返回True，失败返回False
    #     """
    #     try:
    #         # 获取筛选结果
    #         df = self.get_high_amplitude_stocks(date)
            
    #         # 添加筛选日期列
    #         df['filter_date'] = date
            
    #         # 保存到数据库
    #         df.to_sql(table_name, self.engine, if_exists='append', index=False)
            
    #         print(f"数据已成功保存到表 '{table_name}'")
    #         return True
    #     except Exception as e:
    #         print(f"保存数据时发生错误: {str(e)}")
    #         return False

def get_todays_filtered_stocks(db_url):
    """
    获取今日符合条件的股票的快捷方法

    Parameters:
    -----------
    db_url : str
        数据库连接URL

    Returns:
    --------
    pandas.DataFrame
        今日符合条件的股票数据
    """
    # 获取今日日期
    today = datetime.now().strftime('%Y-%m-%d')
    
    # 创建筛选器实例
    filter = StockFilter(db_url)
    
    # 获取并返回结果
    return filter.get_high_amplitude_stocks(today)

if __name__ == "__main__":
    # 示例调用
    db_url = 'mysql+pymysql://root:Chlry#$.8@192.168.18.199:3306/db_gp_cj'
    
    # 方法1：使用快捷函数获取今日数据
    # today_stocks = get_todays_filtered_stocks(db_url)
    # print("今日符合条件的股票：")
    # print(today_stocks)
    
    # 方法2：使用完整的类，指定日期
    filter = StockFilter(db_url)
    stocks = filter.get_high_amplitude_stocks('2024-06-19')
    print("\n指定日期符合条件的股票：")
    print(stocks)
    
    # 保存结果到数据库  这个是AI自动加的，暂时不需要保存数据库
    # filter.save_filtered_stocks('2024-06-19') 