from AlgorithmImports import *
#endregion

import numpy as np
import pandas as pd

class GapAndGoAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2023, 1, 1)  # Set start date
        self.SetEndDate(2023, 1, 31)  # Set end date
        self.SetCash(100000)  # Set strategy cash

        # Add SPY to the universe and request daily data for it
        self.AddEquity("SPY", Resolution.Daily)
        
        # Use universe selection to select a universe of stocks that are gapping up
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
        self.SelectionCount = 10
        
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)  # Set brokerage model
        
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.AfterMarketOpen("SPY", 30), self.GapAndGo)  # Schedule GapAndGo function to run after market open

    def CoarseSelectionFunction(self, coarse):
        # Select only stocks with a price above $5 and a daily dollar volume above $10,000,000
        selected = [x.Symbol for x in coarse if x.Price > 5 and x.DollarVolume > 10000000]
        return selected
    
    def FineSelectionFunction(self, fine):
        # Select stocks that gapped up at least 1% and have positive earnings and revenue growth
        fine = [x for x in fine if hasattr(x.EarningReports, 'BasicEPSGrowth') and x.EarningReports.BasicEPSGrowth.OneYear > 0 and x.FinancialStatements.RevenueGrowth.OneYear > 0]
        fine = [x for x in fine if (x.Open - x.Close[-2])/x.Close[-2] > 0.01]
        fine = sorted(fine, key=lambda x: (x.Volume, x.Price), reverse=True)[:self.SelectionCount]
        return [x.Symbol for x in fine]

    def GapAndGo(self):
        # Loop over all securities in the portfolio
        for security in self.Portfolio.Values:
            # Check if the security is currently tradable
            if not self.Securities[security.Symbol].IsTradable:
                continue

            # Get daily historical data for the last 2 days
            history = self.History([security.Symbol], 2, Resolution.Daily)

            # Get opening and closing prices for the last 2 days
            opens = history.loc[security.Symbol].open.values
            closes = history.loc[security.Symbol].close.values

            # Check if the open price today is greater than the close price of yesterday
            if opens[-1] > closes[-2]:
                # Calculate the percentage change between the two days
                pct_change = (closes[-1] - closes[-2]) / closes[-2]

                # If the percentage change is greater than 1% (i.e., a gap), execute a trade
                if pct_change > 0.01:
                    # Calculate the stop loss and take profit levels based on the current price
                    current_price = self.Securities[security.Symbol].Price
                    stop_loss = np.array([current_price * 0.98]).item()
                    take_profit = np.array([current_price * 1.02]).item()
                    quantity = 100
                    # Place the market order with stop loss and take profit levels
                    #, False, None, stop_loss, take_profit
                    self.ticket = self.MarketOrder(security.Symbol, quantity)
                    self.LimitOrder(security.Symbol, -quantity, self.ticket.AverageFillPrice* 1.02)
                    self.StopMarketOrder(security.Symbol, -quantity, self.ticket.AverageFillPrice* 0.98)
    
    def OnOrderEvent(self, orderEvent):
        order = self.Transactions.GetOrderById(orderEvent.OrderId)
        if order.Status == OrderStatus.Filled:
            if order.Type == OrderType.Limit or order.Type == OrderType.StopMarket:
                self.Transactions.CancelOpenOrders(order.Symbol)