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,7,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.
        # (This way we can use all the powerful builtin datafame methods)
        
        # Select 20 largest dollar_volume stocks with prices higher than 'price_above'.
        # Use 'has_fundamentals' to filter out ETFs and the like.
        # This function expects to return a list of symbols so use the 'tolist' method.
        
        data_df = to_dataframe(coarse_data)
        my_universe = (data_df.
            query("(price > 5) and has_fundamentals").
            nlargest(10, 'dollar_volume').
            index.tolist())
            
        # See how many securities are found in our universe
        context.Log("universe size is: {}".format(len(my_universe)))
        return my_universe
        
    
    def OnData(context, slice):
        # See how many securities we have in our Securities list
        context.Log("securities size is: {}".format(context.Securities.Count))


def to_dataframe(data_c):
    # Helper function to make a dataframe from the coarse object. 
    # Then we can use all the handy dataframe methods.
    # Set the 'coerce_float' parameter to get meaningful datatypes
    # otherwise all the fields will be un-useful generic objects.
    data = [(
            stock.Price, 
            stock.Volume, 
            stock.DollarVolume, 
            stock.HasFundamentalData)
            for stock in data_c]
            
    symbols = [stock.Symbol for stock in data_c ]
    labels = ['price', 'volume', 'dollar_volume', 'has_fundamentals']
    
    data_df = pd.DataFrame.from_records(
            data, 
            index=symbols, 
            columns=labels, 
            coerce_float=True)
            
    return data_df