import pandas as pd
import pandas_ta as ta
from backtesting import Strategy
from backtesting.lib import crossover


class RsiStrategy(Strategy):
    # Default parameters (can be overridden by optimizer)
    rsi_period = 14
    rsi_lower = 30
    rsi_upper = 70
    # Position sizing and ATR stop parameters
    position_size_pct = 0.1
    min_shares = 1
    atr_period = 14
    stop_loss_atr_multiplier = 3.0

    def init(self):
        self.rsi = self.I(ta.rsi, pd.Series(self.data.Close), length=self.rsi_period)
        self.atr = self.I(
            ta.atr,
            pd.Series(self.data.High),
            pd.Series(self.data.Low),
            pd.Series(self.data.Close),
            length=self.atr_period,
        )

    def next(self):
        if crossover(self.rsi, self.rsi_lower):
            price = float(self.data.Close[-1])
            max_invest = max(0, self.equity * self.position_size_pct)
            shares = int(max_invest // price)
            if shares < self.min_shares:
                shares = self.min_shares

            atr_value = float(self.atr[-1]) if not pd.isna(self.atr[-1]) else None
            if atr_value and atr_value > 0:
                stop_price = round(price - (self.stop_loss_atr_multiplier * atr_value), 6)
                if stop_price >= price:
                    stop_price = round(price * 0.99, 6)
            else:
                stop_price = round(price * 0.98, 6)

            self.buy(size=shares, sl=stop_price)
        elif crossover(self.rsi_upper, self.rsi):
            if self.position:
                self.position.close()


class CrossEmaStrategy(Strategy):
    """EMA crossover strategy with ATR-based stop-loss."""
    short_ema = 12
    long_ema = 26
    position_size_pct = 0.1
    min_shares = 1
    atr_period = 14
    stop_loss_atr_multiplier = 3.0

    def init(self):
        self.ema_short = self.I(ta.ema, pd.Series(self.data.Close), length=self.short_ema)
        self.ema_long = self.I(ta.ema, pd.Series(self.data.Close), length=self.long_ema)
        self.atr = self.I(
            ta.atr,
            pd.Series(self.data.High),
            pd.Series(self.data.Low),
            pd.Series(self.data.Close),
            length=self.atr_period,
        )

    def next(self):
        price = float(self.data.Close[-1])
        # Buy when short EMA crosses above long EMA
        if crossover(self.ema_short, self.ema_long):
            max_invest = max(0, self.equity * self.position_size_pct)
            shares = int(max_invest // price)
            if shares < self.min_shares:
                shares = self.min_shares

            atr_value = float(self.atr[-1]) if not pd.isna(self.atr[-1]) else None
            if atr_value and atr_value > 0:
                stop_price = round(price - (self.stop_loss_atr_multiplier * atr_value), 6)
                if stop_price >= price:
                    stop_price = round(price * 0.99, 6)
            else:
                stop_price = round(price * 0.98, 6)

            self.buy(size=shares, sl=stop_price)

        # Sell when short EMA crosses below long EMA
        elif crossover(self.ema_long, self.ema_short):
            if self.position:
                self.position.close()