import pandas as pd
import numpy as np
import pymysql
from datetime import datetime, timedelta
from tqdm import tqdm
import matplotlib.pyplot as plt
import os

# 添加中文字体支持
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

class StockBacktester:
    def __init__(self, db_config):
        """
        初始化股票回测器
        :param db_config: 数据库配置字典
        """
        self.db_config = db_config
        # 测试连接
        try:
            conn = pymysql.connect(
                host=self.db_config['host'],
                port=self.db_config['port'],
                user=self.db_config['user'],
                password=self.db_config['password'],
                database=self.db_config['database']
            )
            conn.close()
            print("数据库连接成功")
        except Exception as e:
            print(f"数据库连接错误: {e}")
            raise
            
        self.initial_capital = 300000  # 30万本金
        self.total_position = 1000000  # 100万建仓规模
        self.borrowed_capital = 700000  # 70万借入资金
        self.annual_interest_rate = 0.05  # 年化5%利息
        self.daily_interest_rate = self.annual_interest_rate / 365
        self.commission_rate = 0.05  # 5%手续费
        self.min_holding_days = 7  # 最少持仓天数
        self.take_profit_pct = 0.20  # 固定20%止盈

    def get_stock_data(self, symbol, start_date, end_date):
        """
        获取指定股票从起始日期到结束日期的数据
        :param symbol: 股票代码
        :param start_date: 起始日期
        :param end_date: 结束日期
        :return: 股票数据DataFrame
        """
        try:
            conn = pymysql.connect(
                host=self.db_config['host'],
                port=self.db_config['port'],
                user=self.db_config['user'],
                password=self.db_config['password'],
                database=self.db_config['database']
            )
            
            query = f"""
            SELECT 
                symbol,
                timestamp,
                CAST(open as DECIMAL(10,2)) as open,
                CAST(high as DECIMAL(10,2)) as high,
                CAST(low as DECIMAL(10,2)) as low,
                CAST(close as DECIMAL(10,2)) as close,
                CAST(volume as DECIMAL(20,0)) as volume
            FROM gp_day_data 
            WHERE symbol = %s 
            AND timestamp BETWEEN %s AND %s
            ORDER BY timestamp
            """
            
            # 使用pandas读取sql
            df = pd.read_sql(query, conn, params=(symbol, start_date, end_date))
            conn.close()
            
            # 确保数值列的类型正确
            numeric_columns = ['open', 'high', 'low', 'close', 'volume']
            for col in numeric_columns:
                if col in df.columns:
                    df[col] = pd.to_numeric(df[col], errors='coerce')
            
            # 删除任何包含 NaN 的行
            df = df.dropna()
            
            if df.empty:
                print(f"股票 {symbol} 在指定时间范围内没有数据")
            else:
                print(f"成功获取股票 {symbol} 数据，共 {len(df)} 条记录")
                
            return df
            
        except Exception as e:
            print(f"获取股票 {symbol} 数据时出错: {str(e)}")
            return pd.DataFrame()

    def backtest_stock(self, symbol, buy_dates, end_date):
        """
        回测单个股票的表现
        :param symbol: 股票代码
        :param buy_dates: 买入日期列表，格式为['YYYY-MM-DD', 'YYYY-MM-DD', ...]
        :param end_date: 结束日期字符串，格式为'YYYY-MM-DD'
        :return: 回测结果
        """
        # 找到最早的买入日期
        buy_dates = sorted(buy_dates)
        earliest_buy_date = buy_dates[0]
        
        # 获取股票数据
        df = self.get_stock_data(symbol, earliest_buy_date, end_date)
        
        if df.empty:
            print(f"股票 {symbol} 在指定时间范围内没有数据")
            return None, None
            
        if len(df) < 2:
            print(f"股票 {symbol} 在指定时间范围内数据不足")
            return None, None
            
        # 将时间戳转换为datetime
        df['timestamp'] = pd.to_datetime(df['timestamp'])
        
        # 转换买入日期和结束日期为datetime对象
        buy_dates_obj = [datetime.strptime(buy_date, '%Y-%m-%d') for buy_date in buy_dates]
        end_date_obj = datetime.strptime(end_date, '%Y-%m-%d')
        
        # 初始化结果列表
        results = []
        positions = []  # 记录所有持仓
        cumulative_profit = 0  # 总累计收益
        
        # 记录统计变量
        total_holding_days = 0
        total_trades = 0
        profitable_trades = 0
        loss_trades = 0
        entry_count = 0  # 成功建仓次数
        
        # 创建时间戳到索引的映射，便于查找
        timestamp_to_idx = {row['timestamp']: idx for idx, row in df.iterrows()}
        
        # 为每个买入日期找到对应的交易日并建仓
        for buy_date_obj in buy_dates_obj:
            # 找到买入日期当天或之后的第一个交易日
            buy_idx = -1
            nearest_date = None
            
            for date in df['timestamp']:
                if date >= buy_date_obj:
                    buy_idx = timestamp_to_idx[date]
                    nearest_date = date
                    break
            
            if buy_idx == -1:
                print(f"股票 {symbol} 在买入日期 {buy_date_obj.strftime('%Y-%m-%d')} 之后没有交易数据")
                continue
            
            entry_row = df.iloc[buy_idx]
            entry_price = float(entry_row['open'])
            entry_date = entry_row['timestamp']
            
            positions.append({
                'entry_price': entry_price,
                'entry_date': entry_date,
                'closed': False
            })
            entry_count += 1
            
            print(f"建仓 #{entry_count}: {symbol} 在 {entry_date.strftime('%Y-%m-%d')} 以 {entry_price:.2f} 价格买入")
        
        # 创建初始结果
        if len(df) > 0:
            first_day = df.iloc[0]
            results.append({
                'timestamp': first_day['timestamp'],
                'profit': 0,
                'position': len(positions) > 0,
                'action': 'hold'
            })
        
        # 遍历所有交易日，检查平仓条件
        for i in range(1, len(df)):
            current_row = df.iloc[i]
            current_date = current_row['timestamp']
            current_price = (float(current_row['high']) + float(current_row['low'])) / 2
            
            day_result = {
                'timestamp': current_date,
                'profit': cumulative_profit,
                'position': False,
                'action': 'hold'
            }
            
            # 检查是否是买入日期，设置持仓标记
            if any(not pos['closed'] for pos in positions):
                day_result['position'] = True
            
            # 检查所有持仓是否需要平仓
            positions_to_close = []
            
            for idx, pos in enumerate(positions):
                if pos['closed']:
                    continue
                    
                pos_days_held = (current_date - pos['entry_date']).days
                pos_size = self.total_position / pos['entry_price']
                
                # 检查是否需要平仓
                should_close = False
                close_type = None
                close_price = current_price
                
                if pos_days_held >= self.min_holding_days:
                    # 检查止盈
                    if float(current_row['high']) >= pos['entry_price'] * (1 + self.take_profit_pct):
                        should_close = True
                        close_type = 'sell_profit'
                        close_price = pos['entry_price'] * (1 + self.take_profit_pct)
                
                # 强制在结束日期平仓
                if current_date >= end_date_obj and not should_close:
                    should_close = True
                    close_type = 'final_sell'
                
                if should_close:
                    profit = (close_price - pos['entry_price']) * pos_size
                    commission = profit * self.commission_rate if profit > 0 else 0
                    interest_cost = self.borrowed_capital * self.daily_interest_rate * pos_days_held
                    net_profit = profit - commission - interest_cost
                    
                    # 更新统计数据
                    total_holding_days += pos_days_held
                    total_trades += 1
                    if net_profit > 0:
                        profitable_trades += 1
                    else:
                        loss_trades += 1
                    
                    cumulative_profit += net_profit
                    
                    positions[idx]['closed'] = True
                    positions_to_close.append((idx, pos, close_type, close_price, net_profit, pos_days_held))
            
            # 处理平仓
            for idx, pos, close_type, close_price, net_profit, pos_days_held in positions_to_close:
                day_result['action'] = close_type
                day_result['profit'] = cumulative_profit
                
                print(f"平仓 #{idx+1}: {symbol} 在 {current_date.strftime('%Y-%m-%d')} 以 {close_price:.2f} 价格卖出")
                print(f"持仓天数: {pos_days_held}, 净利润: {net_profit:.2f}")
            
            # 更新日结果
            day_result['position'] = any(not pos['closed'] for pos in positions)
            results.append(day_result)
        
        # 计算统计数据
        avg_holding_days = total_holding_days / total_trades if total_trades > 0 else 0
        win_rate = profitable_trades / total_trades * 100 if total_trades > 0 else 0
        
        # 统计信息
        stats = {
            'symbol': symbol,
            'total_trades': total_trades,
            'profitable_trades': profitable_trades,
            'loss_trades': loss_trades,
            'win_rate': win_rate,
            'avg_holding_days': avg_holding_days,
            'final_profit': cumulative_profit,
            'entry_count': entry_count
        }
        
        print(f"\n交易统计 - {symbol}")
        print(f"总交易次数: {total_trades}")
        print(f"盈利交易: {profitable_trades}次")
        print(f"亏损交易: {loss_trades}次")
        print(f"胜率: {win_rate:.1f}%")
        print(f"平均持仓天数: {avg_holding_days:.1f}")
        print(f"累计收益: {cumulative_profit:.2f}")
        print(f"建仓次数: {entry_count}\n")
        
        return pd.DataFrame(results), stats

    def backtest_multiple_stocks(self, stocks_buy_dates, end_date):
        """
        回测多只股票
        :param stocks_buy_dates: 字典，键为股票代码，值为买入日期列表
        :param end_date: 所有股票的共同结束日期
        :return: 回测结果字典
        """
        results = {}
        stats_list = []
        
        print(f"\n开始回测，共 {len(stocks_buy_dates)} 支股票，结束日期: {end_date}")
        
        for symbol, buy_dates in tqdm(stocks_buy_dates.items(), desc="回测进度"):
            try:
                result_df, stats = self.backtest_stock(symbol, buy_dates, end_date)
                
                if result_df is not None and stats is not None:
                    results[symbol] = result_df
                    stats_list.append(stats)
                    
            except Exception as e:
                print(f"\n处理股票 {symbol} 时出错: {str(e)}")
                import traceback
                traceback.print_exc()
                continue
        
        # 保存统计结果
        if stats_list:
            stats_df = pd.DataFrame(stats_list)
            
            if not os.path.exists('results'):
                os.makedirs('results')
                
            stats_df.to_csv('results/backtest_summary.csv', index=False)
            print(f"\n回测统计已保存至 results/backtest_summary.csv")
            
            # 计算总体统计
            total_profit = stats_df['final_profit'].sum()
            avg_win_rate = stats_df['win_rate'].mean()
            avg_holding_days = stats_df['avg_holding_days'].mean()
            total_trades = stats_df['total_trades'].sum()
            profitable_stocks = len(stats_df[stats_df['final_profit'] > 0])
            total_entries = stats_df['entry_count'].sum()
            
            print(f"\n总体回测统计:")
            print(f"总收益: {total_profit:.2f}")
            print(f"平均胜率: {avg_win_rate:.1f}%")
            print(f"平均持仓天数: {avg_holding_days:.1f}")
            print(f"总交易次数: {total_trades}")
            print(f"总建仓次数: {total_entries}")
            print(f"盈利股票数: {profitable_stocks}/{len(stats_df)}")
        
        return results, stats_list

    def plot_results(self, results, symbol):
        """
        绘制单个股票的收益走势图
        """
        if results is None or results.empty:
            print(f"股票 {symbol} 没有可绘制的结果")
            return None
            
        plt.figure(figsize=(12, 6))
        plt.plot(results['timestamp'], results['profit'], label='累计收益')
        plt.title(f'股票 {symbol} 回测结果')
        plt.xlabel('日期')
        plt.ylabel('收益 (元)')
        plt.legend()
        plt.grid(True)
        plt.xticks(rotation=45)
        plt.tight_layout()
        
        # 保存图表
        if not os.path.exists('results/charts'):
            os.makedirs('results/charts')
        plt.savefig(f'results/charts/{symbol}_backtest.png')
        
        return plt

    def plot_all_stocks(self, all_results, stats_list):
        """
        绘制所有股票的收益走势图
        """
        if not all_results:
            print("没有可绘制的结果")
            return None, {}
            
        plt.figure(figsize=(15, 8))
        
        # 记录所有股票的收益
        total_profits = {}
        max_profit = float('-inf')
        min_profit = float('inf')
        
        # 绘制每支股票的曲线
        for symbol, results in all_results.items():
            if results is None or results.empty:
                continue
                
            profits = results['profit']
            plt.plot(results['timestamp'], profits, label=symbol, alpha=0.5, linewidth=1)
            
            # 更新最大最小收益
            if not profits.empty:
                max_profit = max(max_profit, profits.max())
                min_profit = min(min_profit, profits.min())
            
            # 记录总收益
            if not profits.empty:
                total_profits[symbol] = profits.iloc[-1]
        
        # 计算统计信息
        if stats_list:
            stats_df = pd.DataFrame(stats_list)
            total_profit = stats_df['final_profit'].sum()
            avg_win_rate = stats_df['win_rate'].mean()
            avg_holding_days = stats_df['avg_holding_days'].mean()
            total_trades = stats_df['total_trades'].sum()
            profitable_stocks = len(stats_df[stats_df['final_profit'] > 0])
            total_entries = stats_df['entry_count'].sum()
            
            # 添加统计信息
            stats_text = (
                f"回测统计:\n"
                f"总收益: {total_profit:.2f}\n"
                f"平均胜率: {avg_win_rate:.1f}%\n"
                f"平均持仓天数: {avg_holding_days:.1f}\n"
                f"总交易次数: {total_trades}\n"
                f"总建仓次数: {total_entries}\n"
                f"盈利股票数: {profitable_stocks}/{len(stats_df)}"
            )
            
            plt.text(0.02, 0.95, stats_text,
                    transform=plt.gca().transAxes,
                    bbox=dict(facecolor='white', alpha=0.8, edgecolor='gray'),
                    verticalalignment='top',
                    fontsize=10)
        
        plt.title('所有股票回测结果')
        plt.xlabel('日期')
        plt.ylabel('收益 (元)')
        plt.grid(True)
        plt.xticks(rotation=45)
        
        # 保存图表
        if not os.path.exists('results/charts'):
            os.makedirs('results/charts')
        plt.savefig('results/charts/all_stocks_backtest.png')
        plt.tight_layout()
        
        return plt, total_profits

def run_backtest(stocks_buy_dates, end_date):
    """
    运行回测的便捷方法
    :param stocks_buy_dates: 字典，键为股票代码，值为买入日期列表，
                         例如 {'SH600522': ['2022-05-10', '2022-06-10'], 'SZ002340': ['2022-06-15']}
    :param end_date: 所有股票共同的结束日期，格式为'YYYY-MM-DD'
    :return: 回测结果
    """
    # 数据库连接配置
    db_config = {
        'host': '192.168.18.199',
        'port': 3306,
        'user': 'root',
        'password': 'Chlry#$.8',
        'database': 'db_gp_cj'
    }
    
    try:
        # 创建results目录
        if not os.path.exists('results'):
            os.makedirs('results')
        
        # 创建回测器实例
        backtester = StockBacktester(db_config)
        
        # 运行回测
        results, stats_list = backtester.backtest_multiple_stocks(stocks_buy_dates, end_date)
        
        # 绘制图表
        if results:
            # 绘制所有股票的图表
            backtester.plot_all_stocks(results, stats_list)
            
            # 绘制每支股票的单独图表
            for symbol, result in results.items():
                backtester.plot_results(result, symbol)
        
        return results, stats_list
        
    except Exception as e:
        print(f"回测运行出错: {str(e)}")
        import traceback
        traceback.print_exc()
        return None, None

if __name__ == "__main__":
    # 测试回测 - 现在使用股票对应多个买入日期
    stocks_buy_dates = {
        'SH600522': ['2022-05-10', '2022-06-10'],  # 两个买入日期
        'SZ002340': ['2022-06-15'],                # 一个买入日期
        'SH601615': ['2022-07-20', '2022-08-01']   # 两个买入日期
    }
    end_date = '2022-10-20'  # 所有股票共同的结束日期
    
    run_backtest(stocks_buy_dates, end_date) 