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

class BasicTemplateAlgorithm(QCAlgorithm):
    
    def Initialize(context):
        context.my_universe = []
        
        # Set the initial cash, start, and end dates for backtesting
        context.SetCash(100000)
        context.SetStartDate(2011,5,1)
        context.SetEndDate(2011,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, context.universe_filter_fine)
        
        # Schedule the 'trade' method to run
        # Add SPY to our data subscription so we can schedule from it (use defaults)
        context.AddEquity('SPY')
        context.AddEquity('AAPL')

        context.Schedule.On(
            context.DateRules.EveryDay('SPY'), 
            context.TimeRules.AfterMarketOpen('SPY', 10),       
            Action(context.my_trade))
        
    def universe_filter_course(context, coarse_data):
        # First convert the 'course' data to a pandas dataframe.

        # 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.
        columns = ['Price', 'DollarVolume', 'HasFundamentalData', 'Volume']
        column_names = ['price', 'dollar_volume', 'has_fundamentals', 'volume']
        data_df = to_dataframe(coarse_data, columns, column_names)
        
        my_universe = (data_df.
            query("(price > 5) and has_fundamentals").
            nlargest(1, 'dollar_volume'))
            
        # See how many securities are found in our universe
        #context.Log("course size is: {}".format(my_universe.volume))
        
        return my_universe.index.tolist()
            
 
 
    def universe_filter_fine(context, fine_data):
        # First convert the 'fine' data to a pandas dataframe.
        symbols = [stock.Symbol for stock in fine_data]
        data = [(float(stock.Price),)
            for stock in fine_data]
        
        labels = ['ev_to_ebitda']
        
        data_df = pd.DataFrame.from_records(
            data, 
            index=symbols, 
            columns=labels,)
            
        context.Log("data is: {}".format((data_df.ev_to_ebitda.values)))
        
        context.my_universe = (data_df.
            nlargest(1, 'ev_to_ebitda').
            index.tolist())  
        
        context.Log("fine size is: {}".format(len(context.my_universe)))
        
        return context.my_universe
            

    def my_trade(context):
        # See how many securities we have in our Securities list
        '''
        for universe in context.UniverseManager.Keys:
            context.Log("universe name: {}".
                format(universe))
        '''        
        for universe in context.UniverseManager.Values:
            context.Log("universe count: {}".
                format(universe.Members.Count))
                
        if context.my_universe:       
            data_df = context.History(['AAPL'], 3)
            context.Log("universe is: {}".format((context.my_universe)))
            context.Log("history is: {}".format((data_df.close)))

        if len(context.my_universe) > 0:
            weight = 1.0 / len(context.my_universe)
        else:
            weight = 0.0

        for security in context.my_universe:
            context.SetHoldings(security, weight)


def to_dataframe(data_c, properties, labels):
    # 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 = [[getattr(stock, name) for name in properties] for stock in data_c]
    symbols = [stock.Symbol for stock in data_c ]
    
    data_df = pd.DataFrame.from_records(
            data, 
            index=symbols, 
            columns=labels, 
            coerce_float=True)
            
    return data_df