from datetime import timedelta
from itertools import combinations
import pandas as pd

class BearPutSpreadAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 1, 1)
        self.SetCash(10000)
        
        self.spy = self.AddEquity("SPY", Resolution.Minute).Symbol
        option = self.AddOption("SPY", Resolution.Minute)
        self.symbol = option.Symbol
    
        option.SetFilter(-5, 5, timedelta(1), timedelta(90))
        
        self.SetBenchmark(self.spy)
        
        
    def OnData(self, data):
        if self.Portfolio.Invested or \
            not data.OptionChains.ContainsKey(self.symbol) or \
            not data.ContainsKey(self.spy):
            return 
        
        for symbol, chain in data.OptionChains.items():
            contracts = [contract for contract in chain]
            
        if len(contracts) == 0:
            return
        
        underlying_price = data[self.spy].Close
        self.TradeOptions(contracts, underlying_price)
        
 
    def TradeOptions(self, contracts, underlying_price):
        # Get all put contracts
        all_puts = [contract for contract in contracts if contract.Right == OptionRight.Put]
        
        # Get all unique expiries
        expires = set([contract.Expiry for contract in all_puts])
        
        rankings_df = pd.DataFrame()
        for expiry in expires:
            puts = [contract for contract in all_puts if contract.Expiry == expiry]
            if len(puts) == 0:
                continue
            
            for option_combinations in combinations(puts, 2):
                buy_put = option_combinations[0]
                sell_put = option_combinations[1]
                
                strike_width = buy_put.Strike - sell_put.Strike
                net_cost = buy_put.LastPrice - sell_put.LastPrice
                
                # We ignore the contract multiplier here
                max_loss = net_cost 
                max_profit = strike_width - net_cost
                
                # Calculate factors
                inverted_profit_loss_ratio = max_loss / max_profit if max_profit != 0 else float('inf')
                break_even_distance = underlying_price - buy_put.Strike + net_cost
                days_to_expiry = (expiry - self.Time).days
                
                # Save factor results
                row = pd.DataFrame({'inverted_profit_loss_ratio' : [inverted_profit_loss_ratio],
                                    'break_even_distance' : [break_even_distance],
                                    'days_to_expiry' : [days_to_expiry]}, 
                                    index=[option_combinations])
                rankings_df = rankings_df.append(row)
                
        if not rankings_df.empty:
            # Rank put contracts by factors
            selected_contracts = rankings_df.rank().sum(axis=1).idxmin()
            
            # Create Bear Put Spread
            buy_symbol = selected_contracts[0].Symbol
            sell_symbol = selected_contracts[1].Symbol
            
            quantity = self.CalculateOrderQuantity(buy_symbol, 0.5)
            
            if quantity > 0:
                self.Buy(buy_symbol, quantity)
                self.Sell(sell_symbol, quantity)