import pandas as pd
import numpy as np

from cycles.market_fees import MarketFees

class Backtest:
    class Data:
        def __init__(self, initial_usd, df, min1_df, init_strategy_fields) -> None:
            self.initial_usd = initial_usd
            self.usd = initial_usd
            self.max_balance = initial_usd
            self.coin = 0
            self.position = 0
            self.entry_price = 0
            self.entry_time = None
            self.current_trade_min1_start_idx = None
            self.current_min1_end_idx = None
            self.price_open = None
            self.price_close = None
            self.current_date = None
            self.strategies = {}
            self.df = df
            self.min1_df = min1_df

            self = init_strategy_fields(self)

    @staticmethod
    def run(data, entry_strategy, exit_strategy, debug=False):
        """
        Backtest a simple strategy using the meta supertrend (all three supertrends agree).
        Buys when meta supertrend is positive, sells when negative, applies a percentage stop loss.
        
        Parameters:
        - min1_df: pandas DataFrame, 1-minute timeframe data for more accurate stop loss checking (optional)
        - initial_usd: float, starting USD amount
        - stop_loss_pct: float, stop loss as a fraction (e.g. 0.05 for 5%)
        - debug: bool, whether to print debug info
        """
        trade_log = []
        drawdowns = []
        trades = []

        for i in range(1, len(data.df)):
            data.price_open = data.df['open'].iloc[i]
            data.price_close = data.df['close'].iloc[i]
            
            data.current_date = data.df['timestamp'].iloc[i]

            if data.position == 0:
                if entry_strategy(data, i):
                    data, entry_log_entry = Backtest.handle_entry(data)
                    trade_log.append(entry_log_entry)
            elif data.position == 1:
                exit_test_results, data, sell_price = exit_strategy(data, i)
                
                if exit_test_results is not None:
                    data, exit_log_entry = Backtest.handle_exit(data, exit_test_results, sell_price)
                    trade_log.append(exit_log_entry)

            # Track drawdown
            balance = data.usd if data.position == 0 else data.coin * data.price_close

            if balance > data.max_balance:
                data.max_balance = balance

            drawdown = (data.max_balance - balance) / data.max_balance
            drawdowns.append(drawdown)

        # If still in position at end, sell at last close
        if data.position == 1:
            data, exit_log_entry = Backtest.handle_exit(data, "EOD", None)
            trade_log.append(exit_log_entry)

        # Calculate statistics
        final_balance = data.usd
        n_trades = len(trade_log)
        wins = [1 for t in trade_log if t['exit'] is not None and t['exit'] > t['entry']]
        win_rate = len(wins) / n_trades if n_trades > 0 else 0
        max_drawdown = max(drawdowns) if drawdowns else 0
        avg_trade = np.mean([t['exit']/t['entry']-1 for t in trade_log if t['exit'] is not None]) if trade_log else 0

        trades = []
        total_fees_usd = 0.0
        for trade in trade_log:
            if trade['exit'] is not None:
                profit_pct = (trade['exit'] - trade['entry']) / trade['entry']
            else:
                profit_pct = 0.0
            trades.append({
                'entry_time': trade['entry_time'],
                'exit_time': trade['exit_time'],
                'entry': trade['entry'],
                'exit': trade['exit'],
                'profit_pct': profit_pct,
                'type': trade['type'],
                'fee_usd': trade['fee_usd']
            })
            fee_usd = trade.get('fee_usd')
            total_fees_usd += fee_usd

        results = {
            "initial_usd": data.initial_usd,
            "final_usd": final_balance,
            "n_trades": n_trades,
            "win_rate": win_rate,
            "max_drawdown": max_drawdown,
            "avg_trade": avg_trade,
            "trade_log": trade_log,
            "trades": trades,
            "total_fees_usd": total_fees_usd,
        }
        if n_trades > 0:
            results["first_trade"] = {
                "entry_time": trade_log[0]['entry_time'],
                "entry": trade_log[0]['entry']
            }
            results["last_trade"] = {
                "exit_time": trade_log[-1]['exit_time'],
                "exit": trade_log[-1]['exit']
            }
        return results

    @staticmethod
    def handle_entry(data):
        entry_fee = MarketFees.calculate_okx_taker_maker_fee(data.usd, is_maker=False)
        usd_after_fee = data.usd - entry_fee
        
        data.coin = usd_after_fee / data.price_open
        data.entry_price = data.price_open
        data.entry_time = data.current_date
        data.usd = 0
        data.position = 1
        
        trade_log_entry = {
            'type': 'BUY',
            'entry': data.entry_price,
            'exit': None,
            'entry_time': data.entry_time,
            'exit_time': None,
            'fee_usd': entry_fee
        }
        return data, trade_log_entry

    @staticmethod
    def handle_exit(data, exit_reason, sell_price):
        btc_to_sell = data.coin
        exit_price = sell_price if sell_price is not None else data.price_open
        usd_gross = btc_to_sell * exit_price
        exit_fee = MarketFees.calculate_okx_taker_maker_fee(usd_gross, is_maker=False)

        data.usd = usd_gross - exit_fee

        exit_log_entry = {
            'type': exit_reason,
            'entry': data.entry_price,
            'exit': exit_price,
            'entry_time': data.entry_time,
            'exit_time': data.current_date,
            'fee_usd': exit_fee
        }
        data.coin = 0
        data.position = 0
        data.entry_price = 0

        return data, exit_log_entry