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
# 止盈10%，不止损，较上一次止盈或者止损价回调5%之后建仓
# 添加中文字体支持
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"
        return pd.read_sql(query, self.engine)['symbol'].tolist()

    def get_stock_data(self, symbol, start_date, end_date):
        """获取指定股票在时间范围内的数据"""
        query = f"""
        SELECT * FROM gp_day_data 
        WHERE symbol = '{symbol}' 
        AND timestamp BETWEEN '{start_date}' AND '{end_date}'
        ORDER BY timestamp
        """
        return pd.read_sql(query, self.engine)

    def calculate_trade_signals(self, df, take_profit_pct, stop_loss_pct, pullback_entry_pct):
        """计算交易信号和收益"""
        results = []
        position = False
        entry_price = 0
        entry_date = None
        last_exit_price = None
        cumulative_profit = 0
        trades_count = 0
        max_holding_period = timedelta(days=180)  # 最长持仓6个月
        max_loss_amount = -300000  # 最大亏损金额为30万

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

        for index, row in df.iterrows():
            current_price = (float(row['high']) + float(row['low'])) / 2
            day_result = {
                'date': row['timestamp'],
                'profit': 0,
                'position': False,
                'action': 'hold'
            }

            if not position:
                # 检查是否可以建仓
                if last_exit_price is None or (
                    (current_price - last_exit_price) / last_exit_price <= pullback_entry_pct
                ):
                    position = True
                    entry_price = current_price
                    entry_date = row['timestamp']
                    day_result['action'] = 'buy'
                    day_result['position'] = True
                    trades_count += 1
                    print(f"\n>>> 建仓信号 #{trades_count}")
                    print(f"日期: {entry_date.strftime('%Y-%m-%d')}")
                    print(f"建仓价格: {entry_price:.2f}")
                    print(f"建仓数量: {self.total_position/entry_price:.0f}股")
                    if last_exit_price is not None:
                        print(f"距离上次平仓价格变化: {((current_price - last_exit_price) / last_exit_price * 100):.2f}%")
            else:
                # 计算当前持仓收益
                current_profit = (current_price - entry_price) / entry_price
                days_held = (row['timestamp'] - entry_date).days
                holding_period = row['timestamp'] - entry_date

                # 计算当天的未实现收益
                position_size = self.total_position / entry_price  # 计算持仓数量
                day_profit = (current_price - entry_price) * position_size  # 用持仓数量计算收益
                day_result['profit'] = day_profit
                day_result['position'] = True

                # 计算当前亏损（包括利息成本）
                interest_cost = self.borrowed_capital * self.daily_interest_rate * days_held
                current_net_profit = day_profit - interest_cost

                # 打印每日收益情况
                print(f"\n当日交易 - {row['timestamp'].strftime('%Y-%m-%d')}")
                print(f"当前价格: {current_price:.2f}")
                print(f"持仓天数: {days_held}")
                print(f"当日收益: {day_profit:,.2f} ({current_profit*100:.2f}%)")

                # 检查是否满足平仓条件
                if days_held >= self.min_holding_days:
                    should_close = False
                    close_type = None
                    close_price = current_price

                    # 检查止盈
                    if float(row['high']) >= entry_price * (1 + take_profit_pct):
                        should_close = True
                        close_type = 'sell_profit'
                        close_price = entry_price * (1 + take_profit_pct)
                    # 检查最大亏损
                    elif current_net_profit <= max_loss_amount:
                        should_close = True
                        close_type = 'sell_loss'
                        close_price = current_price
                    # 检查持仓时间
                    elif holding_period >= max_holding_period:
                        should_close = True
                        close_type = 'time_limit'
                        close_price = current_price

                    if should_close:
                        position_size = self.total_position / entry_price
                        profit = (close_price - entry_price) * position_size
                        commission = profit * self.commission_rate if profit > 0 else 0
                        interest_cost = self.borrowed_capital * self.daily_interest_rate * days_held
                        net_profit = profit - commission - interest_cost
                        
                        cumulative_profit += net_profit
                        position = False
                        last_exit_price = close_price
                        day_result['action'] = close_type
                        day_result['profit'] = net_profit
                        
                        print(f"\n>>> {'止盈平仓' if close_type == 'sell_profit' else '止损平仓' if close_type == 'sell_loss' else '到期平仓'}")
                        print(f"持仓天数: {days_held}")
                        print(f"持仓数量: {position_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}")

            results.append(day_result)

        # 在最后一天强制平仓
        if position:
            final_price = (float(df.iloc[-1]['high']) + float(df.iloc[-1]['low'])) / 2
            position_size = self.total_position / entry_price
            final_profit = (final_price - entry_price) * position_size
            days_held = (df.iloc[-1]['timestamp'] - entry_date).days
            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
            
            results[-1]['action'] = 'final_sell'
            results[-1]['profit'] = net_profit
            cumulative_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}")

        print(f"\n{'='*50}")
        print(f"交易统计")
        print(f"总交易次数: {trades_count}")
        print(f"累计收益: {cumulative_profit:,.2f}")
        print(f"{'='*50}\n")

        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)
        
        if len(df) < 10:  # 跳过数据少于10条的股票
            print(f"股票 {symbol} 数据不足10条，跳过分析")
            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
        )

    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['date'], 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():
            cumulative_profits = results['profit'].cumsum()
            plt.plot(results['date'], cumulative_profits, label=symbol, alpha=0.5, linewidth=1)
            
            # 更新最大最小收益
            max_profit = max(max_profit, cumulative_profits.max())
            min_profit = min(min_profit, cumulative_profits.min())
            
            # 记录最终收益
            total_profits[symbol] = cumulative_profits.iloc[-1]
        
        # 计算所有股票的平均收益曲线
        # 首先找到共同的日期范围
        all_dates = set()
        for results in all_results.values():
            all_dates.update(results['date'])
        all_dates = sorted(list(all_dates))
        
        # 创建一个包含所有日期的DataFrame
        avg_profits = pd.DataFrame(index=all_dates)
        
        # 对每支股票，填充所有日期的收益
        for symbol, results in all_results.items():
            daily_profits = pd.Series(index=all_dates, data=0.0)
            cumulative = 0
            last_cumulative = 0
            for date, profit in zip(results['date'], results['profit']):
                cumulative += profit
                daily_profits[date] = cumulative
                last_cumulative = cumulative
            # 对于没有交易的日期，使用最后一个累计收益填充
            daily_profits.fillna(method='ffill', inplace=True)
            daily_profits.fillna(0, inplace=True)  # 对开始之前的日期填充0
            avg_profits[symbol] = daily_profits
        
        # 计算并绘制平均收益曲线
        avg_line = avg_profits.mean(axis=1)
        plt.plot(all_dates, avg_line, 'r-', label='Average', 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}\nAvg: {avg_line.iloc[-1]:,.0f}', 
                transform=plt.gca().transAxes, bbox=dict(facecolor='white', alpha=0.8))
        
        # 获取前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(list(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 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.10  # 10%止盈
        stop_loss_pct = -0.99   # -10%止损
        pullback_entry_pct = -0.05  # -5%回调建仓
        
        # 运行分析
        results = analyzer.analyze_all_stocks(
            start_date,
            end_date,
            take_profit_pct,
            stop_loss_pct,
            pullback_entry_pct
        )
        
        # 保存结果
        for symbol, result in results.items():
            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()
        
        # 绘制所有股票的收益曲线
        plt, total_profits = analyzer.plot_all_stocks(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()
            
    except Exception as e:
        print(f"程序运行出错: {str(e)}")
        raise

if __name__ == "__main__":
    main() 