from AlgorithmImports import *

class SurpriseTradingAlgorithm(QCAlgorithm):

    def Initialize(self):
        # Enter the symbol you want to trade here
        self.symbol = "pep"
        
        # Add the asset you want to trade and subscribe to Estimize data
        self.AddEquity(self.symbol)
        self.estimize_release_symbol = self.AddData(EstimizeRelease, self.symbol).Symbol
        
        # Set start and end dates for historical data
        self.SetStartDate(2010, 1, 1)
        self.SetEndDate(2022, 12, 31)
        
        # Set cash amount for trading
        self.SetCash(100000)
        
        # Initialize a variable to track positions
        self.position_opened = False

    def OnData(self, slice: Slice):
        if self.estimize_release_symbol in slice:
            estimize_data = slice[self.estimize_release_symbol]
        
            # Check if the data is not None before performing operations
            if estimize_data.Eps is not None and estimize_data.ConsensusEpsEstimate is not None:
                # Calculate the market surprise
                market_surprise = estimize_data.Eps - estimize_data.ConsensusEpsEstimate
                
                # Define a threshold for significant surprises (e.g., 0.05)
                threshold = 0.02

                if market_surprise > threshold and not self.position_opened:
                    # If the surprise is positive and no position is open, sell short and hold for 5 days
                    self.SetHoldings(self.symbol, 1.0)
                    self.position_opened = True
                    self.Schedule.On(self.DateRules.EveryDay(self.symbol), self.TimeRules.AfterMarketOpen(self.symbol, 5), self.ClosePosition)

                elif market_surprise < -threshold and not self.position_opened:
                    # If the surprise is negative and no position is open, buy and hold for 5 days
                    self.SetHoldings(self.symbol, -1.0)
                    self.position_opened = True
                    self.Schedule.On(self.DateRules.EveryDay(self.symbol), self.TimeRules.AfterMarketOpen(self.symbol, 5), self.ClosePosition)

    def ClosePosition(self):
        # Close the position after 5 days
        self.Liquidate(self.symbol)
        self.position_opened = False