from AlgorithmImports import *

class MeanReversionShortAlgorithm(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2022, 1, 1)
        self.SetEndDate(2023, 1, 1)
        self.SetCash(100000)
        
        # Define a list of 10 well-known, highly liquid stocks
        self.tickers = [
            'AAPL', 'MSFT', 'GOOGL', 'AMZN', 'NVDA', 'META', 'TSLA', 'JNJ', 'JPM', 'V'
        ]

        # Add Equity data for the tickers
        self.symbols = []
        for ticker in self.tickers:
            self.symbols.append(self.AddEquity(ticker, Resolution.Daily).Symbol)
        
        # Add SPY for scheduling and market checks
        self.spy = self.AddEquity("SPY", Resolution.Daily).Symbol
        
        # Dictionary to store previous close prices
        self.previous_close = {}
        
        # Warm-up the algorithm to get historical data for previous close prices
        self.SetWarmup(1)
        
        # Schedule event to sell all positions at 3:55 PM
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.At(15, 55), self.Rebalance)
        
    def OnData(self, data):
        if self.IsWarmingUp:
            return
        
        # Avoid trading near market close
        if self.Time.hour == 15 and self.Time.minute >= 55:
            return
        
        # Collect previous close prices
        for symbol in self.symbols:
            if symbol in data.Bars:
                self.previous_close[symbol] = data.Bars[symbol].Close
        
        # Calculate overnight percentage change
        overnight_changes = {}
        for symbol in self.symbols:
            if symbol in data.Bars and symbol in self.previous_close:
                open_price = data.Bars[symbol].Open
                prev_close_price = self.previous_close[symbol]
                overnight_change_pct = (open_price - prev_close_price) / prev_close_price * 100
                overnight_changes[symbol] = overnight_change_pct
        
        # Find stocks with greater than 2% but less than 3% overnight upmove
        top_movers = [symbol for symbol, change in overnight_changes.items() if 2 < change < 3]
        
        # Open short positions for the stocks with greater than 2% but less than 3% upmove
        for symbol in top_movers:
            self.SetHoldings(symbol, -1.0 / len(top_movers))  # Short position
        
        # Reset previous close prices
        self.previous_close = {}
    
    def Rebalance(self):
        if self.IsWarmingUp:
            return
        
        # Ensure market is open before liquidating positions
        if not self.Securities[self.spy].Exchange.DateTimeIsOpen(self.Time):
            return
        
        # Liquidate all positions before market close
        for kvp in self.Portfolio:
            if kvp.Value.Invested:
                self.Liquidate(kvp.Key)
 
        # Log for debugging purposes
        self.Debug(f"Positions liquidated at {self.Time}")

# Ensure the code runs correctly
from AlgorithmImports import *