# region imports
from AlgorithmImports import *
from datetime import datetime
from datetime import timedelta


# endregion


class FatBlueAlligator(QCAlgorithm):

    def Initialize(self):

        self.ticker = "AAPL"
        self.SetStartDate(2022, 1, 1)
        self.MIN_DAYS_TO_EXP = 3  # don't consider contracts less than X days to expiry
        self.CALENDAR_CAL_FAR_EXPIRY = 21  # at least DAYS away from near term for calendar 
        self.STRATEGY_TYPE = 3 # 1 - stock; 2 - straddle;  3 - calendar ;

        datesToWatchStr = '27/01/2022,28/04/2022,28/07/2022,27/10/2022,02/02/2023'

        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)
        self.SetWarmUp(30, Resolution.Daily)
        self.SetCash(1000000)
        self.AddEquity(self.ticker, Resolution.Minute)
        self.option = self.AddOption(self.ticker, Resolution.Minute)
        self.symbol = self.option.Symbol

        self.option.SetFilter( lambda option_filter_universe: option_filter_universe.IncludeWeeklys().Strikes(-1, 1).Expiration(self.MIN_DAYS_TO_EXP, 270)) 
       

        self.datesToWatch = []
        for i in datesToWatchStr.split(","):
            self.datesToWatch.append(datetime.strptime(i, '%d/%m/%Y').date())

        self.dateSell = None
        self.invested = False

    def OnData(self, slice: Slice) -> None:

        if not self.IsWarmingUp and self.Time.hour == 10 and self.Time.minute == 00 and self.dateSell == None:
            
            if self.Time.date()+timedelta(1) in self.datesToWatch:

                self.Log(f"signal is triggered on\t {self.ticker}")

                chain = slice.OptionChains.get(self.symbol, None)

                if not chain:
                    self.Log("options chain is not ready")
                    return

                contract_days = [] # collecting contracts dates list for custom far term selection
                for i in chain:
                    if i.Expiry.date() not in contract_days:
                        contract_days.append(i.Expiry.date())

                contract_days = sorted(contract_days)
                self.Log(f"Contracts are with following expiry dates\t{contract_days}")

                expiry = sorted(chain, key=lambda contract: contract.Expiry, reverse=False)[0].Expiry
                underlyingLastPrice = sorted(chain, key=lambda contract: contract.Expiry, reverse=False)[0].UnderlyingLastPrice # perhaps there is a better way to get underlyng price

                far_expiry=''
                for i in contract_days:
                    if (expiry.date() + timedelta(self.CALENDAR_CAL_FAR_EXPIRY)) <= i:
                        far_expiry = i
                        break

                strikes = [contract.Strike for contract in chain if contract.Expiry == expiry]
                strike = sorted(strikes, key=lambda strike: abs(strike - chain.Underlying.Price))[0]
                
                self.Log(f"{self.symbol}\tstrategy is based on\t\t Stock Price:{underlyingLastPrice}\tstrike:{strike}\tnear term:{str(expiry.date())}\tfar term:{far_expiry}")

                if self.STRATEGY_TYPE == 1: # cash
                    self.Buy(self.ticker, 100)

                if self.STRATEGY_TYPE == 2: # straddle
                    staddle_option_strategy = OptionStrategies.Straddle(self.symbol, strike, expiry)
                    self.Buy(staddle_option_strategy, 1)

                if self.STRATEGY_TYPE == 3: # calendar
                    calendar_option_strategy = OptionStrategies.CallCalendarSpread(self.symbol, strike, expiry, far_expiry)
                    self.Buy(calendar_option_strategy, 1)
                    self.Log("Call Spread 'Buy' has been completed")


                self.dateSell = self.Time.date() + timedelta(1) # sell the next day after earnings announced
                self.invested = True 

        if self.invested and self.Time.date() > self.dateSell:
            self.Liquidate()
            self.invested = False
            self.dateSell = None
            self.Log(f"position has been liquidated")