import pandas as pd
import numpy as np
import os
from src.strategy.prepare import fetch_his_kline
import src.config.config as config
import src.crawler.em.stock as em_stock
from  MyTT import *              #myTT麦语言工具函数指标库


default_rsi_dir = f"{config.global_stock_data_dir}/em_strategy"

def select_stocks(stock_map):
    all_stock_info = []
    for stock_code, stock_name in stock_map.items():
        df = fetch_his_kline(stock_code)
        close_prices = df['close'].values
        
        # 使用 MyTT 库计算不同周期的 RSI
        rsi_6 = RSI(close_prices, 6)
        rsi_12 = RSI(close_prices, 12)
        rsi_24 = RSI(close_prices, 24)

        df['rsi_6'] = rsi_6
        df['rsi_12'] = rsi_12
        df['rsi_24'] = rsi_24

        # 生成买入和卖出信号
        df['buy_signal'] = ((df['rsi_6'] <= 30) & (df['rsi_24'] > 50)) | (df['rsi_6'] <=  15)
        df['sell_signal'] = (df['rsi_6'] > 70) & (df['rsi_24'] > 50)

        # 获取最大日期的数据
        latest_data = df[df['date'] == df['date'].max()]
        if latest_data['buy_signal'].values[0]:
            print(f"current stock {stock_code} {stock_name} has buy signal.")
        else:
            continue

        # 过滤出有买入信号的记录
        buy_dates = df[df['buy_signal']]['date'].tolist()

        if buy_dates:
            max_gains = []
            for buy_date in buy_dates:
                buy_index = df[df['date'] == buy_date].index[0]
                best_max_gain = -np.inf
                for holding_period in range(7, 32):
                    end_index = min(buy_index + holding_period, len(df) - 1)
                    buy_price = df.loc[buy_index, 'close']
                    future_prices = df.loc[buy_index + 1:end_index, 'close']
                    if not future_prices.empty:
                        max_price = future_prices.max()
                        gain = (max_price - buy_price) / buy_price
                        if gain > best_max_gain:
                            best_max_gain = gain
                max_gains.append(best_max_gain)

            # 评估策略
            feasible_count = sum([gain > 0 for gain in max_gains])
            excellent_count = sum([gain > 0.05 for gain in max_gains])
            total_buys = len(max_gains)

            if total_buys > 0:
                feasible_ratio = feasible_count / total_buys
                excellent_ratio = excellent_count / total_buys
            else:
                feasible_ratio = 0
                excellent_ratio = 0

            stock_info = {
                '股票代码': stock_code,
                '股票名称': stock_name,
                '赢率': feasible_ratio,
                '优秀率': excellent_ratio,
                '买入信号次数': total_buys
            }
            print(stock_info)
            print('\n\n')
            all_stock_info.append(stock_info)

    # 筛选出赢率较高和优秀率较高的股票
    result_df = pd.DataFrame(all_stock_info)
    recommended_stocks = result_df[(result_df['赢率'] > 0.5) & (result_df['优秀率'] > 0.3)]

    print("推荐买入股：")
    print(recommended_stocks)

    return recommended_stocks

    # 过滤出有买入或卖出信号的记录
    # selected_df = df[df['buy_signal'] | df['sell_signal']]
    #return selected_df


if __name__ == "__main__":
    os.makedirs(default_rsi_dir, exist_ok=True)

    codes = ['105.QFIN', '105.FUTU']
    # 从网络上获取
    stock_map = em_stock.code_by_fs('hk_famous', em_stock.em_market_fs_types['hk_famous'])
    if stock_map:
        select_stocks(stock_map)

    for code in []:
        print(f"caculate stock {code}")
        result = select_stocks(code)
        full_path = f"{default_rsi_dir}/{code}.csv"
        result.to_csv(full_path, index=False, encoding='utf-8')
        print(f'task complete! see result in {full_path}\n\n')