# region imports
from AlgorithmImports import *
# endregion
from QuantConnect.Algorithm import QCAlgorithm
from QuantConnect.Data.Market import TradeBar
from QuantConnect.Orders import OrderTicket
from datetime import timedelta

class GapUpBreakout(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 12, 1)
        self.SetEndDate(2020, 12, 10)
        self.SetCash(100000)

        # Define the universes for all US equities
        self.AddUniverse(self.CoarseSelectionFunction)

        self.symbol_data = {}
        self.previous_second = -1 # Initialize previous second

    def CoarseSelectionFunction(self, coarse):
        # Select all US equities that have fundamental data
        return [x.Symbol for x in coarse if x.HasFundamentalData and x.Market in [
            "NYS", "ASE", "ARCA", "NCM", "NAS", "NGM", "NMS"
        ]]

    def OnData(self, data):
        for symbol, tradebar in data.TradeBars.items():
            if symbol not in self.symbol_data:
                continue

            sd = self.symbol_data[symbol]

            if not sd.bought_today:
                continue

            current_price = tradebar.Close

            # Sell conditions
            if sd.entry_price is not None:
                take_profit_level = sd.open_price * 1.60
                stop_loss_level = sd.open_price * 1.10

                if current_price >= take_profit_level:
                    self.Log(f"Selling {symbol} at take profit: {current_price}")
                    self.Liquidate(symbol)
                    sd.bought_today = False
                    sd.entry_price = None
                    continue
                elif current_price <= stop_loss_level:
                    self.Log(f"Selling {symbol} at stop loss: {current_price}")
                    self.Liquidate(symbol)
                    sd.bought_today = False
                    sd.entry_price = None
                    continue

        # Check for buy opportunities every second
        if self.Time.second != self.previous_second:
            self.previous_second = self.Time.second
            for symbol in self.ActiveSecurities.Keys:
                if symbol not in data.Open:
                    continue
                if symbol not in data.TradeBars:
                    continue

                open_price = data.Open[symbol]
                current_price = data.TradeBars[symbol].Close

                if symbol not in self.symbol_data:
                    previous_day = self.History(symbol, 2, Resolution.Daily)
                    if previous_day.shape[0] < 2:
                        continue
                    previous_close = previous_day['close'].iloc[-2]

                    if open_price >= 2.75 and open_price >= previous_close * 1.10:
                        self.symbol_data[symbol] = SymbolData(open_price)
                else:
                    sd = self.symbol_data[symbol]
                    if not sd.bought_today and sd.open_price is not None:
                        breakout_level = sd.open_price * 1.20
                        if current_price > breakout_level:
                            quantity = self.CalculateOrderQuantity(symbol, 0.01) # Allocate 1% of portfolio
                            self.Log(f"Buying {symbol} at breakout: {current_price}")
                            order_ticket = self.Buy(symbol, quantity)
                            sd.bought_today = True
                            sd.entry_price = current_price

    def OnEndOfDay(self):
        for symbol, sd in self.symbol_data.items():
            if sd.bought_today:
                self.Log(f"Selling {symbol} at end of day.")
                self.Liquidate(symbol)
                sd.bought_today = False
                sd.entry_price = None

class SymbolData:
    def __init__(self, open_price):
        self.open_price = open_price
        self.bought_today = False
        self.entry_price = None

    def __repr__(self):
        return f"Open: {self.open_price}, Bought: {self.bought_today}, Entry: {self.entry_price}"