
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Any
from abc import ABC, abstractmethod
import warnings
warnings.filterwarnings('ignore')

class Strategy(ABC):
    """策略基类"""
    
    def __init__(self):
        self.name = self.__class__.__name__
    
    @abstractmethod
    def on_bar(self, data: Dict[str, pd.DataFrame], broker, current_date: datetime):
        """
        每个交易日调用的策略逻辑
        
        Args:
            data: 历史数据 {股票代码: DataFrame}
            broker: 交易接口
            current_date: 当前日期
        """
        pass
    
    def on_start(self, broker):
        """回测开始时调用"""
        pass
    
    def on_end(self, broker):
        """回测结束时调用"""
        pass

class Broker:
    """简化的交易接口"""
    
    def __init__(self, initial_cash: float = 1000000, commission_rate: float = 0.0003):
        self.initial_cash = initial_cash
        self.cash = initial_cash
        self.commission_rate = commission_rate
        self.positions = {}  # {股票代码: 数量}
        self.avg_costs = {}  # {股票代码: 平均成本}
        self.trades = []     # 交易记录
        self.daily_values = []  # 每日净值
        
    def buy(self, symbol: str, amount: float, price: Optional[float] = None) -> bool:
        """
        买入股票
        
        Args:
            symbol: 股票代码
            amount: 金额或数量，如果 < 1 则视为金额比例，否则视为股数
            price: 价格，为None时使用当前价格
        """
        if amount <= 0:
            return False
            
        if amount < 1:  # 按比例买入
            buy_amount = self.cash * amount
            shares = int(buy_amount / price / 100) * 100  # 按手买入
        else:  # 按股数买入
            shares = int(amount / 100) * 100  # 按手买入
            buy_amount = shares * price
            
        commission = buy_amount * self.commission_rate
        total_cost = buy_amount + commission
        
        if self.cash >= total_cost and shares >= 100:
            # 更新持仓
            old_shares = self.positions.get(symbol, 0)
            old_cost = self.avg_costs.get(symbol, 0.0)
            
            new_shares = old_shares + shares
            new_avg_cost = ((old_shares * old_cost) + buy_amount) / new_shares
            
            self.positions[symbol] = new_shares
            self.avg_costs[symbol] = new_avg_cost
            self.cash -= total_cost
            
            # 记录交易
            self.trades.append({
                'date': datetime.now(),
                'symbol': symbol,
                'action': 'buy',
                'shares': shares,
                'price': price,
                'amount': buy_amount,
                'commission': commission
            })
            
            return True
        return False
    
    def sell(self, symbol: str, amount: float = 1.0, price: Optional[float] = None) -> bool:
        """
        卖出股票
        
        Args:
            symbol: 股票代码  
            amount: 比例或数量，如果 <= 1 则视为比例，否则视为股数
            price: 价格，为None时使用当前价格
        """
        current_shares = self.positions.get(symbol, 0)
        if current_shares == 0:
            return False
            
        if amount <= 1:  # 按比例卖出
            sell_shares = int(current_shares * amount / 100) * 100
        else:  # 按股数卖出
            sell_shares = min(int(amount / 100) * 100, current_shares)
            
        if sell_shares <= 0:
            return False
            
        sell_amount = sell_shares * price
        commission = sell_amount * self.commission_rate
        net_proceeds = sell_amount - commission
        
        # 更新持仓
        self.positions[symbol] = current_shares - sell_shares
        if self.positions[symbol] == 0:
            del self.positions[symbol]
            if symbol in self.avg_costs:
                del self.avg_costs[symbol]
                
        self.cash += net_proceeds
        
        # 记录交易
        self.trades.append({
            'date': datetime.now(),
            'symbol': symbol,
            'action': 'sell',
            'shares': sell_shares,
            'price': price,
            'amount': sell_amount,
            'commission': commission
        })
        
        return True
    
    def get_position(self, symbol: str) -> int:
        """获取持仓数量"""
        return self.positions.get(symbol, 0)
    
    def get_cash(self) -> float:
        """获取可用现金"""
        return self.cash
    
    def get_total_value(self, prices: Dict[str, float]) -> float:
        """计算总资产"""
        stock_value = sum(shares * prices.get(symbol, 0) 
                         for symbol, shares in self.positions.items())
        return self.cash + stock_value
    
    def get_positions(self) -> Dict[str, int]:
        """获取所有持仓"""
        return self.positions.copy()

def backtest(data, strategy: Strategy, broker: Optional[Broker] = None, 
             start_time: str = '', end_time: str = '') -> Dict[str, Any]:
    """
    简化的回测函数
    
    Args:
        data: 数据对象或数据字典
        strategy: 策略对象
        broker: 交易接口，为None时使用默认设置
        start_time: 开始时间 'YYYY-MM-DD'
        end_time: 结束时间 'YYYY-MM-DD'
    
    Returns:
        回测结果字典
    
    Examples:
        # 基本用法
        from qka.core.data import data
        from qka.core.backtest import backtest, Strategy
        
        class MyStrategy(Strategy):
            def on_bar(self, data, broker, current_date):
                for symbol, df in data.items():
                    if len(df) >= 20:
                        price = df['close'].iloc[-1]
                        ma20 = df['close'].rolling(20).mean().iloc[-1]
                        
                        if price > ma20 and broker.get_position(symbol) == 0:
                            broker.buy(symbol, 0.2, price)  # 买入20%资金
                        elif price < ma20 and broker.get_position(symbol) > 0:
                            broker.sell(symbol, 1.0, price)  # 全部卖出
        
        # 获取数据
        data_obj = data('akshare', stocks=['000001', '600000'])
        hist_data = data_obj.get(start_time='2023-01-01', end_time='2023-12-31')
        
        # 运行回测
        result = backtest(hist_data, MyStrategy(), start_time='2023-01-01', end_time='2023-12-31')
        print(f"总收益率: {result['total_return']:.2%}")
    """
      # 处理数据
    if hasattr(data, 'get'):  # 数据对象
        historical_data = data.get(period='1d', start_time=start_time, end_time=end_time)
    else:  # 直接传入的数据字典
        historical_data = data
    
    if not historical_data:
        raise ValueError("未获取到有效数据")
    
    # 初始化broker
    if broker is None:
        broker = Broker()
    
    # 获取所有交易日期
    all_dates = set()
    for symbol, df in historical_data.items():
        if isinstance(df, pd.DataFrame) and not df.empty:
            dates = df.index
            if start_time:
                start_dt = pd.to_datetime(start_time)
                dates = dates[dates >= start_dt]
            if end_time:
                end_dt = pd.to_datetime(end_time)
                dates = dates[dates <= end_dt]
            all_dates.update(dates)
    
    trading_days = sorted(list(all_dates))
    
    if not trading_days:
        raise ValueError("没有找到有效的交易日期")
    
    # 策略初始化
    strategy.on_start(broker)
    
    # 逐日回测
    for current_date in trading_days:
        broker.current_date = current_date
        
        # 准备当日数据（截止到当前日期的历史数据）
        current_data = {}
        current_prices = {}
        
        for symbol, df in historical_data.items():
            if current_date in df.index:
                # 获取截止到当前日期的所有历史数据
                historical_slice = df.loc[:current_date].copy()
                current_data[symbol] = historical_slice
                current_prices[symbol] = df.loc[current_date, 'close']
        
        # 执行策略
        if current_data:
            try:
                strategy.on_bar(current_data, broker, current_date)
            except Exception as e:
                print(f"策略执行错误 {current_date}: {e}")
                continue
        
        # 记录当日资产状况
        total_value = broker.get_total_value(current_prices)
        broker.daily_values.append({
            'date': current_date,
            'cash': broker.cash,
            'total_value': total_value,
            'positions': broker.get_positions().copy(),
            'prices': current_prices.copy()
        })
    
    # 策略结束
    strategy.on_end(broker)
    
    # 计算回测结果
    return _calculate_performance(broker)

def _calculate_performance(broker: Broker) -> Dict[str, Any]:
    """计算回测绩效指标"""
    
    if not broker.daily_values:
        return {
            'error': '没有有效的回测数据'
        }
    
    # 提取数据
    dates = [record['date'] for record in broker.daily_values]
    values = [record['total_value'] for record in broker.daily_values]
    
    # 基本统计
    initial_value = broker.initial_cash
    final_value = values[-1] if values else initial_value
    total_return = (final_value - initial_value) / initial_value
    
    # 计算年化收益率
    days = len(values)
    if days > 0:
        annual_return = (final_value / initial_value) ** (252 / days) - 1
    else:
        annual_return = 0
    
    # 计算波动率和夏普比率
    returns = pd.Series(values).pct_change().dropna()
    if len(returns) > 1:
        volatility = returns.std() * np.sqrt(252)
        sharpe_ratio = annual_return / volatility if volatility > 0 else 0
    else:
        volatility = 0
        sharpe_ratio = 0
    
    # 计算最大回撤
    peak = pd.Series(values).expanding().max()
    drawdown = (pd.Series(values) - peak) / peak
    max_drawdown = drawdown.min()
    
    # 交易统计
    trades = broker.trades
    total_trades = len(trades)
    total_commission = sum(trade['commission'] for trade in trades)
    
    # 简化的胜率计算
    profitable_trades = 0
    if total_trades > 0:
        buy_trades = [t for t in trades if t['action'] == 'buy']
        sell_trades = [t for t in trades if t['action'] == 'sell']
        
        for sell_trade in sell_trades:
            symbol = sell_trade['symbol']
            sell_date = sell_trade['date']
            
            # 查找最近的买入记录
            recent_buy = None
            for buy_trade in reversed(buy_trades):
                if (buy_trade['symbol'] == symbol and 
                    buy_trade['date'] <= sell_date):
                    recent_buy = buy_trade
                    break
            
            if recent_buy:
                profit = (sell_trade['price'] - recent_buy['price']) * sell_trade['shares']
                profit -= (sell_trade['commission'] + recent_buy['commission'])
                if profit > 0:
                    profitable_trades += 1
    
    win_rate = profitable_trades / len(sell_trades) if sell_trades else 0
    
    return {
        'initial_capital': initial_value,
        'final_value': final_value,
        'total_return': total_return,
        'annual_return': annual_return,
        'volatility': volatility,
        'sharpe_ratio': sharpe_ratio,
        'max_drawdown': max_drawdown,
        'total_trades': total_trades,
        'total_commission': total_commission,
        'win_rate': win_rate,
        'trading_days': days,
        'daily_values': broker.daily_values,
        'trades': trades,
        'positions': broker.get_positions()
    }