import pandas as pd
import numpy as np
import os


class StockTradingStrategy:
    def __init__(self, initial_cash=100000, initial_shares=10000, trade_volume=400, threshold=0.05):
        self.initial_cash = initial_cash
        self.cash = initial_cash
        self.shares = initial_shares
        self.trade_volume = trade_volume
        self.threshold = threshold  # 将阈值提取为参数
        self.stamp_tax_rate = 0.001
        self.commission = 5
        self.transfer_fee = 1

    def calculate_fees(self, price, volume, is_buy):
        total_amount = price * volume
        if is_buy:
            return self.commission + self.transfer_fee
        else:
            return self.commission + self.transfer_fee + (total_amount * self.stamp_tax_rate)

    def run_backtest(self, df):
        # 重置环境
        self.cash = self.initial_cash
        self.shares = 10000

        # 确保按日期排序
        df = df.sort_values('日期').reset_index(drop=True)

        for i in range(1, len(df)):
            open_price = df.iloc[i]['开盘价(元)']
            prev_open_price = df.iloc[i - 1]['开盘价(元)']

            if prev_open_price == 0: continue
            price_change = (open_price - prev_open_price) / prev_open_price

            # 卖出逻辑
            if price_change > self.threshold:
                if self.shares >= self.trade_volume:
                    fee = self.calculate_fees(open_price, self.trade_volume, False)
                    self.cash += (open_price * self.trade_volume) - fee
                    self.shares -= self.trade_volume

            # 买入逻辑
            elif price_change < -self.threshold:
                cost = (open_price * self.trade_volume) + self.calculate_fees(open_price, self.trade_volume, True)
                if self.cash >= cost:
                    self.cash -= cost
                    self.shares += self.trade_volume

        # 计算最终市值
        final_price = df.iloc[-1]['收盘价(元)']
        return self.cash + (self.shares * final_price)


def run_parameter_sweep():
    if not os.path.exists('个股数据.csv'):
        print("错误：找不到 '个股数据.csv'")
        return

    df = pd.read_csv('个股数据.csv', encoding='utf-8')

    # 参数空间
    volumes = range(400, 1100, 100)  # 400 到 1000
    thresholds = [round(x, 2) for x in np.arange(0.01, 0.11, 0.01)]  # 1% 到 10%

    results = []

    print(f"开始进行 {len(volumes) * len(thresholds)} 次参数组合扫描...")

    for vol in volumes:
        for thr in thresholds:
            strategy = StockTradingStrategy(trade_volume=vol, threshold=thr)
            final_assets = strategy.run_backtest(df)
            results.append({
                '每次交易股数': vol,
                '涨跌阈值': f"{int(thr * 100)}%",
                '最终资产': round(final_assets, 2),
                '总收益率': round((final_assets / 100000 - 1) * 100, 2)
            })

    # 转换为 DataFrame 并整理为矩阵
    results_df = pd.DataFrame(results)
    pivot_table = results_df.pivot(index='每次交易股数', columns='涨跌阈值', values='总收益率')
    # --- 新增代码：解除显示截断限制 ---
    pd.set_option('display.max_columns', None)  # 显示所有列
    pd.set_option('display.max_rows', None)  # 显示所有行
    pd.set_option('display.width', 1000)  # 设置显示宽度，防止换行
    # ----------------------------------

    print("\n=== 参数交叉寻优结果（总收益率 %）===")
    print(pivot_table)
    print("\n💡 提示：收益率最高的单元格即为在该历史行情下的最优参数组合。")

    pivot_table.to_csv('参数寻优对比结果.csv')


if __name__ == "__main__":
    run_parameter_sweep()