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, 1, 5)
        self.SetCash(100000)  # Set Strategy Cash
        
        self.SetBenchmark("SPY")
        
        self.stock = self.AddEquity("SPY", Resolution.Hour)
        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 = .05
        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.Portfolio.Cash * (1-self.optionPercent) /100 // self.stock.Price >= 1):
     #       number = (self.Portfolio.Cash * (1-self.optionPercent) / self.stock.Price // 100)
     #       self.protectionNumber += number
     #       self.MarketOrder(self.stockSymbol, number)
            
        while (self.putContract is  None):
            self.Debug("cycling puts")
            self.putContract = self.GetPutContract()
            

        while (self.callContract is None):
            self.Debug("cycling calls")
            self.callContract = self.GetCallContract()
     
        putBid = self.Securities[self.putContract.Symbol].BidPrice
        putAsk = self.Securities[self.putContract.Symbol].AskPrice
        callBid = self.Securities[self.callContract.Symbol].BidPrice
        callAsk = self.Securities[self.callContract.Symbol].AskPrice
            
        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.Debug(self.callTicket)
            self.callPrice = self.callTicket.AverageFillPrice
            
        if (callBid > self.callPrice * 2):
            self.Liquidate(self.callContract.Symbol)
            self.callContract = None

            
        if (putBid > self.maxPutPrice):
            self.maxPutPrice = putBid
            
        if (self.maxPutPrice * .7 >= putAsk):
            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 OnSecuritiesChanged(self, changes):
        self.Debug(f'{changes}')
    
    def GetPutContract(self):
        
        max_price = self.Portfolio.TotalPortfolioValue * self.optionPercent/ 100 / (self.protectionNumber * self.putsPerHundred)
        
        contracts = self.OptionChainProvider.GetOptionContractList(self.stock.Symbol, self.Time)
        contracts = [x for x in contracts if 330 <= (x.ID.Date - self.Time).days <= 400]
        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]
        x = 1
        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:
            if i.ID.StrikePrice >= max_price:
                call = self.AddOptionContract(i)
                return call
        
        return None    
        
        
    def LogData(self):
        if self.putContract is not None and self.callContract is  None:
            self.Debug("time: {}, stock price: {}, ask price: {}, bid price: {}, max price: {}, protection number: {}, cash, unsettled: {}, {}, call ask: , call bid".format(
                self.Time, self.stock.Price, self.Securities[self.putContract.Symbol].AskPrice, self.Securities[self.putContract.Symbol].BidPrice, 
                self.maxPutPrice, self.protectionNumber, self.Portfolio.Cash, self.Portfolio.UnsettledCash)) #self.Securities[self.callContract.Symbol].AskPrice,
                #self.Securities[self.callContract.Symbol].BidPrice))
                
    def OnSecuritiesChanged(self, changes):
        self.Debug(f'{changes}')