from QuantConnect.Securities.Option import OptionPriceModels

class CallTailendSetDates(QCAlgorithm):

    #designed so that in backtesting, if there isn't data or it can't find an option, it should bug out. Final product will be human supervised
    def Initialize(self):
        self.SetStartDate(2014, 1, 1)
        self.SetCash(100000)  # Set Strategy Cash
        stock = self.AddEquity("SPY", Resolution.Minute)
        stock.SetDataNormalizationMode(DataNormalizationMode.Raw)
        self.stock_symbol = stock.Symbol
        
        bond = self.AddEquity("TLT", Resolution.Minute)
        bond.SetDataNormalizationMode(DataNormalizationMode.Raw)
        self.bond_symbol = bond.Symbol
        
        option = self.AddOption("SPY", Resolution.Minute)
        option.SetFilter(lambda universe: universe.Strikes(-30, 30).Expiration(timedelta(366), timedelta(400)))
        option.PriceModel = OptionPriceModels.CrankNicolsonFD()
        self.option_symbol = option.Symbol
        
        self.purchaseDates = [None, None]
        self.callOptions = [None, None]
        self.putOption = None
        
        self.callWeight = .15
        self.putWeight = 0.02
        self.bondWeight = 1 - self.callWeight - self.putWeight
        
        self.month = 0
        self.shouldCheckIndex = -1
        self.shouldRebalance = True
        
        self.SetWarmUp(timedelta(1))

    def OnData(self, data):
        
        if self.IsWarmingUp or self.Time.hour < 9 and self.Time.minute < 30: return
        
        if self.Time.month != self.month:
            self.month = self.Time.month
            if self.month == 6:
                self.shouldCheckIndex = 0
            elif self.month == 12:
                self.shouldCheckIndex = 1
                
        if self.shouldCheckIndex >= 0:
            self.ShouldRebalance(self.shouldCheckIndex)
            if self.shouldRebalance:
                self.Rebalance(self.shouldCheckIndex)
                self.shouldCheckIndex = -1

            
    def GetStockOptions(self, optionType, portfolioPercent):
        optionchain = None
        for kvp in self.CurrentSlice.OptionChains:
            if kvp.Key != self.option_symbol: continue
            optionchain = kvp.Value
        maxValue = self.Portfolio.TotalPortfolioValue * portfolioPercent * (-1 if optionType == OptionRight.Call else 1)
        optionchain = [x for x in optionchain if x.Right == optionType]
        x = 1
        optionchain = sorted(optionchain, key = lambda x: (abs((x.Expiry - self.Time).days - 365), maxValue // (x.AskPrice * 100) * x.Greeks.Delta))
        return optionchain[0]
    
    def ShouldRebalance(self, index):
        self.shouldRebalance = not self.callOptions[index] or self.Portfolio[
            self.callOptions[index].Symbol].UnrealizedProfit < 0 or (self.Time - self.purchaseDates[index]).days > 0
        
        
    def Rebalance(self, index):
        self.Debug(f"Rebalance {self.Time}")
        if self.callOptions[index]:
            self.Liquidate(self.callOptions[index].Symbol)
        if self.putOption:
            self.Liquidate(self.putOption.Symbol)
            
        self.callOptions[index] = self.GetStockOptions(OptionRight.Call, self.callWeight/2)
        self.purchaseDates[index] = self.Time
        self.putOption = self.GetStockOptions(OptionRight.Put, self.putWeight)
            
        if self.callOptions[index - 1]:
            self.SetHoldings(self.callOptions[index - 1].Symbol, self.callWeight / 2)
        self.SetHoldings(self.bond_symbol, self.bondWeight)
        self.SetHoldings(self.putOption.Symbol, self.putWeight)
        self.SetHoldings(self.callOptions[index].Symbol, self.callWeight / 2)
        x = 1
        
    def OnEndOfDay(self):
        for each in self.Portfolio:
            if each.Value.Invested and each.Value.Type == SecurityType.Option:
                self.Debug(f"{each.Key}: {self.Portfolio[each.Key].AbsoluteHoldingsValue}")