from QuantConnect.Algorithm import *
from QuantConnect.Securities.Option import OptionPriceModels
import pandas as pd

class BasicTemplateOptionsFilterUniverseAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2019, 1, 1)
        self.SetEndDate(2019, 4, 1)
        self.SetCash(100000)

        self.goog_option = self.AddOption("GOOG")
        self.aapl_option = self.AddOption("AAPL")
        
        self.goog_option.SetFilter(-3, +1, 15, 60) 
        self.aapl_option.SetFilter(-3, +1, 15, 60) 
        
    def OnData(self,slice):
        
        if self.IsWarmingUp or self.Portfolio.Invested:
            return 
        
        self.combinations = pd.DataFrame()
        for symbol, chain in slice.OptionChains.items():
        
            for contract in chain:
                row = pd.DataFrame({'symbol': [contract.Symbol], 'right': [contract.Right], 'expiry': [contract.Expiry], 'strike': [contract.Strike]})
                self.combinations = self.combinations.append(row)
                
        self.Log(f"\n{self.combinations.to_string()}")
        
        # Calculate number of chains per combination
        chains_per_combo = pd.DataFrame()
        df = self.combinations
        for i, row in df.iterrows():
            num_chains = df[(df.right == row.right) & (df.expiry == row.expiry) & (df.strike == row.strike)].shape[0]
            if chains_per_combo.shape[0] == 0 or \
                chains_per_combo[(chains_per_combo.right == row.right) & (chains_per_combo.expiry == row.expiry) & (chains_per_combo.strike == row.strike)].shape[0] == 0:
                sub_row = pd.DataFrame({'chains':[num_chains], 'right': [row.right], 'expiry': [row.expiry], 'strike': [row.strike]})
                chains_per_combo = chains_per_combo.append(sub_row)
        
        self.Log(f"\n{chains_per_combo.to_string()}")
        
        self.Quit()