import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from datetime import datetime

plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False


class MeanReversionStrategy:
    def __init__(self,
                 initial_cash=100000,
                 initial_shares=10000,
                 lookback_days=20,
                 std_multiplier=1.0,
                 trade_volume=400):
        """
        均值回归策略

        Args:
            initial_cash: 初始资金
            initial_shares: 初始持股
            lookback_days: 回看天数（计算均值和标准差）
            std_multiplier: 标准差倍数（带宽）
            trade_volume: 每次交易数量
        """
        self.cash = initial_cash
        self.shares = initial_shares
        self.lookback_days = lookback_days
        self.std_multiplier = std_multiplier
        self.trade_volume = trade_volume

        # 交易费用
        self.stamp_tax_rate = 0.001
        self.commission = 5
        self.transfer_fee = 1

        self.trade_records = []
        self.daily_positions = []
        self.signals = []

    def calculate_fees(self, price, volume, is_buy):
        """计算交易费用"""
        total_amount = price * volume
        total_fee = 0

        if is_buy:
            total_fee = self.commission + self.transfer_fee
        else:
            stamp_tax = total_amount * self.stamp_tax_rate
            total_fee = self.commission + self.transfer_fee + stamp_tax

        return total_fee

    def execute_trade(self, trade_date, price, is_buy):
        """执行交易"""
        volume = self.trade_volume
        total_amount = price * volume

        fee = self.calculate_fees(price, volume, is_buy)
        total_cost = total_amount + fee if is_buy else fee

        if is_buy:
            if self.cash < total_cost:
                print(f"警告：{trade_date} 资金不足，无法买入 {volume} 股")
                return False

            self.cash -= total_cost
            self.shares += volume
            trade_type = "买入"

        else:
            if self.shares < volume:
                print(f"警告：{trade_date} 持仓不足，无法卖出 {volume} 股")
                return False

            self.cash += total_amount - fee
            self.shares -= volume
            trade_type = "卖出"

        record = {
            '日期': trade_date,
            '交易类型': trade_type,
            '价格': round(price, 4),
            '数量': volume,
            '成交金额': round(total_amount, 2),
            '费用': round(fee, 2),
            '现金余额': round(self.cash, 2),
            '持仓数量': self.shares
        }
        self.trade_records.append(record)

        return True

    def calculate_signal(self, prices):
        """
        计算交易信号
        当价格低于均值-标准差*倍数时买入
        当价格高于均值+标准差*倍数时卖出
        """
        if len(prices) < self.lookback_days:
            return 0, None, None

        recent_prices = prices[-self.lookback_days:]
        mean_price = np.mean(recent_prices)
        std_price = np.std(recent_prices)
        current_price = prices[-1]

        upper_bound = mean_price + self.std_multiplier * std_price
        lower_bound = mean_price - self.std_multiplier * std_price

        # 信号: 1买入, -1卖出, 0持有
        signal = 0
        if current_price <= lower_bound:
            signal = 1  # 买入信号
        elif current_price >= upper_bound:
            signal = -1  # 卖出信号

        return signal, upper_bound, lower_bound

    def run_strategy(self, df, train_days=3000):
        """
        运行均值回归策略
        """
        print("开始运行均值回归策略...")
        print("=" * 80)

        # 数据清洗
        df = df.copy()
        df = df[df['开盘价(元)'].notna()]
        df = df[df['收盘价(元)'].notna()]
        df = df.sort_values('日期').reset_index(drop=True)

        total_days = len(df)
        print(f"总交易日: {total_days} 天")
        print(f"训练集: 前 {train_days} 天")
        print(f"测试集: 后 {total_days - train_days} 天")
        print("=" * 80)

        # 划分数据集
        train_df = df.iloc[:train_days].copy()
        test_df = df.iloc[train_days:].copy()

        print(f"\n训练集参数:")
        print(f"  回看天数: {self.lookback_days}")
        print(f"  标准差倍数: {self.std_multiplier}")
        print(f"  每次交易量: {self.trade_volume} 股")
        print("=" * 80)

        # 在训练集上计算统计参数
        train_prices = train_df['收盘价(元)'].values
        mean_price_train = np.mean(train_prices)
        std_price_train = np.std(train_prices)

        print(f"\n训练集统计:")
        print(f"  均值: {mean_price_train:.4f}")
        print(f"  标准差: {std_price_train:.4f}")
        print(f"  上限: {mean_price_train + self.std_multiplier * std_price_train:.4f}")
        print(f"  下限: {mean_price_train - self.std_multiplier * std_price_train:.4f}")
        print("=" * 80)

        # 在测试集上运行策略
        print("\n开始执行测试集交易...")
        test_prices = test_df['收盘价(元)'].values

        # 初始化历史价格列表
        price_history = list(train_prices[-self.lookback_days:])

        for idx, row in test_df.iterrows():
            current_date = row['日期']
            current_price = row['收盘价(元)']
            open_price = row['开盘价(元)']

            if pd.isna(current_price):
                continue

            price_history.append(current_price)

            # 记录每日持仓
            self.daily_positions.append({
                '日期': current_date,
                '持仓数量': self.shares,
                '现金余额': round(self.cash, 2),
                '持仓市值': round(self.shares * current_price, 2),
                '总资产': round(self.cash + self.shares * current_price, 2),
                '收盘价': round(current_price, 4)
            })

            # 计算信号
            signal, upper_bound, lower_bound = self.calculate_signal(price_history)

            # 记录信号
            if signal != 0:
                self.signals.append({
                    '日期': current_date,
                    '价格': current_price,
                    '信号': '买入' if signal == 1 else '卖出',
                    '上限': upper_bound,
                    '下限': lower_bound,
                    '持仓数量': self.shares
                })

            # 执行交易
            if signal == 1:  # 买入信号
                self.execute_trade(current_date, current_price, is_buy=True)
            elif signal == -1:  # 卖出信号
                self.execute_trade(current_date, current_price, is_buy=False)

            # 保持价格历史长度
            if len(price_history) > self.lookback_days:
                price_history = price_history[-self.lookback_days:]

        # 计算最终结果
        final_cash = self.cash
        final_shares = self.shares
        final_price = test_prices[-1] if len(test_prices) > 0 else 0
        final_value = final_cash + final_shares * final_price

        # 计算基准收益
        initial_price = train_prices[-1] if len(train_prices) > 0 else test_prices[0]
        benchmark_return = (final_price - initial_price) / initial_price * 100
        strategy_return = (final_value - 100000) / 100000 * 100

        results = {
            '交易次数': len(self.trade_records),
            '买入次数': sum(1 for r in self.trade_records if r['交易类型'] == '买入'),
            '卖出次数': sum(1 for r in self.trade_records if r['交易类型'] == '卖出'),
            '初始资金': 100000,
            '最终现金': round(final_cash, 2),
            '最终持股': final_shares,
            '最终价格': round(final_price, 4),
            '最终总资产': round(final_value, 2),
            '策略收益': round(strategy_return, 2),
            '基准收益': round(benchmark_return, 2),
            '超额收益': round(strategy_return - benchmark_return, 2)
        }

        self.print_results(results)
        self.save_records()
        self.plot_results(test_df, train_prices)

        return results

    def print_results(self, results):
        """打印结果"""
        print("\n" + "=" * 80)
        print("策略回测结果")
        print("=" * 80)
        print(f"交易统计:")
        print(f"  总交易次数: {results['交易次数']}")
        print(f"  买入次数: {results['买入次数']}")
        print(f"  卖出次数: {results['卖出次数']}")
        print(f"\n资金变化:")
        print(f"  初始资金: {results['初始资金']:,} 元")
        print(f"  最终现金: {results['最终现金']:,.2f} 元")
        print(f"  最终持股: {results['最终持股']} 股")
        print(f"  最终价格: {results['最终价格']:.4f} 元")
        print(f"  最终总资产: {results['最终总资产']:,.2f} 元")
        print(f"\n收益对比:")
        print(f"  策略收益: {results['策略收益']:.2f}%")
        print(f"  基准收益: {results['基准收益']:.2f}%")
        print(f"  超额收益: {results['超额收益']:.2f}%")
        print("=" * 80)

        # 显示最近10笔交易
        if len(self.trade_records) > 0:
            print(f"\n最近10笔交易记录:")
            recent_trades = self.trade_records[-10:]
            df_trades = pd.DataFrame(recent_trades)
            print(df_trades.to_string(index=False))

    def save_records(self):
        """保存交易记录"""
        if self.trade_records:
            df_trades = pd.DataFrame(self.trade_records)
            df_trades.to_csv('均值回归交易记录.csv', index=False, encoding='utf-8-sig')
            print(f"\n交易记录已保存到 '均值回归交易记录.csv'")

        if self.daily_positions:
            df_pos = pd.DataFrame(self.daily_positions)
            df_pos.to_csv('均值回归每日持仓.csv', index=False, encoding='utf-8-sig')
            print(f"每日持仓已保存到 '均值回归每日持仓.csv'")

    def plot_results(self, test_df, train_prices):
        """绘制结果图表"""
        try:
            fig, axes = plt.subplots(3, 1, figsize=(14, 12))

            # 图1: 价格与买卖信号
            ax1 = axes[0]
            test_prices = test_df['收盘价(元)'].values
            test_dates = test_df['日期'].values

            # 转换日期格式用于绘图
            dates = pd.to_datetime(test_dates)

            ax1.plot(dates, test_prices, 'b-', label='收盘价', linewidth=1.5)

            # 标记买卖点
            buy_points = [(pd.to_datetime(s['日期']), s['价格']) for s in self.signals if s['信号'] == '买入']
            sell_points = [(pd.to_datetime(s['日期']), s['价格']) for s in self.signals if s['信号'] == '卖出']

            if buy_points:
                buy_dates, buy_prices = zip(*buy_points)
                ax1.scatter(buy_dates, buy_prices, c='g', marker='^', s=100, label='买入信号', zorder=5)
            if sell_points:
                sell_dates, sell_prices = zip(*sell_points)
                ax1.scatter(sell_dates, sell_prices, c='r', marker='v', s=100, label='卖出信号', zorder=5)

            ax1.set_title('均值回归策略 - 价格与交易信号', fontsize=14, fontweight='bold')
            ax1.set_xlabel('日期')
            ax1.set_ylabel('价格 (元)')
            ax1.legend()
            ax1.grid(True, alpha=0.3)

            # 图2: 累计收益对比
            ax2 = axes[1]
            if self.daily_positions:
                pos_df = pd.DataFrame(self.daily_positions)
                pos_dates = pd.to_datetime(pos_df['日期'])
                total_assets = pos_df['总资产'].values

                # 计算策略累计收益
                strategy_returns = (total_assets / 100000 - 1) * 100
                ax2.plot(pos_dates, strategy_returns, 'g-', label='策略收益', linewidth=2)

                # 计算基准累计收益
                initial_price = train_prices[-1] if len(train_prices) > 0 else test_prices[0]
                benchmark_returns = (pos_df['收盘价'].values / initial_price - 1) * 100
                ax2.plot(pos_dates, benchmark_returns, 'b--', label='基准收益 (买入持有)', linewidth=2)

                ax2.set_title('累计收益对比', fontsize=14, fontweight='bold')
                ax2.set_xlabel('日期')
                ax2.set_ylabel('累计收益 (%)')
                ax2.legend()
                ax2.grid(True, alpha=0.3)

                # 添加零线
                ax2.axhline(y=0, color='k', linestyle='-', alpha=0.2)

            # 图3: 资产曲线与资金利用率
            ax3 = axes[2]
            if self.daily_positions:
                pos_df = pd.DataFrame(self.daily_positions)
                pos_dates = pd.to_datetime(pos_df['日期'])

                # 总资产曲线
                ax3.plot(pos_dates, pos_df['总资产'], 'purple', label='总资产', linewidth=2)
                ax3.fill_between(pos_dates, 100000, pos_df['总资产'], alpha=0.2)

                # 现金余额
                ax3.plot(pos_dates, pos_df['现金余额'], 'orange', label='现金余额', linewidth=1.5)

                ax3.set_title('资产变化曲线', fontsize=14, fontweight='bold')
                ax3.set_xlabel('日期')
                ax3.set_ylabel('资产 (元)')
                ax3.legend()
                ax3.grid(True, alpha=0.3)

                # 添加初始资金线
                ax3.axhline(y=100000, color='k', linestyle='--', alpha=0.3, label='初始资金')

            plt.tight_layout()
            plt.savefig('均值回归策略结果.png', dpi=300, bbox_inches='tight')
            print("结果图表已保存到 '均值回归策略结果.png'")
            plt.show()

        except Exception as e:
            print(f"绘图失败: {e}")


def optimize_parameters(df):
    """
    参数优化：寻找最优的回看天数和标准差倍数
    """
    print("\n" + "=" * 80)
    print("开始参数优化...")
    print("=" * 80)

    best_return = -float('inf')
    best_params = {}
    results_list = []

    # 参数网格
    lookback_options = [10, 15, 20, 25, 30, 40, 50]
    std_options = [0.5, 0.8, 1.0, 1.2, 1.5, 2.0]

    # 训练集天数
    train_days = 3000

    df_clean = df.copy()
    df_clean = df_clean[df_clean['收盘价(元)'].notna()]
    df_clean = df_clean.sort_values('日期').reset_index(drop=True)

    train_df = df_clean.iloc[:train_days]
    test_df = df_clean.iloc[train_days:]

    if len(test_df) < 100:
        print("测试数据不足，无法进行优化")
        return None

    print(f"优化参数范围:")
    print(f"  回看天数: {lookback_options}")
    print(f"  标准差倍数: {std_options}")
    print(f"  训练集大小: {train_days} 天")
    print(f"  测试集大小: {len(test_df)} 天")
    print("-" * 80)

    total_combinations = len(lookback_options) * len(std_options)
    count = 0

    for lookback in lookback_options:
        for std_mult in std_options:
            count += 1
            print(f"  测试组合 {count}/{total_combinations}: 回看={lookback}, 标准差={std_mult:.1f}")

            # 运行策略
            strategy = MeanReversionStrategy(
                initial_cash=100000,
                initial_shares=10000,
                lookback_days=lookback,
                std_multiplier=std_mult,
                trade_volume=400
            )

            # 重置交易状态
            strategy.cash = 100000
            strategy.shares = 10000
            strategy.trade_records = []
            strategy.daily_positions = []
            strategy.signals = []

            # 运行策略
            try:
                # 直接运行测试集
                test_prices = test_df['收盘价(元)'].values
                train_prices = train_df['收盘价(元)'].values

                price_history = list(train_prices[-lookback:])

                for idx, row in test_df.iterrows():
                    current_price = row['收盘价(元)']
                    if pd.isna(current_price):
                        continue

                    price_history.append(current_price)
                    signal, _, _ = strategy.calculate_signal(price_history)

                    if signal == 1:
                        strategy.execute_trade(row['日期'], current_price, is_buy=True)
                    elif signal == -1:
                        strategy.execute_trade(row['日期'], current_price, is_buy=False)

                    if len(price_history) > lookback:
                        price_history = price_history[-lookback:]

                final_price = test_prices[-1]
                final_value = strategy.cash + strategy.shares * final_price
                strategy_return = (final_value - 100000) / 100000 * 100

                results_list.append({
                    '回看天数': lookback,
                    '标准差倍数': std_mult,
                    '交易次数': len(strategy.trade_records),
                    '策略收益': round(strategy_return, 2),
                    '最终资产': round(final_value, 2)
                })

                if strategy_return > best_return:
                    best_return = strategy_return
                    best_params = {'lookback': lookback, 'std_mult': std_mult}

            except Exception as e:
                print(f"    组合出错: {e}")
                continue

    # 显示优化结果
    print("\n" + "=" * 80)
    print("参数优化结果")
    print("=" * 80)

    # 排序显示
    results_df = pd.DataFrame(results_list)
    results_df = results_df.sort_values('策略收益', ascending=False)

    print("\n最优参数:")
    print(f"  回看天数: {best_params['lookback']}")
    print(f"  标准差倍数: {best_params['std_mult']:.1f}")
    print(f"  最优收益: {best_return:.2f}%")

    print("\n\n前10名参数组合:")
    print(results_df.head(10).to_string(index=False))

    # 保存优化结果
    results_df.to_csv('参数优化结果.csv', index=False, encoding='utf-8-sig')
    print("\n优化结果已保存到 '参数优化结果.csv'")

    return best_params


def main():
    """主函数"""
    try:
        # 读取数据
        df = pd.read_csv('个股数据.csv', encoding='utf-8')
        print(f"数据加载成功！共 {len(df)} 条记录")
        print(f"日期范围: {df['日期'].iloc[0]} 到 {df['日期'].iloc[-1]}")
        print("-" * 80)

        # 数据预处理
        df['日期'] = pd.to_datetime(df['日期'])

        # 可选：参数优化
        print("\n是否进行参数优化？(y/n)")
        choice = input().strip().lower()

        if choice == 'y' or choice == 'yes':
            best_params = optimize_parameters(df)
            if best_params:
                print(f"\n使用优化后的参数: 回看={best_params['lookback']}, 标准差={best_params['std_mult']}")
                lookback = best_params['lookback']
                std_mult = best_params['std_mult']
            else:
                print("使用默认参数")
                lookback = 20
                std_mult = 1.0
        else:
            lookback = 20
            std_mult = 1.0

        # 运行主策略
        strategy = MeanReversionStrategy(
            initial_cash=100000,
            initial_shares=10000,
            lookback_days=lookback,
            std_multiplier=std_mult,
            trade_volume=400
        )

        results = strategy.run_strategy(df, train_days=3000)

        print("\n" + "=" * 80)
        print("策略执行完成！")
        print("=" * 80)

    except FileNotFoundError:
        print("错误：找不到 '个股数据.csv' 文件！")
        print("请确保数据文件在当前目录下。")
    except Exception as e:
        print(f"运行过程中出现错误：{e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()