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, 3, 10)    #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
        self.tradesOn = True
        # Scheduled event on Mondays at 3:45PM
        self.Schedule.On(self.DateRules.EveryDay("SPY"),
           self.TimeRules.BeforeMarketClose("SPY", 1),
           self.TurnOn)

        self.Schedule.On(self.DateRules.Every([DayOfWeek.Friday]), \
            self.TimeRules.At(15, 45), \
            self.SellTime)
                 
        option.SetFilter(lambda universe: universe.IncludeWeeklys().Strikes(0, 0).Expiration(datetime.timedelta(15), datetime.timedelta(20)))

    def OnData(self, slice):
        
        ## Add a boolean condition to see if the algo. should continue to trade that day
        if not self.Portfolio.Invested and self.tradesOn:
            
            if slice != None:
                for i in slice.OptionChains:
                    if i.Key != self.symbol: continue
                    optionchain = i.Value
            
                    self.Log(f"underlying price: {optionchain.Underlying.Price}")
                    
                    puts = [x for x in optionchain if x.Right == 1]

                    # get closest ATM put
                    self.ATMput = sorted(puts, key = lambda x: x.Strike, reverse = True)[0]
                    self.Log(f'{self.Time} :: Put expiry: {self.ATMput.Expiry}')
                    self.Log(f'{self.Time} :: Put Symbol: {self.ATMput.Symbol}')
                    for symbol in self.Securities.Keys:
                        if self.Securities[symbol].Invested:
                            self.Log(f'{self.Time} :: Liquidating {symbol}')
                    
                    self.Liquidate()
                    self.SetHoldings(self.ATMput.Symbol, self.Allocate)
                    self.Log(f'{self.Time} :: Purchased Put Option with ' + str(abs((self.ATMput.Expiry - self.Time).days)) + ' days to expiration')
       
    def SellTime(self):
        for symbol in self.Securities.Keys:
            if self.Securities[symbol].Invested:
                self.Log(f'{self.Time} :: Liquidating {symbol}')
        self.Liquidate()
        
        ## This means that the OnData() section won't execute any of the code inside the if() statement from 15:46 - Market Close
        self.tradesOn = False

    def TurnOn(self):
        self.tradesOn = True