import math
from datetime import timedelta

from QuantConnect.Securities.Option import OptionPriceModels
from QuantConnect.Orders import OrderStatus

class BullCallSpreadAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2017, 1, 1)
        self.SetEndDate(2017, 2, 1)
        self.SetCash(10000)
        self.equity = self.AddEquity("SPY", Resolution.Minute)
        self.option = self.AddOption("SPY", Resolution.Minute)
        self.symbol = self.option.Symbol
    
        # set our strike/expiry filter for this option chain
        self.option.SetFilter(-20, 1, timedelta(15), timedelta(45))
        self.option.PriceModel = OptionPriceModels.CrankNicolsonFD()
        
        self.SetWarmUp(TimeSpan.FromDays(4))
        
        # use the underlying equity SPY as the benchmark
        self.SetBenchmark(self.equity.Symbol)
        self.Schedule.On(self.DateRules.EveryDay(self.symbol),
                         self.TimeRules.BeforeMarketClose(self.symbol, 30),
                         self.TradeOptions)
                         
        self.risk_factor = 0.5
        
    def OnData(self,slice):
        if self.equity.Symbol in self.Portfolio and self.Portfolio[self.equity.Symbol].Quantity != 0:
            self.Sell(self.equity.Symbol, self.Portfolio[self.equity.Symbol].Quantity)
        
        optionchain = slice.OptionChains
        valid = False
        for i in slice.OptionChains:
            if i.Key != self.symbol:
                continue
            if i.Value.Contracts.Count != 0:
                valid = True
                break
            else:
                self.Log("Chain is empty")
        if not valid:
            self.optionchain = None
            return
        self.optionchain = optionchain
 
    def TradeOptions(self):
        if self.IsWarmingUp:
            return
        if self.Portfolio.Invested:
            return
        if not self.optionchain:
            self.Log("No chain")
            return
        chain = None
        for i in self.optionchain:
            if i.Key != self.symbol: 
                continue
            chain = i.Value
            break
        
        if chain is None:
            self.Log("Chain not found")
            return
        
        if chain.Contracts.Count == 0:
            self.Log("Chain empty")
            return
            
        expiry = sorted(chain, key = lambda x: x.Expiry, reverse=True)[0].Expiry
        self.Log("Expiry is {}".format(expiry))
        puts = [i for i in chain if i.Expiry == expiry and i.Right == OptionRight.Put]
        if not puts:
            self.Log("No contracts for given expiry")
            return
        puts_by_delta = sorted(puts, key=lambda x: x.Greeks.Delta, reverse=True)    
        
        low_put = None
        high_put = None
            
        for contract in puts_by_delta:
            if abs(contract.Greeks.Delta) > 0.40 and low_put is None:
                low_put = contract
            if abs(contract.Greeks.Delta) > 0.45 and high_put is None:
                high_put = contract
                break
            
        if high_put and low_put:
            self.Log("Low Put: {}, High Put: {}".format(low_put.Greeks.Delta, high_put.Greeks.Delta))
        else:
            self.Log("Lowest Put: {}, Highest Put: {}".format(puts_by_delta[0].Greeks.Delta, puts_by_delta[-1].Greeks.Delta))
            return
        
        spread_risk = float(high_put.Strike - low_put.Strike) * 100.0
        if spread_risk == 0:
            self.Log("Strikes are the same.")
            return
        amount_to_invest = float(self.Portfolio.Cash) * self.risk_factor
        num_shares = math.floor(amount_to_invest / float(spread_risk))
        self.Log("Spread risk: {}, amount_to_invest: {}, num_shares: {}".format(spread_risk, amount_to_invest, num_shares))
        if num_shares > 0:
            if self.Sell(high_put.Symbol, num_shares).Status != OrderStatus.Filled:
                self.Log("Couldn't fill sell order")
                return
            if self.Buy(low_put.Symbol, num_shares).Status != OrderStatus.Filled:
                self.Log("Couldn't fill buy order, reversing sell order")
                self.Buy(high_put.Symbol, num_shares)
                return
            
        self.Schedule.On(self.DateRules.On(expiry.year, expiry.month, expiry.day),
                         self.TimeRules.BeforeMarketClose(self.symbol, 15),
                         self.CloseIfOTMNearExpiry)
    
    def CloseIfOTMNearExpiry(self):
        if self.Portfolio.TotalUnrealizedProfit < 0:
            for i in self.Portfolio:
                symbol = i.Key
                security = i.Value
                if security.Quantity == 0:
                    continue
                if security.IsLong:
                    self.Sell(symbol, security.Quantity)
                else:
                    self.Buy(symbol, security.Quantity)
            
    
    def OnOrderEvent(self, orderEvent):
        self.Log(str(orderEvent))