import backtrader as bt
import math

class NAS100ScalpingStrategy(bt.Strategy):
    params = (
        ("lookback_period", 20),
        ("volume_threshold", 1.5),
        ("momentum_threshold", 0.02),
        ("risk_percent", 0.01),    # 1% risk per trade
        ("max_daily_loss", 0.03),  # 3% max daily loss
        ("atr_period", 14),
        ("initial_rr_ratio", 1.5),  # Initial risk-reward ratio
        ("trail_start_rr", 1.0),    # Start trailing after 1:1 R:R
        ("trail_atr_mult", 1.5),    # Trailing stop ATR multiplier
        ("partial_take_rr", 2.0),   # Take partial profits at 2:1 R:R
        ("partial_take_percent", 0.5),  # Take 50% of position as partial
        # New parameter for volatility expansion threshold
        ("volatility_threshold", 1.8),
    )
    
    def __init__(self):
        # Existing technical indicators
        self.vol_sma = bt.indicators.SimpleMovingAverage(self.datas[0].volume, period=20)
        self.immediate_support = bt.indicators.Lowest(self.datas[0].low, period=10)
        self.immediate_resistance = bt.indicators.Highest(self.datas[0].high, period=10)
        self.avg_range = bt.indicators.SimpleMovingAverage(
            self.datas[0].high - self.datas[0].low, period=self.p.lookback_period
        )
        self.avg_range10 = bt.indicators.SimpleMovingAverage(
            self.datas[0].high - self.datas[0].low, period=10
        )
        self.atr = bt.indicators.ATR(self.datas[0], period=self.p.atr_period)
        self.rsi = bt.indicators.RSI(self.datas[0].close, period=14)
        self.ema_fast = bt.indicators.EMA(self.datas[0].close, period=9)
        self.ema_slow = bt.indicators.EMA(self.datas[0].close, period=21)
        self.ema_trend = bt.indicators.EMA(self.datas[0].close, period=50)  # Longer term trend
        
        # State variables for signals and risk management
        self.prev_consolidation = False
        self.last_roc = None
        self.order = None
        self.position_type = 0  # 1 for long, -1 for short
        self.entry_price = None
        self.initial_stop_loss = None
        self.trailing_stop = None
        self.initial_take_profit = None
        self.partial_taken = False
        self.risk_amount = 0
        
        # Daily loss tracking and performance
        self.daily_pnl = 0.0
        self.current_date = None
        self.trade_count = 0
        self.winning_trades = 0
        self.max_winner_r = 0
        
        # NEW: Leading indicator tracking lists
        self.range_history = []
        self.tick_direction = []  # To track tick-by-tick direction

    def next(self):
        # Update daily pnl and date
        current_day = self.datas[0].datetime.date(0)
        if self.current_date != current_day:
            self.daily_pnl = 0.0
            self.current_date = current_day
        
        # Daily loss protection
        if self.daily_pnl <= -self.broker.getvalue() * self.p.max_daily_loss:
            if self.position:
                self.close()
                self.log("Daily loss limit reached. Closing position.")
            return
        
        # Avoid processing if an order is pending
        if self.order:
            return
        
        # Ensure sufficient data history
        if len(self.datas[0]) < max(self.p.lookback_period, self.p.atr_period, 50):
            return
        
        # NEW: Update extra leading indicators (range and tick direction)
        self._update_leading_indicators()
        
        # Calculate risk per trade
        account_value = self.broker.getvalue()
        self.risk_amount = account_value * self.p.risk_percent
        
        # Original momentum and acceleration calculations
        roc1 = self._calculate_momentum()
        acceleration = self._calculate_acceleration(roc1)
        
        # Gather original signals
        momentum_signal = self._get_momentum_signal(roc1, acceleration)
        volume_signal = self._get_volume_signal()
        breakout_signal = self._get_breakout_signal()
        trend_signal = self._get_trend_signal()
        rsi_signal = self._get_rsi_signal()
        long_trend_signal = self._get_long_trend_signal()
        
        # Count signals from original indicators
        bullish_signals, bearish_signals = self._count_signals(
            momentum_signal, volume_signal, breakout_signal, trend_signal, rsi_signal, long_trend_signal
        )
        
        # NEW: Compute additional leading indicator signals
        momentum_accel = self._get_momentum_acceleration_signal(acceleration)
        vol_expansion = self._get_volatility_expansion_signal()
        order_flow = self._get_order_flow_signal()
        momentum_shift = self._get_momentum_shift_signal()
        structure_signal = self._get_price_structure_signal()
        
        # Add extra signals into totals
        if momentum_accel > 0:
            bullish_signals += 1
        elif momentum_accel < 0:
            bearish_signals += 1
            
        if vol_expansion > 0:
            bullish_signals += 1
        elif vol_expansion < 0:
            bearish_signals += 1
            
        if order_flow > 0:
            bullish_signals += 1
        elif order_flow < 0:
            bearish_signals += 1
            
        if momentum_shift > 0:
            bullish_signals += 1
        elif momentum_shift < 0:
            bearish_signals += 1
            
        if structure_signal > 0:
            bullish_signals += 1
        elif structure_signal < 0:
            bearish_signals += 1
        
        # Decide on entry if not in a position
        if not self.position:
            self._check_entry_signals(bullish_signals, bearish_signals)
        else:
            self._manage_position()

    def _update_leading_indicators(self):
        # Update range_history with current candle range
        current_range = self.datas[0].high[0] - self.datas[0].low[0]
        self.range_history.append(current_range)
        if len(self.range_history) > 10:
            self.range_history.pop(0)
        
        # Update tick_direction based on price change from previous close
        if len(self.datas[0]) >= 2:
            if self.datas[0].close[0] > self.datas[0].close[-1]:
                tick = 1
            elif self.datas[0].close[0] < self.datas[0].close[-1]:
                tick = -1
            else:
                tick = 0
            self.tick_direction.append(tick)
            if len(self.tick_direction) > 5:
                self.tick_direction.pop(0)
    
    def _calculate_momentum(self):
        if len(self.datas[0]) > 1 and self.datas[0].close[-1] > 0:
            return (self.datas[0].close[0] - self.datas[0].close[-1]) / self.datas[0].close[-1]
        return 0

    def _calculate_acceleration(self, roc1):
        acceleration = 0
        if self.last_roc is not None:
            acceleration = roc1 - self.last_roc
        self.last_roc = roc1
        return acceleration

    def _get_momentum_signal(self, roc1, acceleration):
        if roc1 > self.p.momentum_threshold and acceleration > 0:
            return 1
        elif roc1 < -self.p.momentum_threshold and acceleration < 0:
            return -1
        return 0

    def _get_volume_signal(self):
        # Original volume signal based on SMA
        if self.vol_sma[0] <= 0:
            return 0
        vol_spike = self.datas[0].volume[0] / self.vol_sma[0]
        if vol_spike > self.p.volume_threshold:
            if self.datas[0].close[0] > self.datas[0].open[0]:
                return 1
            elif self.datas[0].close[0] < self.datas[0].open[0]:
                return -1
        return 0

    def _get_breakout_signal(self):
        cur_range = self.datas[0].high[0] - self.datas[0].low[0]
        avg_range_val = self.avg_range[0] if self.avg_range[0] > 0 else 1
        cur_compression = cur_range / avg_range_val
        consolidation = cur_compression < 0.7
        
        breakout_signal = 0
        if self.prev_consolidation and not consolidation:
            if self.datas[0].close[0] > self.datas[0].open[0]:
                breakout_signal = 1
            elif self.datas[0].close[0] < self.datas[0].open[0]:
                breakout_signal = -1
        self.prev_consolidation = consolidation
        return breakout_signal

    def _get_trend_signal(self):
        if self.ema_fast[0] > self.ema_slow[0]:
            return 1
        elif self.ema_fast[0] < self.ema_slow[0]:
            return -1
        return 0

    def _get_long_trend_signal(self):
        if self.datas[0].close[0] > self.ema_trend[0]:
            return 1
        elif self.datas[0].close[0] < self.ema_trend[0]:
            return -1
        return 0

    def _get_rsi_signal(self):
        if self.rsi[0] < 35:
            return 1
        elif self.rsi[0] > 65:
            return -1
        return 0

    def _count_signals(self, momentum_signal, volume_signal, breakout_signal, trend_signal, rsi_signal, long_trend_signal):
        bullish_signals = 0
        bearish_signals = 0
        for signal in [momentum_signal, volume_signal, breakout_signal, trend_signal, rsi_signal, long_trend_signal]:
            if signal == 1:
                bullish_signals += 1
            elif signal == -1:
                bearish_signals += 1
        
        # Add support/resistance confluence
        if self.datas[0].close[0] > self.immediate_resistance[0]:
            bullish_signals += 1
        elif self.datas[0].close[0] < self.immediate_support[0]:
            bearish_signals += 1
            
        return bullish_signals, bearish_signals

    # NEW: Additional leading indicator methods
    def _get_momentum_acceleration_signal(self, acceleration):
        if acceleration > 0.0001:
            return 1
        elif acceleration < -0.0001:
            return -1
        return 0
    
    def _get_volatility_expansion_signal(self):
        if len(self.range_history) < 5:
            return 0
        current_range = self.range_history[-1]
        avg_range = sum(self.range_history[:-1]) / (len(self.range_history) - 1)
        if avg_range > 0 and current_range/avg_range > self.p.volatility_threshold:
            # When price is near the top of the candle, consider it bullish expansion,
            # and bearish if near the bottom.
            current_price = self.datas[0].close[0]
            range_pos = (current_price - self.datas[0].low[0]) / (self.datas[0].high[0] - self.datas[0].low[0])
            if range_pos > 0.7:
                return 1
            elif range_pos < 0.3:
                return -1
        return 0

    def _get_order_flow_signal(self):
        # Use volume weighted by candle position as a proxy for order flow pressure.
        if (self.datas[0].high[0] - self.datas[0].low[0]) == 0:
            return 0
        close_pos = (self.datas[0].close[0] - self.datas[0].low[0]) / (self.datas[0].high[0] - self.datas[0].low[0])
        buying_power = close_pos * self.datas[0].volume[0]
        selling_power = (1 - close_pos) * self.datas[0].volume[0]
        if buying_power > selling_power * 1.5:
            return 1
        elif selling_power > buying_power * 1.5:
            return -1
        return 0

    def _get_momentum_shift_signal(self):
        if len(self.tick_direction) < 4:
            return 0
        recent = self.tick_direction[-4:]
        if recent.count(1) >= 3:
            return 1
        elif recent.count(-1) >= 3:
            return -1
        return 0

    def _get_price_structure_signal(self):
        if len(self.datas[0]) < 10:
            return 0
        recent_highs = [self.datas[0].high[-i] for i in range(1, 6)]
        recent_lows = [self.datas[0].low[-i] for i in range(1, 6)]
        current_high = self.datas[0].high[0]
        current_low = self.datas[0].low[0]
        if current_high > max(recent_highs):
            return 1
        elif current_low < min(recent_lows):
            return -1
        return 0

    def _calculate_position_size(self, stop_distance):
        if stop_distance <= 0:
            return 0
        position_value = self.risk_amount / stop_distance
        position_size = position_value / self.datas[0].close[0]
        min_size = 0.01
        max_size = self.broker.getvalue() * 0.1 / self.datas[0].close[0]
        return max(min_size, min(position_size, max_size))

    def _check_entry_signals(self, bullish_signals, bearish_signals):
        current_price = self.datas[0].close[0]
        # For a long entry, require bullish confluence (using a threshold of 4)
        if bullish_signals >= 4 and self.rsi[0] < 70 and current_price > self.ema_trend[0]:
            stop_distance = current_price - (self.immediate_support[0] * 0.9985)
            if stop_distance > 0:
                position_size = self._calculate_position_size(stop_distance)
                if position_size > 0:
                    initial_tp_distance = stop_distance * self.p.initial_rr_ratio
                    self.order = self.buy(size=position_size)
                    self.position_type = 1
                    self.entry_price = current_price
                    self.initial_stop_loss = current_price - stop_distance
                    self.trailing_stop = self.initial_stop_loss
                    self.initial_take_profit = current_price + initial_tp_distance
                    self.partial_taken = False
                    self.log(f"LONG ENTRY: Size={position_size:.4f} @ {current_price:.2f} | SL={self.initial_stop_loss:.2f} | Risk=${self.risk_amount:.2f} | Signals={bullish_signals}")
        # For a short entry, require bearish confluence
        elif bearish_signals >= 4 and self.rsi[0] > 30 and current_price < self.ema_trend[0]:
            stop_distance = (self.immediate_resistance[0] * 1.0015) - current_price
            if stop_distance > 0:
                position_size = self._calculate_position_size(stop_distance)
                if position_size > 0:
                    initial_tp_distance = stop_distance * self.p.initial_rr_ratio
                    self.order = self.sell(size=position_size)
                    self.position_type = -1
                    self.entry_price = current_price
                    self.initial_stop_loss = current_price + stop_distance
                    self.trailing_stop = self.initial_stop_loss
                    self.initial_take_profit = current_price - initial_tp_distance
                    self.partial_taken = False
                    self.log(f"SHORT ENTRY: Size={position_size:.4f} @ {current_price:.2f} | SL={self.initial_stop_loss:.2f} | Risk=${self.risk_amount:.2f} | Signals={bearish_signals}")

    def _manage_position(self):
        current_price = self.datas[0].close[0]
        if self.position_type == 1:
            self._manage_long_position(current_price)
        elif self.position_type == -1:
            self._manage_short_position(current_price)

    def _manage_long_position(self, current_price):
        risk_per_share = self.entry_price - self.initial_stop_loss
        current_profit = current_price - self.entry_price
        r_multiple = current_profit / risk_per_share if risk_per_share > 0 else 0
        if r_multiple > self.max_winner_r:
            self.max_winner_r = r_multiple
        
        # Partial profit-taking
        if (not self.partial_taken and r_multiple >= self.p.partial_take_rr and self.position.size > 0.02):
            partial_size = self.position.size * self.p.partial_take_percent
            self.sell(size=partial_size)
            self.partial_taken = True
            self.log(f"PARTIAL PROFIT: Sold {partial_size:.4f} @ {current_price:.2f} | R={r_multiple:.2f}")
        
        # Trail stop after reaching a set R ratio
        if r_multiple >= self.p.trail_start_rr:
            atr_trail = current_price - (self.atr[0] * self.p.trail_atr_mult)
            breakeven_plus = self.entry_price + (risk_per_share * 0.2)
            new_trailing_stop = max(atr_trail, breakeven_plus, self.trailing_stop)
            if new_trailing_stop > self.trailing_stop:
                self.trailing_stop = new_trailing_stop
                self.log(f"TRAIL UPDATE: New trailing stop @ {self.trailing_stop:.2f} | R={r_multiple:.2f}")
        
        exit_long = False
        exit_reason = ""
        if current_price <= self.trailing_stop:
            exit_long = True
            exit_reason = f"Trailing Stop (R={r_multiple:.2f})"
        elif r_multiple < 0.5:
            if self.rsi[0] > 80 or (self.ema_fast[0] < self.ema_slow[0] and r_multiple < 0):
                exit_long = True
                exit_reason = f"Early Exit (R={r_multiple:.2f})"
        elif r_multiple > 3.0 and current_price < self.ema_fast[0]:
            exit_long = True
            exit_reason = f"Trend Reversal on Big Winner (R={r_multiple:.2f})"
        
        if exit_long:
            self.order = self.close()
            self.log(f"CLOSING LONG @ {current_price:.2f} - {exit_reason}")

    def _manage_short_position(self, current_price):
        risk_per_share = self.initial_stop_loss - self.entry_price
        current_profit = self.entry_price - current_price
        r_multiple = current_profit / risk_per_share if risk_per_share > 0 else 0
        if r_multiple > self.max_winner_r:
            self.max_winner_r = r_multiple
        
        if (not self.partial_taken and r_multiple >= self.p.partial_take_rr and abs(self.position.size) > 0.02):
            partial_size = abs(self.position.size) * self.p.partial_take_percent
            self.buy(size=partial_size)
            self.partial_taken = True
            self.log(f"PARTIAL PROFIT: Covered {partial_size:.4f} @ {current_price:.2f} | R={r_multiple:.2f}")
        
        if r_multiple >= self.p.trail_start_rr:
            atr_trail = current_price + (self.atr[0] * self.p.trail_atr_mult)
            breakeven_plus = self.entry_price - (risk_per_share * 0.2)
            new_trailing_stop = min(atr_trail, breakeven_plus, self.trailing_stop)
            if new_trailing_stop < self.trailing_stop:
                self.trailing_stop = new_trailing_stop
                self.log(f"TRAIL UPDATE: New trailing stop @ {self.trailing_stop:.2f} | R={r_multiple:.2f}")
        
        exit_short = False
        exit_reason = ""
        if current_price >= self.trailing_stop:
            exit_short = True
            exit_reason = f"Trailing Stop (R={r_multiple:.2f})"
        elif r_multiple < 0.5:
            if self.rsi[0] < 20 or (self.ema_fast[0] > self.ema_slow[0] and r_multiple < 0):
                exit_short = True
                exit_reason = f"Early Exit (R={r_multiple:.2f})"
        elif r_multiple > 3.0 and current_price > self.ema_fast[0]:
            exit_short = True
            exit_reason = f"Trend Reversal on Big Winner (R={r_multiple:.2f})"
        
        if exit_short:
            self.order = self.close()
            self.log(f"CLOSING SHORT @ {current_price:.2f} - {exit_reason}")

    def notify_order(self, order):
        if order.status in [order.Completed, order.Canceled, order.Margin]:
            self.order = None

    def notify_trade(self, trade):
        if trade.isclosed:
            self.trade_count += 1
            self.daily_pnl += trade.pnl
            risk_per_share = 0
            r_multiple = 0
            if self.entry_price is not None and self.initial_stop_loss is not None:
                risk_per_share = abs(self.entry_price - self.initial_stop_loss)
                if risk_per_share > 0 and abs(trade.size) > 0:
                    r_multiple = trade.pnl / (risk_per_share * abs(trade.size))
                else:
                    r_multiple = 0
            if trade.pnl > 0:
                self.winning_trades += 1
            win_rate = (self.winning_trades / self.trade_count) * 100 if self.trade_count > 0 else 0
            self.log(f"TRADE CLOSED #{self.trade_count} | PnL: ${trade.pnl:.2f} | R={r_multiple:.2f} | Win Rate: {win_rate:.1f}% | Daily PnL: ${self.daily_pnl:.2f}")
            self.position_type = 0
            self.entry_price = None
            self.initial_stop_loss = None
            self.trailing_stop = None
            self.initial_take_profit = None
            self.partial_taken = False

    def log(self, txt):
        dt = self.datas[0].datetime.date(0)
        t = self.datas[0].datetime.time(0)
        print(f"{dt} {t} | {txt}")

    def stop(self):
        win_rate = (self.winning_trades / self.trade_count) * 100 if self.trade_count > 0 else 0
        print(f"\n=== ENHANCED STRATEGY RESULTS ===")
        print(f"Total Trades: {self.trade_count}")
        print(f"Winning Trades: {self.winning_trades}")
        print(f"Win Rate: {win_rate:.2f}%")
        print(f"Max R Multiple Achieved: {self.max_winner_r:.2f}")
        print(f"Final Portfolio Value: ${self.broker.getvalue():.2f}")
        print(f"================================\n")
ROI

-1.57%

-$1,568.15

Win Ratio

14.96%

146 of 976 positions

Max Drawdown

2.18%

$2,182.74

Time in Market

90.68%

24798 of 27347 candles


IMPROVE THIS BASDE ON RESULTS