#region imports
from AlgorithmImports import *
#endregion
class NadionParticleThrustAssembly(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2022, 8, 1)  # Set Start Date
        self.SetCash(100000)  # Set Strategy Cash
        self.AddUniverse(self.MyCoarseFilterFunction)
        self.quandl_symbol_by_symbol = {}
        self.symbols_wo_data = ["BRK.A", "BRK.B"]
            
    def MyCoarseFilterFunction(self, coarse):
        # Sort by dollar volume
        top_dollar_volume = sorted(coarse, key=lambda x: x.DollarVolume, reverse=True)[:50]
        top_dollar_volume = [x.Symbol for x in top_dollar_volume]
        
        # Subscribe to Quandl datafeed for most liquid symbols
        for symbol in top_dollar_volume:
            if symbol in self.quandl_symbol_by_symbol:
                continue
            if symbol.Value not in self.symbols_wo_data:
                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 top_dollar_volume:
                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 = sorted(short_volume_by_symbol.items(), key=lambda x: x[1], reverse=True)
        self.symbols = [symbol for symbol, _ in sorted_by_short_volume[: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'