import datetime
from QuantConnect.Securities.Option import OptionHolding

class StaddleStrategy(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2017, 1, 1)
        self.SetEndDate(2019, 3, 21)
        self.SetCash(10000)
        
        TICKER = 'SPY'
        
        self.underlying = self.AddEquity(TICKER, Resolution.Minute)
        self.option = self.AddOption(TICKER, Resolution.Minute)
        self.buy_qty = 1
        self.rolling_days = 1
        
        # Don't adjust by split and dividends, because the options strike price
        # is never adjusted. I need the real price to compare with the option
        # strike price
        self.underlying.SetDataNormalizationMode(DataNormalizationMode.Raw)
        
        self.option.SetFilter(-5, 5, datetime.timedelta(25), datetime.timedelta(60))
        self.SetBenchmark(TICKER)
        
        self.Schedule.On(
            self.DateRules.EveryDay(TICKER),
            self.TimeRules.BeforeMarketClose(TICKER, 60),
            Action(self.check_sell_staddle),
        )
        
        self.Schedule.On(
            self.DateRules.EveryDay(TICKER),
            self.TimeRules.BeforeMarketClose(TICKER, 59),
            Action(self.check_buy_staddle),
        )
        
        # Set the OnData slice. Mandatory for options.
        self.slice = None
        
        self.SetBenchmark(TICKER)
 

    def OnData(self, slice):
        self.slice = slice
        
    def OnEndOfDay(self):
        # Log leverage
        # account_leverage = self.Portfolio.TotalAbsoluteHoldingsCost / self.Portfolio.TotalPortfolioValue
        account_leverage = self.Portfolio.TotalHoldingsValue / self.Portfolio.TotalPortfolioValue

        self.Plot("Leverage", "Leverage", account_leverage)

    def check_buy_staddle(self):
        if not self.Portfolio.Invested:
            self.buy_staddle()

    def buy_staddle(self):
        call, put  = self._get_staddle_option_contracts()
        if put is None or call is None:
            self.Log('No contract with the same Strike price')
            return
        
        self.Log('Selected contracts for Staddle: %s and %s' % (
            call.Symbol.Value, put.Symbol.Value))
    
        self.Log('Buy %d opt for each contract' % self.buy_qty)
        self.MarketOrder(call.Symbol, self.buy_qty)
        self.MarketOrder(put.Symbol, self.buy_qty)
        
    def check_sell_staddle(self):
        # Get all options in portfolio
        all_options = list(filter(
            lambda opt: isinstance(opt, OptionHolding),
            self.Portfolio.Values)
        )
        holding_options = list(filter(lambda opt: opt.Quantity > 0,
                                      all_options))
                                      
        # Only touch my options, and not other asset of the portfolio
        holding_options = list(filter(
            lambda opt: opt.Symbol.Underlying == self.option.Underlying.Symbol,
            holding_options
        ))
        
        if len(holding_options) == 0:
            # No contract to sell
            return

        if len(holding_options) != 2:
            raise Exception('Expected to have 2 different options, but have %d' %
                len(holding_options))

        opt1 = holding_options[0]
        opt2 = holding_options[1]
        
        days_to_expire = (opt1.Security.Expiry - self.Time).days
        if days_to_expire <= self.rolling_days:
            self.Log('Options %s and %s expires in %d days' % (
                opt1.Symbol.Value, opt2.Symbol.Value, days_to_expire))
            self.sell_staddle([opt1, opt2])
        # else:
        #     self.Log('X-Options %s and %s expires in %d days' % (
        #         opt1.Symbol.Value, opt2.Symbol.Value, days_to_expire))

            
    def sell_staddle(self, options):
        for opt in options:
            self.MarketOrder(opt.Symbol, -1 * opt.Quantity)
                    
        
    def _get_staddle_option_contracts(self):
        if self.slice is None:
            self.Log('No slice. This should be a QuantConnect isssue')
            return None, None
        
        if self.slice.OptionChains is None:
            self.Log('No self.slice.OptionChains. This should be a QuantConnect isssue')
            return None, None
        
        underlying_chains = list(filter(
            lambda chain: chain.Key.Underlying == self.option.Underlying.Symbol,
            self.slice.OptionChains,
        ))

        if len(underlying_chains) == 0:
            self.Log('No option contract for underlying %s.' %
                     self.option.Underlying.Symbol.Value)
            return None, None
        
        underlying_chain = underlying_chains[0].Value

        # Filter options without volume
        opt_volume = list(filter(
            lambda opt: opt.Volume > 0,
            underlying_chain,
        ))
        
        # Get the puts
        # Keep only the puts
        puts = list(filter(
            lambda opt: opt.Right == OptionRight.Put,
            underlying_chain,
        ))

        # Keep only the calls
        calls = list(filter(
            lambda opt: opt.Right == OptionRight.Call,
            underlying_chain,
        ))

        put_symbols_str = ', '.join([opt.Symbol.Value for opt in puts])
        call_symbols_str = ', '.join([opt.Symbol.Value for opt in calls])
        self.Log('calls available=%s' % call_symbols_str)
        self.Log('puts available=%s' % put_symbols_str)
        
        # Sort the puts by distance against the underlying price. I sort by
        # puts because they use to be the contracts with lower volume
        underlying_price =  self.option.Underlying.Price
        priority_puts = sorted(
            puts,
            key=lambda opt: abs(opt.Strike - underlying_price),
            reverse=False,
        )
        priority_puts_str = ', '.join([opt.Symbol.Value for opt in priority_puts])
        
        self.Log('underlying price=%f' % underlying_price)
        self.Log('priority puts: %s' % priority_puts_str)
        
        put, call = self._get_pair_contracts(priority_puts, calls)
        
        return call, put
        
    def _get_pair_contracts(self, list_one, list_two):
        '''Find two contracts on different list with the same Strike price'''
        for contract_one in list_one:
            for contract_two in list_two:
                if (contract_one.Strike == contract_two.Strike and
                    contract_one.Expiry == contract_two.Expiry):
                    return contract_one, contract_two
                    
        # No contracts with same strike were found
        return None, None