import numpy as np
import pandas as pd
import datetime

from System import *
from QuantConnect import *
from QuantConnect.Algorithm import *
from QuantConnect.Data import *

class bugs(QCAlgorithm):

    def Initialize(self):

        self.SetStartDate(2018, 1, 1)   #Set Start Date
        self.SetEndDate(2018, 2, 31)   #Set End Date
        self.SetCash(25000)             #Set Strategy Cash
        
        # Find more symbols here: http://quantconnect.com/data
        equity = self.AddEquity("SPY", Resolution.Minute)
        option = self.AddOption("SPY", Resolution.Minute)
        self.symbol = option.Symbol
        self.ATMcall = None
        self.ATMput = None
        self.Allocate = 0.10
        
        # Scheduled event on Mondays at 3:45PM
        self.Schedule.On(self.DateRules.Every([DayOfWeek.Monday]), \
            self.TimeRules.At(3, 45), \
            self.SellTime)
                 
        option.SetFilter(lambda universe: universe.IncludeWeeklys().Strikes(0, 0).Expiration(datetime.timedelta(4), datetime.timedelta(8)))

    def OnData(self, slice):
        
        if not self.Portfolio.Invested:
            
            if slice != None:
                for i in slice.OptionChains:
                    if i.Key != self.symbol: continue
                    optionchain = i.Value
            
                    self.Log("underlying price:" + str(optionchain.Underlying.Price))
                    
                    calls = [x for x in optionchain if x.Right == 0]
                    puts = [x for x in optionchain if x.Right == 1]
                    
                    # get closest ATM call
                    self.ATMcall = sorted(calls, key = lambda x: x.Strike, reverse = True)[0]
                    self.Log(str(self.ATMcall))
                    
                    # get closest ATM put
                    self.ATMput = sorted(puts, key = lambda x: x.Strike, reverse = True)[0]
                    self.Log(str(self.ATMput))
                    
                    self.Liquidate()
                    self.MarketOrder(self.ATMput.Symbol, 5)
                    self.MarketOrder(self.ATMcall.Symbol, 5)
                
    def SellTime(self):
        self.Liquidate()
        self.Log("Sold at : {0}".format(self.Time))