# region imports
from AlgorithmImports import *
from datetime import time, timedelta
# endregion

class PreMarketGapEMABounceStrategy(QCAlgorithm):
    """
    Pre-market Gap and EMA Bounce Strategy
    
    Strategy Logic:
    1. Calculate pre-market high/low levels (4:00 AM - 9:30 AM ET)
    2. Detect breakouts during regular hours (9:30 AM - 4:00 PM ET)
    3. Wait for EMA bounce confirmation before entry
    4. Manage positions with 2% stop loss and 4% take profit
    """
    
    def initialize(self):
        # Set strategy parameters
        self.set_start_date(2023, 1, 1)
        self.set_end_date(2024, 12, 31)
        self.set_cash(100000)
        
        # Add SPY with minute resolution and extended market hours
        self.spy = self.add_equity("SPY", Resolution.MINUTE, 
                                  extended_market_hours=True).symbol
        
        # Strategy parameters - MATCHING ORIGINAL NOTEBOOK EXACTLY
        self.ema_period = 20  # FIXED: Changed from 9 to 20 to match profitable notebook
        self.confirmation_bars = 1  # Reduced from 2 to 1 for more immediate entries like notebook
        self.stop_loss_pct = 0.02
        self.take_profit_pct = 0.04
        self.max_risk_per_trade = 0.02
        
        # Initialize EMA indicator
        self.ema = ExponentialMovingAverage("EMA", self.ema_period)
        
        # State variables
        self.premarket_high = None
        self.premarket_low = None
        self.current_date = None
        self.breakout_direction = None
        self.breakout_confirmed = False
        self.confirmation_count = 0
        self.waiting_for_bounce = False
        self.position_entry_price = None
        self.stop_loss_price = None
        self.take_profit_price = None
        self.previous_bar = None  # Store previous bar for bounce detection
        self.previous_ema = None  # Store previous EMA value
        self.daily_trades = 0  # Track trades per day
        self.max_daily_trades = 1  # Limit to 1 trade per day like notebook
        self.previous_close = None  # Store previous day's close for gap calculation
        self.entry_date = None  # Track when position was entered for max hold period
        
        # Signal quality filters - more reasonable thresholds
        self.min_gap_size = 0.004  # Minimum 0.4% gap (was too restrictive at 0.5%)
        self.max_gap_size = 0.03   # Maximum 3% gap
        self.min_volume = 800000   # Minimum 800K volume for breakouts (was too high at 1M)
        
        # Market hours
        self.premarket_start = time(4, 0)
        self.market_open = time(9, 30)
        self.market_close = time(16, 0)
        self.entry_cutoff = time(15, 0)
        
        # Tracking variables
        self.premarket_bars = []
        self.last_bar_time = None
        
        # Schedule daily reset
        self.schedule.on(self.date_rules.every_day(self.spy),
                        self.time_rules.at(4, 0),
                        self.reset_daily_variables)
        
        # Schedule premarket level calculation
        self.schedule.on(self.date_rules.every_day(self.spy),
                        self.time_rules.at(9, 25),
                        self.calculate_premarket_levels)
        
        # Schedule end of day exit - only exit losing positions, let winners run overnight
        self.schedule.on(self.date_rules.every_day(self.spy),
                        self.time_rules.before_market_close(self.spy, 1),  # 1 minute before close
                        self.end_of_day_exit)
    
    def reset_daily_variables(self):
        """Reset variables at the start of each trading day"""
        self.premarket_high = None
        self.premarket_low = None
        self.breakout_direction = None
        self.breakout_confirmed = False
        self.confirmation_count = 0
        self.waiting_for_bounce = False
        self.premarket_bars = []
        self.previous_bar = None  # Reset previous bar tracking
        self.previous_ema = None  # Reset previous EMA tracking
        self.daily_trades = 0  # Reset daily trade counter
        self.previous_close = None  # Reset for gap calculation
        # Note: Don't reset entry_date here - let it persist for overnight holds
        
        if self.current_date != self.time.date():
            self.current_date = self.time.date()
            self.debug(f"Starting new trading day: {self.current_date}")
            self.debug(f"Daily trade limit: {self.max_daily_trades}")
    
    def calculate_premarket_levels(self):
        """Calculate premarket high and low levels with gap size filtering"""
        if len(self.premarket_bars) > 0:
            highs = [bar.high for bar in self.premarket_bars]
            lows = [bar.low for bar in self.premarket_bars]
            
            self.premarket_high = max(highs)
            self.premarket_low = min(lows)
            
            # Get previous day's close for gap calculation
            if self.previous_close is None:
                # Get historical data for previous close
                history = self.history(self.spy, 2, Resolution.DAILY)
                if len(history) >= 1:
                    self.previous_close = history['close'].iloc[-1]
            
            # Calculate gap sizes
            if self.previous_close is not None:
                gap_up = (self.premarket_high - self.previous_close) / self.previous_close
                gap_down = (self.previous_close - self.premarket_low) / self.previous_close
                max_gap = max(gap_up, gap_down)
                
                # Only proceed if gap is within acceptable range
                if self.min_gap_size <= max_gap <= self.max_gap_size:
                    self.debug(f"Valid gap detected - PMH: {self.premarket_high:.2f}, PML: {self.premarket_low:.2f}, Gap: {max_gap*100:.2f}%")
                else:
                    self.debug(f"Gap out of range - Gap: {max_gap*100:.2f}%, Min: {self.min_gap_size*100:.1f}%, Max: {self.max_gap_size*100:.1f}%")
                    # Reset premarket levels to prevent trading
                    self.premarket_high = None
                    self.premarket_low = None
                    return
            
            self.debug(f"Premarket levels - High: {self.premarket_high:.2f}, Low: {self.premarket_low:.2f}")
        else:
            self.debug("No premarket data available")
    
    def on_data(self, data):
        """Main data processing function"""
        if self.spy not in data.bars:
            return
        
        bar = data.bars[self.spy]
        current_time = self.time.time()
        
        # Update EMA and store previous values
        if self.ema.is_ready:
            self.previous_ema = self.ema.current.value
        self.ema.update(bar.end_time, bar.close)
        
        # Store premarket bars
        if self.premarket_start <= current_time < self.market_open:
            self.premarket_bars.append(bar)
            self.previous_bar = bar  # Track previous bar
            return
        
        # Process regular market hours only
        if not (self.market_open <= current_time <= self.market_close):
            return
        
        # Skip if we don't have premarket levels yet
        if self.premarket_high is None or self.premarket_low is None:
            return
        
        # Skip if EMA is not ready
        if not self.ema.is_ready:
            self.previous_bar = bar  # Still track previous bar
            return
        
        # Manage existing position
        if self.portfolio.invested:
            self.manage_position(bar)
            self.previous_bar = bar  # Update previous bar
            return
        
        # Skip late entries (avoid entries too close to market open and close)
        if current_time < time(9, 35) or current_time > self.entry_cutoff:
            self.previous_bar = bar  # Update previous bar
            return
        
        # Skip if daily trade limit reached
        if self.daily_trades >= self.max_daily_trades:
            self.previous_bar = bar  # Update previous bar
            return
        
        # Look for breakout signals
        if not self.breakout_confirmed:
            self.check_for_breakout(bar)
        
        # Look for EMA bounce entry
        if self.waiting_for_bounce:
            self.check_for_ema_bounce_entry(bar)
        
        # Update previous bar for next iteration
        self.previous_bar = bar
    
    def check_for_breakout(self, bar):
        """Check for breakout above premarket high or below premarket low with volume filter"""
        # Add volume filter for breakouts
        if bar.volume < self.min_volume:
            return
        
        # Debug logging for breakout detection
        self.debug(f"Checking breakout - PMH: {self.premarket_high:.2f}, PML: {self.premarket_low:.2f}, High: {bar.high:.2f}, Low: {bar.low:.2f}, Close: {bar.close:.2f}, Volume: {bar.volume}")
        
        # Check for breakout above premarket high
        if bar.high > self.premarket_high and self.breakout_direction != "short":
            if self.breakout_direction != "long":
                self.breakout_direction = "long"
                self.confirmation_count = 0
                self.debug(f"Potential LONG breakout detected - High: {bar.high:.2f} > PMH: {self.premarket_high:.2f}")
            
            # Check for confirmation - simplified to just require close in breakout direction
            if bar.close > self.premarket_high:  # Close above PMH for long confirmation
                self.confirmation_count += 1
                self.debug(f"Long breakout confirmation {self.confirmation_count}/{self.confirmation_bars} - Close: {bar.close:.2f} > PMH: {self.premarket_high:.2f}")
                if self.confirmation_count >= self.confirmation_bars:
                    self.breakout_confirmed = True
                    self.waiting_for_bounce = True
                    self.debug(f"LONG breakout CONFIRMED after {self.confirmation_count} bars")

        # Check for breakout below premarket low
        elif bar.low < self.premarket_low and self.breakout_direction != "long":
            if self.breakout_direction != "short":
                self.breakout_direction = "short"
                self.confirmation_count = 0
                self.debug(f"Potential SHORT breakout detected - Low: {bar.low:.2f} < PML: {self.premarket_low:.2f}")
            
            # Check for confirmation - simplified to just require close in breakout direction
            if bar.close < self.premarket_low:  # Close below PML for short confirmation
                self.confirmation_count += 1
                self.debug(f"Short breakout confirmation {self.confirmation_count}/{self.confirmation_bars} - Close: {bar.close:.2f} < PML: {self.premarket_low:.2f}")
                if self.confirmation_count >= self.confirmation_bars:
                    self.breakout_confirmed = True
                    self.waiting_for_bounce = True
                    self.debug(f"SHORT breakout CONFIRMED after {self.confirmation_count} bars")
    
    def check_for_ema_bounce_entry(self, bar):
        """Check for EMA bounce entry signal - EXACT match to original notebook logic"""
        if not self.ema.is_ready or self.previous_bar is None or self.previous_ema is None:
            return False
            
        current_ema = self.ema.current.value
        prev_ema = self.previous_ema
        
        # Debug logging for signal analysis
        self.debug(f"EMA bounce check - Direction: {self.breakout_direction}, Current EMA: {current_ema:.2f}, Prev EMA: {prev_ema:.2f}")
        
        if self.breakout_direction == "long":
            # Realistic EMA bounce conditions - matching actual market behavior
            prev_touched_ema = self.previous_bar.low <= prev_ema * 1.002  # 0.2% buffer - more realistic
            current_above_ema = bar.close > current_ema  # Just needs to be above EMA, no additional buffer
            momentum_up = bar.close > self.previous_bar.close  # Any positive momentum
            
            self.debug(f"Long conditions - Prev touched EMA: {prev_touched_ema} (prev_low: {self.previous_bar.low:.2f} vs {prev_ema * 1.002:.2f})")
            self.debug(f"Current above EMA: {current_above_ema} (close: {bar.close:.2f} vs {current_ema:.2f})")
            self.debug(f"Momentum up: {momentum_up} (close: {bar.close:.2f} vs prev_close: {self.previous_bar.close:.2f})")
            
            if (prev_touched_ema and current_above_ema and momentum_up):
                self.debug(f"LONG EMA bounce signal triggered!")
                self.enter_long_position(bar.close)
        
        elif self.breakout_direction == "short":
            # Realistic EMA bounce conditions - matching actual market behavior
            prev_touched_ema = self.previous_bar.high >= prev_ema * 0.998  # 0.2% buffer - more realistic
            current_below_ema = bar.close < current_ema  # Just needs to be below EMA, no additional buffer
            momentum_down = bar.close < self.previous_bar.close  # Any negative momentum
            
            self.debug(f"Short conditions - Prev touched EMA: {prev_touched_ema} (prev_high: {self.previous_bar.high:.2f} vs {prev_ema * 0.998:.2f})")
            self.debug(f"Current below EMA: {current_below_ema} (close: {bar.close:.2f} vs {current_ema:.2f})")
            self.debug(f"Momentum down: {momentum_down} (close: {bar.close:.2f} vs prev_close: {self.previous_bar.close:.2f})")
            
            if (prev_touched_ema and current_below_ema and momentum_down):
                self.debug(f"SHORT EMA bounce signal triggered!")
                self.enter_short_position(bar.close)
    
    def calculate_simple_stop_loss(self, entry_price, direction):
        """Calculate simple 2% stop loss (matching original strategy)"""
        stop_pct = 0.02  # Simple 2% stop loss
        
        if direction == "long":
            return entry_price * (1 - stop_pct)
        else:
            return entry_price * (1 + stop_pct)
    
    def enter_long_position(self, entry_price):
        """Enter long position with simple 2% risk management and proper orders"""
        # Simple stop loss calculation
        stop_price = self.calculate_simple_stop_loss(entry_price, "long")
        stop_distance = entry_price - stop_price
        
        # Simple position sizing: 2% portfolio risk
        portfolio_value = self.portfolio.total_portfolio_value
        risk_amount = portfolio_value * 0.02  # 2% risk per trade
        
        shares = int(risk_amount / stop_distance)
        
        # Limit position size to available cash
        max_shares = int(self.portfolio.cash / entry_price)
        shares = min(shares, max_shares)
        
        if shares > 0:
            # Enter position with market order
            self.market_order(self.spy, shares)
            
            # Set up stop loss and take profit orders
            self.position_entry_price = entry_price
            self.stop_loss_price = stop_price
            self.take_profit_price = entry_price * (1 + self.take_profit_pct)
            self.waiting_for_bounce = False
            self.entry_date = self.time.date()  # Track entry date for max hold period
            
            # Place stop loss and take profit orders
            self.stop_market_order(self.spy, -shares, self.stop_loss_price)
            self.limit_order(self.spy, -shares, self.take_profit_price)
            
            # Increment daily trade counter
            self.daily_trades += 1
            
            self.debug(f"Entered LONG position: {shares} shares at ${entry_price:.2f}")
            self.debug(f"Stop Loss Order: ${self.stop_loss_price:.2f}, Take Profit Order: ${self.take_profit_price:.2f}")
            self.debug(f"Daily trades: {self.daily_trades}/{self.max_daily_trades}")
    
    def enter_short_position(self, entry_price):
        """Enter short position with simple 2% risk management and proper orders"""
        # Simple stop loss calculation
        stop_price = self.calculate_simple_stop_loss(entry_price, "short")
        stop_distance = stop_price - entry_price
        
        # Simple position sizing: 2% portfolio risk
        portfolio_value = self.portfolio.total_portfolio_value
        risk_amount = portfolio_value * 0.02  # 2% risk per trade
        
        shares = int(risk_amount / stop_distance)
        
        # Limit position size to available cash
        max_shares = int(self.portfolio.cash / entry_price)
        shares = min(shares, max_shares)
        
        if shares > 0:
            # Enter position with market order (negative shares for short)
            self.market_order(self.spy, -shares)
            
            # Set up stop loss and take profit orders
            self.position_entry_price = entry_price
            self.stop_loss_price = stop_price
            self.take_profit_price = entry_price * (1 - self.take_profit_pct)
            self.waiting_for_bounce = False
            self.entry_date = self.time.date()  # Track entry date for max hold period
            
            # Place stop loss and take profit orders (positive shares to cover short)
            self.stop_market_order(self.spy, shares, self.stop_loss_price)
            self.limit_order(self.spy, shares, self.take_profit_price)
            
            # Increment daily trade counter
            self.daily_trades += 1
            
            self.debug(f"Entered SHORT position: {shares} shares at ${entry_price:.2f}")
            self.debug(f"Stop Loss Order: ${self.stop_loss_price:.2f}, Take Profit Order: ${self.take_profit_price:.2f}")
            self.debug(f"Daily trades: {self.daily_trades}/{self.max_daily_trades}")
    
    def manage_position(self, bar):
        """Position management - orders should handle exits automatically"""
        # With stop loss and take profit orders in place, we don't need manual management
        # This function is kept for any additional monitoring/logging
        if self.portfolio.invested:
            current_price = bar.close
            position = self.portfolio[self.spy]
            
            # Optional: Log position status - only if we have stop loss price set
            if self.stop_loss_price is not None:
                if position.is_long and current_price <= self.stop_loss_price * 1.01:
                    self.debug(f"LONG position near stop loss: current ${current_price:.2f}, stop ${self.stop_loss_price:.2f}")
                elif position.is_short and current_price >= self.stop_loss_price * 0.99:
                    self.debug(f"SHORT position near stop loss: current ${current_price:.2f}, stop ${self.stop_loss_price:.2f}")
            else:
                # Position exists but we don't have stop loss info (likely overnight hold)
                if self.position_entry_price is not None:
                    if position.is_long:
                        pnl_pct = (current_price - self.position_entry_price) / self.position_entry_price
                    else:
                        pnl_pct = (self.position_entry_price - current_price) / self.position_entry_price
                    self.debug(f"Managing overnight position: current ${current_price:.2f}, P&L: {pnl_pct*100:.2f}%")
    
    def end_of_day_exit(self):
        """Smart overnight decision - exit if losing OR if held too long"""
        if self.portfolio.invested and self.position_entry_price is not None:
            position = self.portfolio[self.spy]
            current_price = self.securities[self.spy].price
            
            # Calculate P&L
            if position.is_long:
                pnl_pct = (current_price - self.position_entry_price) / self.position_entry_price
            else:
                pnl_pct = (self.position_entry_price - current_price) / self.position_entry_price
            
            # Calculate days held
            days_held = (self.time.date() - self.entry_date).days if self.entry_date else 0
            
            # Exit if losing more than 0.5% OR if held more than 3 days (prevent buy-and-hold drift)
            should_exit = pnl_pct < -0.005 or days_held > 3
            
            if should_exit:
                self.liquidate(self.spy)
                self.debug(f"End of day exit - P&L: {pnl_pct*100:.2f}%, Days held: {days_held}, Reason: {'Loss' if pnl_pct < -0.005 else 'Max hold period'}")
                # Reset tracking variables when we exit
                self.position_entry_price = None
                self.stop_loss_price = None
                self.take_profit_price = None
                self.entry_date = None
            else:
                self.debug(f"Holding profitable position overnight - P&L: {pnl_pct*100:.2f}%, Days held: {days_held}")
        
        # Store the closing price for next day's gap calculation
        current_price = self.securities[self.spy].price
        if current_price > 0:
            self.previous_close = current_price
            self.debug(f"Stored closing price: {self.previous_close:.2f}")
    
    def on_order_event(self, order_event):
        """Handle order events"""
        if order_event.status == OrderStatus.FILLED:
            self.debug(f"Order filled: {order_event.fill_quantity} shares at ${order_event.fill_price:.2f}")
    
    def on_end_of_day(self):
        """Only reset tracking variables if no position is held"""
        # Don't reset if we're holding a position overnight
        if not self.portfolio.invested:
            self.position_entry_price = None
            self.stop_loss_price = None
            self.take_profit_price = None
            self.entry_date = None
            self.debug("Reset position tracking - no position held")