import pandas as pd
import numpy as np
from sqlalchemy import create_engine
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 StockAnalyzer:
    def __init__(self, db_connection_string):
        """
        Initialize the stock analyzer
        :param db_connection_string: Database connection string
        """
        self.engine = create_engine(db_connection_string)
        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  # 最少持仓天数

    def get_stock_list(self):
        """获取所有股票列表"""
        query = "SELECT DISTINCT gp_code as symbol FROM gp_code_all where mark1 = '1'"
        return pd.read_sql(query, self.engine)['symbol'].tolist()

    def get_stock_data(self, symbol, start_date, end_date, include_history=False):
        """获取指定股票在时间范围内的数据"""
        if include_history:
            # 如果需要历史数据，获取start_date之前240天的数据
            query = f"""
            WITH history_data AS (
                SELECT 
                    symbol,
                    timestamp,
                    CAST(close as DECIMAL(10,2)) as close
                FROM gp_day_data 
                WHERE symbol = '{symbol}' 
                AND timestamp < '{start_date}'
                ORDER BY timestamp DESC
                LIMIT 240
            )
            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 = '{symbol}' 
            AND timestamp BETWEEN '{start_date}' AND '{end_date}'
            UNION ALL
            SELECT 
                symbol,
                timestamp,
                NULL as open,
                NULL as high,
                NULL as low,
                close,
                NULL as volume
            FROM history_data
            ORDER BY timestamp
            """
        else:
            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 = '{symbol}' 
            AND timestamp BETWEEN '{start_date}' AND '{end_date}'
            ORDER BY timestamp
            """
        df = pd.read_sql(query, self.engine)
        
        # 确保数值列的类型正确
        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 的行，但只考虑分析期间需要的列
        if include_history:
            # 历史数据只需要close列不为空
            historical_mask = (df['timestamp'] < start_date) & df['close'].notna()
            analysis_mask = (df['timestamp'] >= start_date) & df[['open', 'high', 'low', 'close', 'volume']].notna().all(axis=1)
            df = df[historical_mask | analysis_mask]
        else:
            df = df.dropna()
        
        return df

    def calculate_trade_signals(self, df, take_profit_pct, stop_loss_pct, pullback_entry_pct, start_date):
        """计算交易信号和收益"""
        # 检查历史数据是否足够
        historical_data = df[df['timestamp'] < start_date]
        if len(historical_data) < 240:
            print(f"历史数据不足240天，跳过分析")
            return pd.DataFrame()

        # 计算240日均线基准价格（使用历史数据）
        ma240_base_price = historical_data['close'].tail(240).mean()
        print(f"240日均线基准价格: {ma240_base_price:.2f}")

        # 只使用分析期间的数据进行交易信号计算
        analysis_df = df[df['timestamp'] >= start_date].copy()
        if len(analysis_df) < 10:  # 确保分析期间有足够的数据
            print(f"分析期间数据不足10天，跳过分析")
            return pd.DataFrame()

        results = []
        positions = []  # 记录所有持仓，每个元素是一个字典，包含入场价格和日期
        max_positions = 5  # 最大建仓次数
        cumulative_profit = 0  # 总累计收益
        trades_count = 0
        max_holding_period = timedelta(days=180)  # 最长持仓6个月
        max_loss_amount = -300000  # 最大亏损金额为30万
        
        # 添加统计变量
        total_holding_days = 0  # 总持仓天数
        max_capital_usage = 0  # 最大资金占用
        total_trades = 0  # 总交易次数
        profitable_trades = 0  # 盈利交易次数
        loss_trades = 0  # 亏损交易次数

        print(f"\n{'='*50}")
        print(f"开始分析股票: {df.iloc[0]['symbol']}")
        print(f"{'='*50}")

        # 记录前一天的收盘价
        prev_close = None

        for index, row in analysis_df.iterrows():
            current_price = (float(row['high']) + float(row['low'])) / 2
            day_result = {
                'timestamp': row['timestamp'],
                'profit': cumulative_profit,  # 记录当前的累计收益
                'position': False,
                'action': 'hold'
            }

            # 更新所有持仓的收益
            positions_to_remove = []

            for pos in positions:
                pos_days_held = (row['timestamp'] - 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(row['high']) >= pos['entry_price'] * (1 + take_profit_pct):
                        should_close = True
                        close_type = 'sell_profit'
                        close_price = pos['entry_price'] * (1 + take_profit_pct)
                    # 检查最大亏损
                    elif (close_price - pos['entry_price']) * pos_size <= max_loss_amount:
                        should_close = True
                        close_type = 'sell_loss'
                    # 检查持仓时间
                    elif (row['timestamp'] - pos['entry_date']) >= max_holding_period:
                        should_close = True
                        close_type = 'time_limit'

                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_to_remove.append(pos)
                    day_result['action'] = close_type
                    day_result['profit'] = cumulative_profit
                    
                    print(f"\n>>> {'止盈平仓' if close_type == 'sell_profit' else '止损平仓' if close_type == 'sell_loss' else '到期平仓'}")
                    print(f"持仓天数: {pos_days_held}")
                    print(f"持仓数量: {pos_size:.0f}股")
                    print(f"平仓价格: {close_price:.2f}")
                    print(f"{'盈利' if profit > 0 else '亏损'}: {profit:,.2f}")
                    print(f"手续费: {commission:,.2f}")
                    print(f"利息成本: {interest_cost:,.2f}")
                    print(f"净{'盈利' if net_profit > 0 else '亏损'}: {net_profit:,.2f}")

            # 移除已平仓的持仓
            for pos in positions_to_remove:
                positions.remove(pos)

            # 检查是否可以建仓
            if len(positions) < max_positions and prev_close is not None:
                # 计算当日跌幅
                daily_drop = (prev_close - float(row['low'])) / prev_close
                
                # 检查价格是否在均线区间内
                price_in_range = (
                    current_price >= ma240_base_price * 0.7 and 
                    current_price <= ma240_base_price * 1.3
                )
                
                # 如果跌幅超过5%且价格在均线区间内
                if daily_drop >= 0.05 and price_in_range:
                    positions.append({
                        'entry_price': current_price,
                        'entry_date': row['timestamp']
                    })
                    trades_count += 1
                    print(f"\n>>> 建仓信号 #{trades_count}")
                    print(f"日期: {row['timestamp'].strftime('%Y-%m-%d')}")
                    print(f"建仓价格: {current_price:.2f}")
                    print(f"建仓数量: {self.total_position/current_price:.0f}股")
                    print(f"当日跌幅: {daily_drop*100:.2f}%")
                    print(f"距离均线: {((current_price/ma240_base_price)-1)*100:.2f}%")

            # 更新前一天收盘价
            prev_close = float(row['close'])

            # 更新日结果
            day_result['position'] = len(positions) > 0
            results.append(day_result)

            # 计算当前资金占用
            current_capital_usage = sum(self.total_position for pos in positions)
            max_capital_usage = max(max_capital_usage, current_capital_usage)

        # 在最后一天强制平仓所有持仓
        if positions:
            final_price = (float(analysis_df.iloc[-1]['high']) + float(analysis_df.iloc[-1]['low'])) / 2
            final_total_profit = 0
            
            for pos in positions:
                position_size = self.total_position / pos['entry_price']
                days_held = (analysis_df.iloc[-1]['timestamp'] - pos['entry_date']).days
                final_profit = (final_price - pos['entry_price']) * position_size
                commission = final_profit * self.commission_rate if final_profit > 0 else 0
                interest_cost = self.borrowed_capital * self.daily_interest_rate * days_held
                net_profit = final_profit - commission - interest_cost
                
                # 更新统计数据
                total_holding_days += days_held
                total_trades += 1
                if net_profit > 0:
                    profitable_trades += 1
                else:
                    loss_trades += 1
                
                final_total_profit += net_profit
                
                print(f"\n>>> 到期强制平仓")
                print(f"持仓天数: {days_held}")
                print(f"持仓数量: {position_size:.0f}股")
                print(f"平仓价格: {final_price:.2f}")
                print(f"毛利润: {final_profit:,.2f}")
                print(f"手续费: {commission:,.2f}")
                print(f"利息成本: {interest_cost:,.2f}")
                print(f"净利润: {net_profit:,.2f}")
            
            results[-1]['action'] = 'final_sell'
            cumulative_profit += final_total_profit  # 更新最终的累计收益
            results[-1]['profit'] = cumulative_profit  # 更新最后一天的累计收益

            # 计算统计数据
            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

            print(f"\n{'='*50}")
            print(f"交易统计")
            print(f"总交易次数: {trades_count}")
            print(f"累计收益: {cumulative_profit:,.2f}")
            print(f"最大资金占用: {max_capital_usage:,.2f}")
            print(f"平均持仓天数: {avg_holding_days:.1f}")
            print(f"胜率: {win_rate:.1f}%")
            print(f"盈利交易: {profitable_trades}次")
            print(f"亏损交易: {loss_trades}次")
            print(f"{'='*50}\n")

            # 将统计数据添加到结果中
            results[-1]['stats'] = {
                'total_trades': total_trades,
                'profitable_trades': profitable_trades,
                'loss_trades': loss_trades,
                'win_rate': win_rate,
                'avg_holding_days': avg_holding_days,
                'max_capital_usage': max_capital_usage,
                'final_profit': cumulative_profit
            }

        return pd.DataFrame(results)

    def analyze_stock(self, symbol, start_date, end_date, take_profit_pct, stop_loss_pct, pullback_entry_pct):
        """分析单个股票"""
        df = self.get_stock_data(symbol, start_date, end_date, include_history=True)
        
        if len(df) < 240:  # 确保总数据量足够
            print(f"股票 {symbol} 总数据不足240天，跳过分析")
            return None
            
        return self.calculate_trade_signals(
            df,
            take_profit_pct=take_profit_pct,
            stop_loss_pct=stop_loss_pct,
            pullback_entry_pct=pullback_entry_pct,
            start_date=start_date
        )

    def analyze_all_stocks(self, start_date, end_date, take_profit_pct, stop_loss_pct, pullback_entry_pct):
        """分析所有股票"""
        stocks = self.get_stock_list()
        all_results = {}
        
        print(f"\n开始分析，共 {len(stocks)} 支股票")
        print(f"分析参数:")
        print(f"- 开始日期: {start_date}")
        print(f"- 结束日期: {end_date}")
        print(f"- 止盈: {take_profit_pct*100}%")
        print(f"- 止损: {stop_loss_pct*100}%")
        print(f"- 回调建仓: {pullback_entry_pct*100}%\n")
        
        for symbol in tqdm(stocks, desc="Analyzing stocks"):
            result = self.analyze_stock(
                symbol,
                start_date,
                end_date,
                take_profit_pct,
                stop_loss_pct,
                pullback_entry_pct
            )
            if result is not None:
                all_results[symbol] = result
                
        return all_results

    def plot_results(self, results, symbol):
        """绘制单个股票的收益走势图"""
        plt.figure(figsize=(12, 6))
        plt.plot(results['timestamp'], results['profit'].cumsum(), label='Cumulative Profit')
        plt.title(f'Stock {symbol} Trading Results')
        plt.xlabel('Date')
        plt.ylabel('Profit (CNY)')
        plt.legend()
        plt.grid(True)
        plt.xticks(rotation=45)
        plt.tight_layout()
        return plt

    def plot_all_stocks(self, all_results):
        """绘制所有股票的收益走势图"""
        plt.figure(figsize=(15, 8))
        
        # 记录所有股票的累计收益
        total_profits = {}
        max_profit = float('-inf')
        min_profit = float('inf')
        
        # 绘制每支股票的曲线
        for symbol, results in all_results.items():
            # 直接使用已经计算好的累计收益
            profits = results['profit']
            plt.plot(results['timestamp'], profits, label=symbol, alpha=0.5, linewidth=1)
            
            # 更新最大最小收益
            max_profit = max(max_profit, profits.max())
            min_profit = min(min_profit, profits.min())
            
            # 记录总收益（使用最后一个累计收益值）
            total_profits[symbol] = profits.iloc[-1]
        
        # 计算所有股票的中位数收益曲线
        # 首先找到共同的日期范围
        all_dates = set()
        for results in all_results.values():
            all_dates.update(results['timestamp'])
        all_dates = sorted(list(all_dates))
        
        # 创建一个包含所有日期的DataFrame
        profits_df = pd.DataFrame(index=all_dates)
        
        # 对每支股票，填充所有日期的收益
        for symbol, results in all_results.items():
            daily_profits = pd.Series(index=all_dates, data=np.nan)
            # 直接使用已经计算好的累计收益
            for date, profit in zip(results['timestamp'], results['profit']):
                daily_profits[date] = profit
            # 对于没有交易的日期，使用最后一个累计收益填充
            daily_profits.fillna(method='ffill', inplace=True)
            daily_profits.fillna(0, inplace=True)  # 对开始之前的日期填充0
            profits_df[symbol] = daily_profits
        
        # 计算并绘制中位数收益曲线
        median_line = profits_df.median(axis=1)
        plt.plot(all_dates, median_line, 'r-', label='Median', linewidth=2)
        
        plt.title('All Stocks Trading Results')
        plt.xlabel('Date')
        plt.ylabel('Profit (CNY)')
        plt.grid(True)
        plt.xticks(rotation=45)
        
        # 添加利润区间标注
        plt.text(0.02, 0.98, 
                f'Profit Range:\nMax: {max_profit:,.0f}\nMin: {min_profit:,.0f}\nMedian: {median_line.iloc[-1]:,.0f}', 
                transform=plt.gca().transAxes, 
                bbox=dict(facecolor='white', alpha=0.8))
        
        # 计算所有股票的统计数据
        total_stats = {
            'total_trades': 0,
            'profitable_trades': 0,
            'loss_trades': 0,
            'total_holding_days': 0,
            'max_capital_usage': 0,
            'total_profit': 0,
            'profitable_stocks': 0,
            'loss_stocks': 0,
            'total_capital_usage': 0,  # 添加总资金占用
            'total_win_rate': 0,       # 添加总胜率
        }

        # 获取前5名和后5名股票
        top_5 = sorted(total_profits.items(), key=lambda x: x[1], reverse=True)[:5]
        bottom_5 = sorted(total_profits.items(), key=lambda x: x[1])[:5]
        
        # 添加排名信息
        rank_text = "Top 5 Stocks:\n"
        for symbol, profit in top_5:
            rank_text += f"{symbol}: {profit:,.0f}\n"
        rank_text += "\nBottom 5 Stocks:\n"
        for symbol, profit in bottom_5:
            rank_text += f"{symbol}: {profit:,.0f}\n"
            
        plt.text(1.02, 0.98, rank_text, transform=plt.gca().transAxes, 
                bbox=dict(facecolor='white', alpha=0.8), verticalalignment='top')

        # 记录每支股票的统计数据用于计算平均值
        stock_stats = []

        for symbol, results in all_results.items():
            if 'stats' in results.iloc[-1]:
                stats = results.iloc[-1]['stats']
                stock_stats.append(stats)  # 记录每支股票的统计数据
                
                total_stats['total_trades'] += stats['total_trades']
                total_stats['profitable_trades'] += stats['profitable_trades']
                total_stats['loss_trades'] += stats['loss_trades']
                total_stats['total_holding_days'] += stats['avg_holding_days'] * stats['total_trades']
                total_stats['max_capital_usage'] = max(total_stats['max_capital_usage'], stats['max_capital_usage'])
                total_stats['total_capital_usage'] += stats['max_capital_usage']
                total_stats['total_profit'] += stats['final_profit']
                total_stats['total_win_rate'] += stats['win_rate']
                if stats['final_profit'] > 0:
                    total_stats['profitable_stocks'] += 1
                else:
                    total_stats['loss_stocks'] += 1

        # 计算总体统计
        total_stocks = total_stats['profitable_stocks'] + total_stats['loss_stocks']
        avg_holding_days = total_stats['total_holding_days'] / total_stats['total_trades'] if total_stats['total_trades'] > 0 else 0
        win_rate = total_stats['profitable_trades'] / total_stats['total_trades'] * 100 if total_stats['total_trades'] > 0 else 0
        stock_win_rate = total_stats['profitable_stocks'] / total_stocks * 100 if total_stocks > 0 else 0
        
        # 计算平均值统计
        avg_trades_per_stock = total_stats['total_trades'] / total_stocks if total_stocks > 0 else 0
        avg_profit_per_stock = total_stats['total_profit'] / total_stocks if total_stocks > 0 else 0
        avg_capital_usage = total_stats['total_capital_usage'] / total_stocks if total_stocks > 0 else 0
        avg_win_rate = total_stats['total_win_rate'] / total_stocks if total_stocks > 0 else 0

        # 添加总体统计信息
        stats_text = (
            f"总体统计:\n"
            f"总交易次数: {total_stats['total_trades']}\n"
            f"总盈利: {total_stats['total_profit']:,.0f}\n"
            # f"最大资金占用: {total_stats['max_capital_usage']:,.0f}\n"
            f"平均持仓天数: {avg_holding_days:.1f}\n"
            f"交易胜率: {win_rate:.1f}%\n"
            f"股票胜率: {stock_win_rate:.1f}%\n"
            f"盈利股票数: {total_stats['profitable_stocks']}\n"
            f"亏损股票数: {total_stats['loss_stocks']}\n"
            f"\n平均值统计:\n"
            f"平均交易次数: {avg_trades_per_stock:.1f}\n"
            f"平均收益: {avg_profit_per_stock:,.0f}\n"
            f"平均资金占用: {avg_capital_usage:,.0f}\n"
            f"平均持仓天数: {avg_holding_days:.1f}\n"
            f"平均胜率: {avg_win_rate:.1f}%"
        )
        
        plt.text(0.02, 0.6, stats_text,
                transform=plt.gca().transAxes,
                bbox=dict(facecolor='white', alpha=0.8, edgecolor='gray'),
                verticalalignment='top',
                fontsize=10)
        
        # 获取前5名和后5名股票
        top_5 = sorted(total_profits.items(), key=lambda x: x[1], reverse=True)[:5]
        bottom_5 = sorted(total_profits.items(), key=lambda x: x[1])[:5]
        
        # 添加排名信息
        rank_text = "Top 5 Stocks:\n"
        for symbol, profit in top_5:
            rank_text += f"{symbol}: {profit:,.0f}\n"
        rank_text += "\nBottom 5 Stocks:\n"
        for symbol, profit in bottom_5:
            rank_text += f"{symbol}: {profit:,.0f}\n"
            
        plt.text(1.02, 0.98, rank_text, transform=plt.gca().transAxes, 
                bbox=dict(facecolor='white', alpha=0.8), verticalalignment='top')
        
        plt.tight_layout()
        return plt, total_profits

    def plot_profit_distribution(self, total_profits):
        """绘制所有股票的收益分布小提琴图"""
        plt.figure(figsize=(12, 8))
        
        # 将收益数据转换为数组，确保使用最终累计收益
        profits = np.array([profit for profit in total_profits.values()])
        
        # 计算统计数据
        mean_profit = np.mean(profits)
        median_profit = np.median(profits)
        max_profit = np.max(profits)
        min_profit = np.min(profits)
        std_profit = np.std(profits)
        
        # 绘制小提琴图
        violin_parts = plt.violinplot(profits, positions=[0], showmeans=True, showmedians=True)
        
        # 设置小提琴图的颜色
        violin_parts['bodies'][0].set_facecolor('lightblue')
        violin_parts['bodies'][0].set_alpha(0.7)
        violin_parts['cmeans'].set_color('red')
        violin_parts['cmedians'].set_color('blue')
        
        # 添加统计信息
        stats_text = (
            f"统计信息:\n"
            f"平均收益: {mean_profit:,.0f}\n"
            f"中位数收益: {median_profit:,.0f}\n"
            f"最大收益: {max_profit:,.0f}\n"
            f"最小收益: {min_profit:,.0f}\n"
            f"标准差: {std_profit:,.0f}"
        )
        
        plt.text(0.65, 0.95, stats_text,
                transform=plt.gca().transAxes,
                bbox=dict(facecolor='white', alpha=0.8, edgecolor='gray'),
                verticalalignment='top',
                fontsize=10)
        
        # 设置图表样式
        plt.title('股票收益分布图', fontsize=14, pad=20)
        plt.ylabel('收益 (元)', fontsize=12)
        plt.xticks([0], ['所有股票'], fontsize=10)
        plt.grid(True, axis='y', alpha=0.3)
        
        # 添加零线
        plt.axhline(y=0, color='r', linestyle='--', alpha=0.3)
        
        # 计算盈利和亏损的股票数量
        profit_count = np.sum(profits > 0)
        loss_count = np.sum(profits < 0)
        total_count = len(profits)
        
        # 添加盈亏比例信息
        ratio_text = (
            f"盈亏比例:\n"
            f"盈利: {profit_count}支 ({profit_count/total_count*100:.1f}%)\n"
            f"亏损: {loss_count}支 ({loss_count/total_count*100:.1f}%)\n"
            f"总计: {total_count}支"
        )
        
        plt.text(0.65, 0.6, ratio_text,
                transform=plt.gca().transAxes,
                bbox=dict(facecolor='white', alpha=0.8, edgecolor='gray'),
                verticalalignment='top',
                fontsize=10)
        
        plt.tight_layout()
        return plt

    def plot_profit_matrix(self, df):
        """绘制止盈比例分析矩阵的结果图"""
        plt.figure(figsize=(15, 10))
        
        # 创建子图
        fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(15, 10))
        
        # 1. 平均收益曲线
        ax1.plot(df.index, df['avg_profit'], 'b-', linewidth=2)
        ax1.set_title('平均收益 vs 止盈比例')
        ax1.set_xlabel('止盈比例 (%)')
        ax1.set_ylabel('平均收益 (元)')
        ax1.grid(True)
        
        # 标记最优止盈比例点
        best_profit_pct = df['avg_profit'].idxmax()
        best_profit = df['avg_profit'].max()
        ax1.plot(best_profit_pct, best_profit, 'ro')
        ax1.annotate(f'最优: {best_profit_pct:.1f}%\n{best_profit:,.0f}元',
                    (best_profit_pct, best_profit),
                    xytext=(10, 10), textcoords='offset points')
        
        # 2. 平均胜率曲线
        ax2.plot(df.index, df['avg_win_rate'], 'g-', linewidth=2)
        ax2.set_title('平均胜率 vs 止盈比例')
        ax2.set_xlabel('止盈比例 (%)')
        ax2.set_ylabel('胜率 (%)')
        ax2.grid(True)
        
        # 3. 平均持仓天数曲线
        ax3.plot(df.index, df['avg_holding_days'], 'r-', linewidth=2)
        ax3.set_title('平均持仓天数 vs 止盈比例')
        ax3.set_xlabel('止盈比例 (%)')
        ax3.set_ylabel('持仓天数')
        ax3.grid(True)
        
        # 4. 平均交易次数和资金占用
        ax4.plot(df.index, df['avg_trades'], 'c-', linewidth=2, label='平均交易次数')
        ax4_twin = ax4.twinx()
        ax4_twin.plot(df.index, df['avg_capital_usage'], 'm-', linewidth=2, label='平均资金占用')
        ax4.set_title('平均交易次数和资金占用 vs 止盈比例')
        ax4.set_xlabel('止盈比例 (%)')
        ax4.set_ylabel('交易次数')
        ax4_twin.set_ylabel('资金占用 (元)')
        
        # 合并两个y轴的图例
        lines1, labels1 = ax4.get_legend_handles_labels()
        lines2, labels2 = ax4_twin.get_legend_handles_labels()
        ax4_twin.legend(lines1 + lines2, labels1 + labels2, loc='upper right')
        ax4.grid(True)
        
        # 调整布局
        plt.tight_layout()
        return plt

    def analyze_profit_matrix(self, start_date, end_date, stop_loss_pct, pullback_entry_pct):
        """分析不同止盈比例的表现矩阵"""
        # 创建结果列表
        matrix_results = []
        
        # 从30%到5%，每次减少1%
        for take_profit_pct in np.arange(0.30, 0.04, -0.01):
            print(f"\n分析止盈比例: {take_profit_pct*100:.1f}%")
            
            # 运行分析
            results = self.analyze_all_stocks(
                start_date,
                end_date,
                take_profit_pct,
                stop_loss_pct,
                pullback_entry_pct
            )
            
            # 过滤掉空结果
            valid_results = {symbol: result for symbol, result in results.items() if result is not None and not result.empty}
            
            if valid_results:
                # 计算统计数据
                total_stats = {
                    'total_trades': 0,
                    'profitable_trades': 0,
                    'loss_trades': 0,
                    'total_holding_days': 0,
                    'total_profit': 0,
                    'total_capital_usage': 0,
                    'total_win_rate': 0,
                    'stock_count': len(valid_results)
                }
                
                # 收集每支股票的统计数据
                for symbol, result in valid_results.items():
                    if 'stats' in result.iloc[-1]:
                        stats = result.iloc[-1]['stats']
                        total_stats['total_trades'] += stats['total_trades']
                        total_stats['profitable_trades'] += stats['profitable_trades']
                        total_stats['loss_trades'] += stats['loss_trades']
                        total_stats['total_holding_days'] += stats['avg_holding_days'] * stats['total_trades']
                        total_stats['total_profit'] += stats['final_profit']
                        total_stats['total_capital_usage'] += stats['max_capital_usage']
                        total_stats['total_win_rate'] += stats['win_rate']
                
                # 计算平均值
                avg_trades = total_stats['total_trades'] / total_stats['stock_count']
                avg_holding_days = total_stats['total_holding_days'] / total_stats['total_trades'] if total_stats['total_trades'] > 0 else 0
                avg_profit = total_stats['total_profit'] / total_stats['stock_count']
                avg_capital_usage = total_stats['total_capital_usage'] / total_stats['stock_count']
                avg_win_rate = total_stats['total_win_rate'] / total_stats['stock_count']
                
                # 添加到结果矩阵
                matrix_results.append({
                    'take_profit_pct': take_profit_pct * 100,
                    'avg_trades': avg_trades,
                    'avg_holding_days': avg_holding_days,
                    'avg_profit': avg_profit,
                    'avg_capital_usage': avg_capital_usage,
                    'avg_win_rate': avg_win_rate,
                    'total_trades': total_stats['total_trades'],
                    'stock_count': total_stats['stock_count']
                })
        
        # 创建DataFrame并格式化
        df = pd.DataFrame(matrix_results)
        df = df.round({
            'take_profit_pct': 1,
            'avg_trades': 1,
            'avg_holding_days': 1,
            'avg_profit': 0,
            'avg_capital_usage': 0,
            'avg_win_rate': 1,
            'total_trades': 0,
            'stock_count': 0
        })
        
        # 设置止盈比例为索引
        df.set_index('take_profit_pct', inplace=True)
        
        # 保存到CSV文件
        df.to_csv('results/profit_matrix_analysis.csv')
        
        # 打印结果
        print("\n止盈比例分析矩阵:")
        print("=" * 120)
        print(df.to_string())
        print("=" * 120)
        
        # 绘制并保存矩阵分析图表
        plt = self.plot_profit_matrix(df)
        plt.savefig('results/profit_matrix_analysis.png', bbox_inches='tight', dpi=300)
        plt.close()
        
        return df

def main():
    # 数据库连接配置
    db_config = {
        'host': '192.168.18.199',
        'port': 3306,
        'user': 'root',
        'password': 'Chlry#$.8',
        'database': 'db_gp_cj'
    }
    
    try:
        connection_string = f"mysql+pymysql://{db_config['user']}:{db_config['password']}@{db_config['host']}:{db_config['port']}/{db_config['database']}"
        
        # 创建results目录
        if not os.path.exists('results'):
            os.makedirs('results')
        
        # 创建分析器实例
        analyzer = StockAnalyzer(connection_string)
        
        # 设置分析参数
        start_date = '2024-01-01'
        end_date = '2024-12-31'
        take_profit_pct = 0.25  # 30%止盈
        stop_loss_pct = -0.99   # -99%止损
        pullback_entry_pct = -0.05  # -5%回调建仓
        
        # 控制是否运行止盈比例分析的标志
        run_profit_matrix = False  # 设置为 True 时运行止盈比例分析
        
        if run_profit_matrix:
            # 运行止盈比例分析
            profit_matrix = analyzer.analyze_profit_matrix(
                start_date,
                end_date,
                stop_loss_pct,
                pullback_entry_pct
            )
            # 使用最优止盈比例
            take_profit_pct = profit_matrix['avg_profit'].idxmax() / 100
            print(f"\n使用最优止盈比例 {take_profit_pct*100:.1f}% 运行详细分析")
        
        # 运行分析
        results = analyzer.analyze_all_stocks(
            start_date,
            end_date,
            take_profit_pct,
            stop_loss_pct,
            pullback_entry_pct
        )
        
        # 保存结果
        for symbol, result in results.items():
            if result is not None and not result.empty:
                result.to_csv(f'results/{symbol}_analysis.csv', index=False)
                plt = analyzer.plot_results(result, symbol)
                plt.savefig(f'results/{symbol}_chart.png')
                plt.close()
            else:
                print(f"跳过空结果: {symbol}")
        
        # 过滤掉空结果
        valid_results = {symbol: result for symbol, result in results.items() if result is not None and not result.empty}
        
        if valid_results:
            # 绘制所有股票的收益曲线
            plt, total_profits = analyzer.plot_all_stocks(valid_results)
            plt.savefig('results/all_stocks_analysis.png', bbox_inches='tight', dpi=300)
            plt.close()
            
            # 绘制收益分布小提琴图
            plt = analyzer.plot_profit_distribution(total_profits)
            plt.savefig('results/profit_distribution.png', bbox_inches='tight', dpi=300)
            plt.close()
        else:
            print("没有有效的分析结果可供绘图")
            
    except Exception as e:
        print(f"程序运行出错: {str(e)}")
        raise

if __name__ == "__main__":
    main() 