from datetime import datetime, timedelta
from decimal import Decimal
from QuantConnect.Securities.Option import OptionPriceModels

class BasicTemplateOptionsAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2015, 1, 1)
        self.SetEndDate(2015, 3, 31)
        self.SetCash(1000000)
        equity = self.AddEquity("SPY", Resolution.Minute)
        option = self.AddOption("SPY", Resolution.Minute)
        self.symbol = option.Symbol
        self.SetWarmup(TimeSpan.FromDays(5))
        
        option.PriceModel = OptionPriceModels.CrankNicolsonFD()

        # set our strike/expiry filter for this option chain
        option.SetFilter(-3, +3, timedelta(28), timedelta(39))
        # use the underlying equity as the benchmark
        self.SetBenchmark(equity.Symbol)
        self.put = "SPY" # Initialize the put contract
        self.call = "SPY" # Initialize the call contract
        self.expiration = datetime.now() + timedelta(28)
        
    def OnData(self,slice):
        if not self.Portfolio[self.put].Invested and not self.Portfolio[self.call].Invested and self.Time.hour != 0 and self.Time.minute == 02:
#            self.Log("Reached OnData 1")
            self.TradeOptions(slice) # sell the options
        if ((slice.Time + timedelta(2)) > self.expiration) and self.Time.hour != 0 and self.Time.minute == 01 and self.Portfolio[self.put].Invested and self.Portfolio[self.call].Invested:
            self.Buy(self.put, self.quantity_put) # buy back the put options
            self.Buy(self.call, self.quantity_call) # buy back the call options
        else: return
#            self.Log("1 Date: {0}".format(datetime.now() + timedelta(21)))
#            self.Log("2 Date: {0}".format(self.expiration))
#            self.Log("3 Date: {0}".format(((datetime.now() + timedelta(21)) > self.expiration)))
        # if the option contract expires, print out the price and position information      
#        if slice.Delistings.Count > 0:
#            if [x.Key == self.put for x in slice.Delistings]:
#                self.Log("stock SPY quantity: {0}".format(self.Portfolio["SPY"].Quantity))
#                self.Log("{0} quantity: {1}".format(self.put.Value, self.Portfolio[self.put].Quantity))
#                self.Log("The stock price at Expiry S(T): {}".format(self.Securities["SPY"].Price))

    def TradeOptions(self,slice):
        if slice.OptionChains.Count == 0: return   
        for i in slice.OptionChains:
#            self.Log("Reached TradeOptions 1")
            if i.Key != self.symbol: continue
#            self.Log("Reached TradeOptions 1.5")
            chain = i.Value
            underlying_price = Decimal(chain.Underlying.Price)
            put_point = Decimal(-0.5)
            call_point = Decimal(0.5)
            put = [x for x in chain if x.Right == 1] # filter the put options contracts
            # sorted the contracts according to their expiration dates and choose the ATM options
            put_contracts = sorted(sorted(put, key = lambda x: abs(put_point - x.Greeks.Delta)), 
                                            key = lambda x: x.Expiry, reverse=True)
#            self.Log("put contracts: {0}, {1}, {2}, {3}, {4}, {5}".format(put_contracts[0], put_contracts[1], put_contracts[2], put_contracts[3], put_contracts[4], self.Securities["SPY"].Price))
            if len(put_contracts) == 0: return
#            self.Log("Reached TradeOptions 2")
            put_contract = put_contracts[0] 
            put_delta = put_contract.Greeks.Delta
            put_iv = put_contract.ImpliedVolatility
            self.put = put_contract.Symbol
            call = [x for x in chain if x.Right == 0] # filter the call options contracts
            # sorted the contracts according to their expiration dates and choose the ATM options
            call_contracts = sorted(sorted(call, key = lambda x: abs(call_point - x.Greeks.Delta)), 
                                            key = lambda x: x.Expiry, reverse=True)
#            self.Log("call contracts: {0}, {1}, {2}, {3}, {4}, {5}".format(call_contracts[0], call_contracts[1], call_contracts[2], call_contracts[3], call_contracts[4], self.Securities["SPY"].Price))
            if len(call_contracts) == 0: return  
#            self.Log("Reached TradeOptions 3")
            call_contract = call_contracts[0]
            call_delta = call_contract.Greeks.Delta
            call_iv = call_contract.ImpliedVolatility
            self.call = call_contract.Symbol
            self.quantity_put = int(self.Portfolio.TotalPortfolioValue / put_contract.Strike / 100)
            self.quantity_call = int(self.Portfolio.TotalPortfolioValue / call_contract.Strike / 100)
            self.Sell(self.put, self.quantity_put) # short the put options
            self.Sell(self.call, self.quantity_call) # short the call options
#            self.Log("Correct Strike vs Current Strike: {0}, {1}".format(self.Securities["SPY"].Price, call_contract.Strike))
            self.Log("Put/Call Delta: {0}, {1}".format(put_delta, call_delta))
            self.Log("Put/Call IV: {0}, {1}".format(put_iv, call_iv))
            self.expiration = put_contract.Expiry
#            self.Log("Reached TradeOptions 4")
  #          if self.Portfolio["SPY"].Quantity == 0:
  #              self.Buy("SPY",1)     # buy 100 the underlying stock
  #              self.Log("The stock price at time 0 S(0): {}".format(self.Securities["SPY"].Price))