from AlgorithmImports import *
from datetime import datetime, timedelta
from collections import deque


class FocusedYellowLemur(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2023, 1, 1)  # Set Start Date
        self.SetEndDate(2024, 5, 13)
        self.SetCash(100000)  # Set Strategy Cash
        
        self.ticker = "QQQ"
        self.res = Resolution.Minute
        self.equitySymbol = None
        self.optionSymbol = None
        

        equity = self.AddEquity(self.ticker, self.res)
        equity.SetDataNormalizationMode(DataNormalizationMode.Raw)
        option = self.AddOption(equity.Symbol)
        option.set_filter(0, +5, 30, 45)
        self.equitySymbol = equity.Symbol
        self.optionSymbol = option.Symbol

        self.consolidator = None
        self.consolidators = dict()

        self.lkb = 20
        self.lkbOption = 20

        self.atmCall = None



# Underlying

    # Underlying
        # Create a minutes=30 QuoteBarConsolidator
        consolidator = QuoteBarConsolidator(timedelta(minutes=30))

        # Create the custom indicator
        self.indicator = CustomIndicator("indicator", self.lkb)

        # Register the custom indicator to update with the consolidated data
        self.RegisterIndicator(self.equitySymbol, self.indicator, consolidator)

        # Subscribe to the consolidated data
        self.SubscriptionManager.AddConsolidator(self.equitySymbol, consolidator)


# Option
        self.atm_indicator = CustomIndicator("atm_indicator", self.lkbOption)
        


        self.AddRiskManagement(MaximumUnrealizedProfitPercentPerSecurity(0.70))


    def OnData(self, data):



        if not data.ContainsKey(self.equitySymbol): return


        option_invested = [x.Key for x in self.Portfolio if x.Value.Invested and x.Value.Type==SecurityType.Option]


        if option_invested:
            if self.Time + timedelta(7) > option_invested[0].ID.Date:
                self.Liquidate(option_invested[0], "Too close to expiration")
            return
        

        self.updateIndicators(data, self.optionSymbol, self.indicator, self.atm_indicator)



    def updateIndicators(self, data, OS, underlying_indicator, atm_indicator):

        calls = None

        chain = data.option_chains.get_value(OS)
        if chain is None:
            return


        calls = [i for i in chain if i.Right == OptionRight.Call]
        

        if not calls: return

        call_contracts = sorted(calls, key = lambda x: x.Strike)
        #self.log("Checking")

        
        if len(call_contracts) < 1: return
            #self.Log("through")

        


        self.atmCall = call_contracts[0]
        

        if not underlying_indicator.IsReady or not atm_indicator.IsReady: return

        
        long_condition = self.securities[self.equitySymbol].Close > underlying_indicator.Value and self.securities[self.atmCall.Symbol].ask_price > atm_indicator.Value


        if not self.Portfolio[self.atmCall.Symbol].IsLong:
            if long_condition:
                self.MarketOrder(self.atmCall.Symbol, 6)
                self.Log("Bought atm call")

        
        self.Plot("Underlying", "Underlying Indicator", underlying_indicator.Value)
        self.Plot("Underlying", "Underlying Value", self.securities[self.equitySymbol].Close)

        self.Plot("Option", "Option Indicator", atm_indicator.Value)
        self.Plot("Option", "Option Value", self.securities[self.atmCall.Symbol].ask_price)



    def OnDataConsolidated(self, sender, quotebar):

        if self.atmCall != None:

            # Check if the quotebar is for thje atm call option
            if quotebar.Symbol == self.atmCall.Symbol:
                self.atm_indicator.Update(quotebar)
                #self.Log("OnDataConsolidated called for Call Option at " + str(self.Time))




    def OnSecuritiesChanged(self, changes):
        for security in changes.AddedSecurities:
            if security.Type == SecurityType.Equity:  # Check if the security is an Equity
                #self.Log(f"{security}")
                self.consolidator = QuoteBarConsolidator(timedelta(minutes=30))
                self.consolidator.DataConsolidated += self.OnDataConsolidated
                self.SubscriptionManager.AddConsolidator(security.Symbol, self.consolidator)
                self.consolidators[security.Symbol] = self.consolidator

            else:
                self.consolidator = QuoteBarConsolidator(timedelta(minutes=30))
                self.consolidator.DataConsolidated += self.OnDataConsolidated
                self.SubscriptionManager.AddConsolidator(security.Symbol, self.consolidator)
                self.consolidators[security.Symbol] = self.consolidator


        for security in changes.RemovedSecurities:
            if security.Symbol in self.consolidators:
                self.consolidator = self.consolidators.pop(security.Symbol)
                self.SubscriptionManager.RemoveConsolidator(security.Symbol, self.consolidator)
                self.consolidator.DataConsolidated -= self.OnDataConsolidated


class CustomIndicator(PythonIndicator):
    def __init__(self, name, period):
        super().__init__()
        self.Name = name
        self.Value = 0
        self.period = period
        self.queue = deque(maxlen=period)


    def Update(self, input):
        if not isinstance(input, QuoteBar):
            raise TypeError('CustomIndicator.Update: input must be a QuoteBar')

        self.queue.append(input.Close)
        
        self.Value = sum(self.queue) / self.period

        return self.IsReady

    @property
    def IsReady(self):
        return len(self.queue) == self.period