#TODO implment spread read for order pricing
import pandas as pd
import numpy as np
from Execution.ImmediateExecutionModel import ImmediateExecutionModel
from Portfolio.EqualWeightingPortfolioConstructionModel import EqualWeightingPortfolioConstructionModel

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


class VentralHorizontalThrustAssembly(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 11, 6) # Set Start Date
        self.SetEndDate(2020, 11, 28)
        self.SetCash(50000)  # Set Strategy Cash
        self.spy = self.AddEquity("SPY", Resolution.Daily)
        self.swir = self.AddEquity("SWIR", Resolution.Daily)
        self.clne = self.AddEquity("CLNE", Resolution.Daily)
        self.tsla = self.AddEquity("TSLA", Resolution.Daily)
        
        self.UniverseSettings.Resolution = Resolution.Daily
        
        
        self.SetSecurityInitializer(self.CustomSecurityInitializer)

        #self.AddAlpha(PsychSignalAlphaModel())
        #self.AddUniverse(self.CoarseSelectionFunction)
        #self.SetExecution(ImmediateExecutionModel())
        #self.SetSecurityInitializer(self.CustomSecurityInitializer)
        
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())

        self.__numberOfSymbols = 100
        
        self.Schedule.On(self.DateRules.EveryDay("SPY"), 
        self.TimeRules.BeforeMarketClose("SPY",16), Action(self.OpenPositions))
        
        
        def CustomSecurityInitializer(self, security):
            security.SetDataNormalizationMode(DataNormalizationMode.Raw)
        
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.At(9, 00), Action(self.BeforeTradingStart))
        
    def CustomSecurityInitializer(self, security):
        security.SetDataNormalizationMode(DataNormalizationMode.Raw)
    

        
    '''    
    def OnSecuritiesChanged(self, changes):
            for i in changes.AddedSecurities:
                self.Securities[i].FeeModel = ConstantFeeModel(0)
                self.Securities[i].SlippageModel = ConstantSlippageModel(0)   
    '''
    def OnData(self, data):
        '''OnData event is the primary entry point for your algorithm. Each new data point will be pumped in here.
            Arguments:
                data: Slice object keyed by symbol containing the stock data
        '''
        
        pass

    # sort the data by daily dollar volume and take the top 'NumberOfSymbols'
    def CoarseSelectionFunction(self, coarse):
        '''
        data_df = to_dataframe(coarse)
        my_universe = (data_df.
            query("(price > 1) and (price < 12.5) and (dollar_volume > 1500000) and has_fundamentals").
            #nlargest(20, 'dollar_volume').
            index.tolist())
        self.universe = my_universe    
        return self.universe
        
        '''
        # sort descending by daily dollar volume
        filterCoarse = [x for x in coarse if x.HasFundamentalData 
                                        and x.Price > 2.5
                                        and x.Price < 99 
                                        and x.DollarVolume > 1000000
                                        and x.DollarVolume < 50000000]
                                        
                                        
        #filterCoarse = [x for x in filteredCoarse if x.HasFundamentalData and x.Price < 12.5]
        #filterCoarse = [x for x in filteredCoarse if x.HasFundamentalData and x.DollarVolume < 1500000]
    
        # return the symbol objects of the top entries from our sorted collection
        #return [ x.Symbol for x in sortedByDollarVolume[:self.__numberOfSymbols] ]
        
        self.universe = [x.Symbol for x in filterCoarse]
        self.Log(len(self.universe))
        return self.universe
        
            
    def OpenPositions(self):
        for i in self.ActiveSecurities.Keys:
                self.Securities[i].FeeModel = ConstantFeeModel(0)
                self.Securities[i].SlippageModel = ConstantSlippageModel(0)
        self.MarketOnCloseOrder("SWIR", 1)
        self.MarketOnCloseOrder("CLNE", 1)
        self.MarketOnCloseOrder("TSLA", 1)

        
        
    def BeforeTradingStart(self):
        self.Liquidate()