import pandas as pd
import numpy as np
from datetime import datetime
import os


class StockTradingStrategy:
    def __init__(self, initial_cash=100000, initial_shares=10000, trade_volume=400):
        """
        初始化交易策略

        Args:
            initial_cash: 初始现金
            initial_shares: 初始持股数量
            trade_volume: 每次买卖的股数
        """
        self.cash = initial_cash
        self.shares = initial_shares
        self.trade_volume = trade_volume

        # 交易费用设置
        self.stamp_tax_rate = 0.001  # 千分之一印花税（卖出时收取）
        self.commission = 5  # 佣金（每笔固定5元）
        self.transfer_fee = 1  # 过户费（每笔固定1元）

        # 交易记录
        self.trade_records = []

        # 历史持仓记录
        self.position_history = []

    def calculate_fees(self, price, volume, is_buy):
        """
        计算交易费用

        Args:
            price: 交易价格
            volume: 交易数量
            is_buy: 是否为买入

        Returns:
            total_fee: 总费用
        """
        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):
        """
        执行交易

        Args:
            trade_date: 交易日期
            price: 交易价格
            is_buy: 是否为买入

        Returns:
            bool: 交易是否成功执行
        """
        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,
            '价格': price,
            '数量': volume,
            '金额': total_amount,
            '费用': fee,
            '现金余额': round(self.cash, 2),
            '持仓数量': self.shares
        }
        self.trade_records.append(record)

        return True

    def run_strategy(self, df):
        """
        运行交易策略

        Args:
            df: 包含股票数据的DataFrame

        Returns:
            results: 交易结果汇总
        """
        print("开始运行交易策略...")
        print(f"初始现金: {self.cash:.2f} 元")
        print(f"初始持股: {self.shares} 股")
        print(f"每次交易: {self.trade_volume} 股")
        print("-" * 60)

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

        # 存储每日持仓价值
        daily_positions = []

        for i in range(len(df)):
            current_row = df.iloc[i]
            current_date = current_row['日期']
            open_price = current_row['开盘价(元)']

            # 记录当前持仓
            current_position_value = self.shares * open_price + self.cash
            daily_positions.append({
                '日期': current_date,
                '持仓数量': self.shares,
                '现金余额': round(self.cash, 2),
                '持仓市值': round(self.shares * open_price, 2),
                '总资产': round(current_position_value, 2)
            })

            # 第一天的开盘价没有前一天比较，跳过交易判断
            if i == 0:
                continue

            prev_row = df.iloc[i - 1]
            prev_open_price = prev_row['开盘价(元)']

            # 计算开盘价变化率
            if prev_open_price == 0:
                continue

            price_change_ratio = (open_price - prev_open_price) / prev_open_price

            # 判断是否触发交易条件
            if price_change_ratio > 0.05:
                # 上涨超过5%，执行卖出
                self.execute_trade(current_date, open_price, is_buy=False)
            elif price_change_ratio < -0.05:
                # 下跌超过5%，执行买入
                self.execute_trade(current_date, open_price, is_buy=True)

        # 生成交易汇总
        total_trades = len(self.trade_records)
        buy_trades = sum(1 for record in self.trade_records if record['交易类型'] == '买入')
        sell_trades = sum(1 for record in self.trade_records if record['交易类型'] == '卖出')

        final_cash = self.cash
        final_shares = self.shares
        final_price = df.iloc[-1]['收盘价(元)'] if not df.empty else 0
        final_value = final_cash + final_shares * final_price

        results = {
            '交易次数': total_trades,
            '买入次数': buy_trades,
            '卖出次数': sell_trades,
            '最终现金': round(final_cash, 2),
            '最终持股': final_shares,
            '最终总资产': round(final_value, 2),
            '总盈亏': round(final_value - 100000, 2),
            '总盈亏率': round((final_value / 100000 - 1) * 100, 2)
        }

        # 输出结果
        print("\n" + "=" * 60)
        print("交易策略运行完成！")
        print("=" * 60)
        print(f"总交易次数: {total_trades}")
        print(f"买入次数: {buy_trades}")
        print(f"卖出次数: {sell_trades}")
        print(f"最终现金: {final_cash:.2f} 元")
        print(f"最终持股: {final_shares} 股")
        print(f"最终总资产: {final_value:.2f} 元")
        print(f"总盈亏: {final_value - 100000:.2f} 元")
        print(f"总盈亏率: {(final_value / 100000 - 1) * 100:.2f}%")
        print("=" * 60)

        # 保存交易记录
        if self.trade_records:
            trade_df = pd.DataFrame(self.trade_records)
            trade_df.to_csv('交易记录.csv', index=False, encoding='utf-8-sig')
            print(f"\n交易记录已保存到 '交易记录.csv'")

        # 保存每日持仓
        if daily_positions:
            position_df = pd.DataFrame(daily_positions)
            position_df.to_csv('每日持仓.csv', index=False, encoding='utf-8-sig')
            print(f"每日持仓已保存到 '每日持仓.csv'")

        return results


def main():
    """
    主函数：读取数据并执行策略
    """
    # 检查数据文件是否存在
    if not os.path.exists('个股数据.csv'):
        print("错误：找不到 '个股数据.csv' 文件！")
        print("请确保数据文件在当前目录下。")
        return

    try:
        # 读取数据
        df = pd.read_csv('个股数据.csv', encoding='utf-8')

        # 检查必要的字段
        required_columns = ['日期', '前收盘价(元)', '开盘价(元)', '最高价(元)',
                            '最低价(元)', '收盘价(元)', '成交量(股)', '成交金额(元)']

        for col in required_columns:
            if col not in df.columns:
                print(f"错误：缺少必要的字段 '{col}'")
                return

        # 创建策略实例并运行
        strategy = StockTradingStrategy(
            initial_cash=100000,
            initial_shares=10000,
            trade_volume=400
        )

        results = strategy.run_strategy(df)

        # 打印详细结果
        print("\n详细交易记录（部分）:")
        if strategy.trade_records:
            trade_df = pd.DataFrame(strategy.trade_records)
            print(trade_df.tail(10).to_string(index=False))

    except Exception as e:
        print(f"运行过程中出现错误：{e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()