class NadionParticleThrustAssembly(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2021, 2, 1)  # Set Start Date
        self.SetCash(100000)  # Set Strategy Cash
        self.SetUniverseSelection(FineFundamentalUniverseSelectionModel(self.CoarseSelectionFunction, self.FineSelectionFunction, None, None))
        self.quandl_symbol_by_symbol = {}
            
    def CoarseSelectionFunction(self, coarse):
        # Sort by dollar volume
        has_fundamental_data = [c for c in coarse if c.HasFundamentalData]
        top_dollar_volume = sorted(coarse, key=lambda x: x.DollarVolume, reverse=True)
        return [x.Symbol for x in top_dollar_volume[:50]]
    
    def FineSelectionFunction(self, fine):

        # Gather SharesOutstanding & Subscribe to Quandl datafeed for most liquid symbols
        shares_outstanding_by_symbol = {}
        for f in fine:
            symbol = f.Symbol
            shares_outstanding_by_symbol[symbol] = f.CompanyProfile.SharesOutstanding
            if symbol in self.quandl_symbol_by_symbol:
                continue
            self.quandl_symbol_by_symbol[symbol] = self.AddData(QuandlFINRAData, f'FINRA/FNSQ_{symbol.Value}', Resolution.Daily).Symbol
        
        # Gather short volume of the most liquid symbols
        short_volume_by_symbol = {}
        symbols_to_remove = []
        for symbol, quandl_symbol in self.quandl_symbol_by_symbol.items():
            if symbol in shares_outstanding_by_symbol:
                history = self.History(quandl_symbol, 1)
                if history.empty or 'shortvolume' not in history.columns:
                    continue
                short_volume_by_symbol[symbol] = history.loc[quandl_symbol].iloc[0]['shortvolume']
            else:
                # Remove Quandl data feed
                symbols_to_remove.append(symbol)
                self.RemoveSecurity(quandl_symbol)
                
        # Remove symbols from the dictionary that don't have an active quandl subscription
        for symbol in symbols_to_remove:
            self.quandl_symbol_by_symbol.pop(symbol, None)
            
        # Return symbols with highest short volume
        sorted_by_short_volume_ratio = sorted(short_volume_by_symbol.items(), key=lambda x: x[1] / shares_outstanding_by_symbol[x[0]], reverse=True)
        self.symbols = [symbol for symbol, _ in sorted_by_short_volume_ratio[:5]]
        return self.symbols
        
    
    def OnData(self, data):
        ## Access the value of the Quandl data using the standard accessor
        for symbol in self.symbols:
            quandl_symbol = self.quandl_symbol_by_symbol[symbol]
            if data.ContainsKey(quandl_symbol) and data[quandl_symbol] is not None:
                self.Log(f"Short interest of {str(symbol)}: {data[quandl_symbol].Value}")

class QuandlFINRAData(PythonQuandl):
    
    def __init__(self):
        ## Rename the Quandl object column to the data we want, which is the 'ShortVolume' column
        ## of the CSV that our API call returns
        self.ValueColumnName = 'ShortVolume'