from QuantConnect.Data.Custom.CBOE import CBOE
from math import sqrt,floor

class ModifiedCollar(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 1, 2)  # Set Start Date
        self.SetEndDate(2020, 4, 1)
        self.SetCash(100000)  # Set Strategy Cash
        
        self.SetBenchmark("SPY")
        
        self.stock = self.AddEquity("SPY", Resolution.Minute)
        self.stockSymbol = self.stock.Symbol
        self.stock.SetDataNormalizationMode(DataNormalizationMode.Raw)
        
        self.vixSymbol = self.AddData(CBOE, "VIX").Symbol
        self.vix = 0
        self.simpleMoving = SimpleMovingAverage(60)
        self.vixSMA = self.SMA(self.vixSymbol, 60, Resolution.Daily)
        self.simpleMoving.Updated += self.OnVixSMA
        
        self.SetSecurityInitializer(lambda x: x.SetMarketPrice(self.GetLastKnownPrice(x)))
        self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.AfterMarketOpen("SPY", 10), self.LogData)
        self.length = None
         
        self.callContract = self.putContract = None
        
        self.maxPutPrice = self.callPrice = 0
        
        self.callTicket = self.putTicket = None
        
        self.optionPercent = .04
        self.optionPercentOfStock = self.optionPercent / (1 - self.optionPercent)
        
        self.putsPerHundred, self.protectionNumber = 2, 1
        
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)
        
        self.SetWarmUp(timedelta(days=60))
        
    def OnVixSMA(self, sender, updated):
        return
        
    def OnData(self, data):
        '''OnData event is the primary entry point for your algorithm. Each new data point will be pumped in here.
            Argumen
                data: Slice object keyed by symbol containing the stock data
        '''
        self.Log("hello")
        for volatility in data.Get(CBOE).Values:
            self.vix = volatility.Value
            self.simpleMoving.Update(self.Time, self.vix)
            
        if self.IsWarmingUp:
            return
    
        if (not self.Portfolio[self.stockSymbol].Invested):
            number = self.CalculateOrderQuantity(self.stockSymbol, 1-self.optionPercent)
            self.protectionNumber = number // 100 
            number = self.protectionNumber * 100
            self.ticket = self.MarketOrder(self.stockSymbol, number)
            
        if (self.putContract is  None):
            self.putContract = self.GetPutContract()

     #   if (self.callContract is None):
     #       self.callContract = self.GetCallContract()
            
        if (self.putContract is not None and not self.Portfolio[self.putContract.Symbol].Invested):
            self.putTicket = self.MarketOrder(self.putContract.Symbol, self.protectionNumber * self.putsPerHundred)
            self.Debug("bought")
            self.maxPutPrice = self.putTicket.AverageFillPrice
           # self.Debug("bought price: {}".format(self.maxPutPrice))
        
     #   if (self.callContract is not None and not self.Portfolio[self.callContract.Symbol].Invested):
     #       self.callTicket = self.MarketOrder(self.callContract.Symbol, -self.protectionNumber)
     #       self.callPrice = self.callTicket.AverageFillPrice
            
     #   if (self.callContract.AskPrice * 2 > self.callPrice):
     #       self.Liquidate(self.callContract.Symbol)
     #      self.callContract = None
            
        if (self.putContract.BidPrice > self.maxPutPrice):
            self.maxPutPrice = self.putContract.BidPrice
            
        if (self.maxPutPrice * .8 >= self.putContract.AskPrice):
            self.putTicket = self.MarketOrder(self.putContract.Symbol, -self.protectionNumber * self.putsPerHundred)
            self.Debug("sold")
           # self.Debug("sold price: {}".format(self.putTicket.AverageFillPrice))
            self.RemoveSecurity(self.putContract.Symbol)
            self.maxPutPrice = 0
            self.putContract = None
            
        
    def GetPutContract(self):
        
        max_price = self.Portfolio.TotalPortfolioValue * self.optionPercent/ 100 / (self.protectionNumber * self.putsPerHundred)
        
        contracts = self.OptionChainProvider.GetOptionContractList(self.stock.Symbol, self.Time)
        x = 1
        contracts = [x for x in contracts if 330 <= (x.ID.Date - self.Time).days <= 400]
        x = 2
        contracts = [x for x in contracts if x.ID.StrikePrice < self.stock.Price]
        contracts = [x for x in contracts if x.ID.OptionRight == OptionRight.Put]
        contracts = sorted(sorted(contracts, key = lambda x: x.ID.StrikePrice, reverse = True), key = lambda x: abs(x.ID.Date - self.Time))
        
        for i in contracts:
            contract = self.AddOptionContract(i)
            if contract.AskPrice < max_price:
                return contract
            else:
                self.RemoveSecurity(i)
        
        return None
        
    def GetCallContract(self):
        
      #  if (self.vix >= self.vixSMA.Current.Value * 1.5):
      #      return None
        
        max_price = floor((1 + (self.vix/100 / sqrt(52) )) * self.Securities[self.stockSymbol].Price) 
        
        contracts = self.OptionChainProvider.GetOptionContractList(self.stock.Symbol, self.Time)
        contracts = [x for x in contracts if 5 <= (x.ID.Date - self.Time).days <= 7]
        contracts = [x for x in contracts if x.ID.StrikePrice > self.stock.Price]
        contracts = [x for x in contracts if x.ID.OptionRight == OptionRight.Call]
        contracts = sorted(sorted(contracts, key = lambda x: x.ID.StrikePrice, reverse = False), key = lambda x: x.ID.Date)
        
        for i in contracts:
            call = self.AddOptionContract(i)
            if i.ID.StrikePrice >= max_price:
                call = self.AddOptionContract(i)
                return call
        
        return None    
        
        
    def LogData(self):
        if self.putContract is not None:
            self.Debug("time: {}, stock price: {}, ask price: {}, bid price: {}, max price: {}, percentage different: {}".format(self.Time, self.stock.Price, self.putContract.AskPrice, self.putContract.BidPrice, self.maxPutPrice, (self.maxPutPrice - self.putContract.AskPrice)/self.maxPutPrice))