from QuantConnect.Data.UniverseSelection import *
import pandas as pd

class BasicTemplateAlgorithm(QCAlgorithm):
    
    def Initialize(context):
        
        # Set the initial cash, start, and end dates for backtesting
        context.SetCash(100000)
        context.SetStartDate(2016,5,1)
        context.SetEndDate(2016,5,10)
        
        # Subscribe to data for securities returned in my universe 
        context.UniverseSettings.Resolution = Resolution.Daily
        context.UniverseSettings.MinimumTimeInUniverse = 0 
        context.AddUniverse(context.universe_filter_course)
        
        
    def universe_filter_course(context, coarse_data):
        # First convert the 'course' data to a pandas dataframe.
        # Use coerce parameter = True to convert data objects to numbers
        symbols = [stock.Symbol for stock in coarse_data ]
        data = [(
                stock.DollarVolume, 
                stock.Price,
                stock.HasFundamentalData)
                for stock in coarse_data]
                
        column_names = ['dollar_volume', 'price', 'has_fundamentals']
                    
        data_df = pd.DataFrame.from_records(
            data, 
            index=symbols, 
            columns=column_names, 
            coerce_float=True)
        
        # Use pandas methods to select the assets we want
        # First find the values of dollar_volume at our desired bounds
        lower_percent = data_df.dollar_volume.quantile(.85)
        upper_percent = data_df.dollar_volume.quantile(1.0)
        
        # Now simply query using those values
        # Filter for has_fundamentals to remove ETFs
        my_universe = (data_df.
            query('has_fundamentals & (dollar_volume >= @lower_percent) & (dollar_volume <= @upper_percent)'))
        
        # See how many securities are found in our universe
        context.Log("{} securities found ".
            format(my_universe.shape[0]))
        
        context.Log("top 3 securities: {}".
            format(my_universe.nlargest(3, 'dollar_volume').dollar_volume))
                       
        context.Log("bottom 3 securities: {}".
            format(my_universe.nsmallest(3, 'dollar_volume').dollar_volume))
        
        # Expects a list of symbols returned
        return my_universe.index.tolist()
 

    def OnData(context, data):
        pass