import pandas as pd
import numpy as np

from cycles.market_fees import MarketFees

class Backtest:
    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.trade_log = []
        self.drawdowns = []
        self.trades = []

        self = init_strategy_fields(self)

    def run(self, entry_strategy, exit_strategy, debug=False):
        """
        Runs the backtest using provided entry and exit strategy functions.

        The method iterates over the main DataFrame (self.df), simulating trades based on the entry and exit strategies. 
        It tracks balances, drawdowns, and logs each trade, including fees. At the end, it returns a dictionary of performance statistics.

        Parameters:
        - entry_strategy: function, determines when to enter a trade. Should accept (self, i) and return True to enter.
        - exit_strategy: function, determines when to exit a trade. Should accept (self, i) and return (exit_reason, sell_price) or (None, None) to hold.
        - debug: bool, whether to print debug info (default: False)

        Returns:
        - dict with keys: initial_usd, final_usd, n_trades, win_rate, max_drawdown, avg_trade, trade_log, trades, total_fees_usd, and optionally first_trade and last_trade.
        """

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

            # check if we are in buy/sell position
            if self.position == 0:
                if entry_strategy(self, i):
                    self.handle_entry()
            elif self.position == 1:
                exit_test_results, sell_price = exit_strategy(self, i)
                
                if exit_test_results is not None:
                    self.handle_exit(exit_test_results, sell_price)

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

            if balance > self.max_balance:
                self.max_balance = balance

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

        # If still in position at end, sell at last close
        if self.position == 1:
            self.handle_exit("EOD", None)
            

        # Calculate statistics
        final_balance = self.usd
        n_trades = len(self.trade_log)
        wins = [1 for t in self.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(self.drawdowns) if self.drawdowns else 0
        avg_trade = np.mean([t['exit']/t['entry']-1 for t in self.trade_log if t['exit'] is not None]) if self.trade_log else 0

        trades = []
        total_fees_usd = 0.0

        for trade in self.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": self.initial_usd,
            "final_usd": final_balance,
            "n_trades": n_trades,
            "win_rate": win_rate,
            "max_drawdown": max_drawdown,
            "avg_trade": avg_trade,
            "trade_log": self.trade_log,
            "trades": trades,
            "total_fees_usd": total_fees_usd,
        }
        if n_trades > 0:
            results["first_trade"] = {
                "entry_time": self.trade_log[0]['entry_time'],
                "entry": self.trade_log[0]['entry']
            }
            results["last_trade"] = {
                "exit_time": self.trade_log[-1]['exit_time'],
                "exit": self.trade_log[-1]['exit']
            }
        return results

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

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

        self.usd = usd_gross - exit_fee

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

        self.trade_log.append(exit_log_entry)