from clr import AddReference
AddReference("System")
AddReference("QuantConnect.Algorithm")
AddReference("QuantConnect.Common")

from System import *
from QuantConnect import *
from QuantConnect.Algorithm import *
from QuantConnect.Securities.Option import OptionStrategies
from datetime import datetime, timedelta

class Option(QCAlgorithm):

    def Initialize(self):
        
        self.SetCash(10000)

        # Start and end dates for the backtest.
        self.SetStartDate(2017,7,1)
        self.SetEndDate(2018,7,1)

        # Add assets
        option = self.AddOption("SPY", Resolution.Minute)
        self.option_symbol = option.Symbol
        self.AddEquity("SPY", Resolution.Minute)
        
        # set our strike/expiry filter for this option chain
        #option.SetFilter(-2, +2, timedelta(2), timedelta(30))
        #option.SetFilter(lambda universe: universe.IncludeWeeklys().Strikes(-2, 2).Expiration(timedelta(0), timedelta(10)))
        option.SetFilter(-2, +2, timedelta(0), timedelta(60))

        # use the underlying equity as the benchmark
        self.SetBenchmark("SPY")
        
        self.Schedule.On(self.DateRules.EveryDay("SPY"), \
            self.TimeRules.AfterMarketOpen("SPY", 10), \
            Action(self.MarketClose))
            
        self.Schedule.On(self.DateRules.EveryDay("SPY"), \
            self.TimeRules.BeforeMarketClose("SPY", 10), \
            Action(self.MarketOpenCall))
            
        self.contract = None

    def OnData(self,slice): #This guy will be automatically called every minute
        if self.Portfolio.Invested: #Hope to speed up by skipping minutes. But still very slow
            return
        self.option_data = slice #Pass option data every minute
        
    def TradeOptions(self, callorput, amount, farorclose): 
        #callorput = 0 means Call; callorput = 1 means Put
        #amount > 0 means long; amount < 0 means short
        #farorclose = True means furthest expiration; farorclose = False means closest expiration
        
        for i in self.option_data.OptionChains:
            chain = i.Value
        
            option = [x for x in chain if x.Right == callorput] #Extract calls or puts
        
            contracts = sorted(sorted(option, \
                key = lambda x: abs(chain.Underlying.Price - x.Strike)), \
                key = lambda x: x.Expiry, reverse=farorclose) #Choose ATM options with the closest

            if len(contracts) == 0: 
                self.Log("No contracts found") 
                continue
    
            self.contract = contracts[0]
            self.MarketOrder(self.contract.Symbol, amount)
            
            return
        
    def MarketOpenCall(self):
        self.TradeOptions(0,1,True)
        
    def MarketOpenPut(self):
        self.TradeOptions(1,-1,False)

    def MarketClose(self):
        #if self.contract is not None and self.Portfolio[self.contract].Invested:
        #    self.Sell(self.contract, 1)

        self.Liquidate()
        try:
            self.stoplossorder.Cancel() #Could add a stop loss order somewhere
        except:
            self.Log("No stoploss order")

    def OnOrderEvent(self, orderEvent):
        self.Log(str(orderEvent))