#region imports
from AlgorithmImports import *
#endregion
#This is a Template of dynamic stock selection.
#You can try your own fundamental factor and ranking method by editing the CoarseSelectionFunction and FineSelectionFunction

from QuantConnect.Data.UniverseSelection import *

class BasicTemplateAlgorithm(QCAlgorithm):
    
    def __init__(self):
    # set the flag for rebalance
        self.reb = 1
    # Number of stocks to pass CoarseSelection process
        self.num_coarse = 100
    # Number of stocks to long/short
        self.num_fine = 10
        self.symbols = None
        self.mom = {}           # Dict of Momentum indicator keyed by Symbol
        self.lookback = 252     # Momentum indicator lookback period

    def Initialize(self):
        self.SetCash(100000)
        self.SetStartDate(2009, 7, 1)  # Set Start Date
        self.SetEndDate(2019, 7, 1)    # Set Start Date   
        # self.SetStartDate(2000,1,1)
    # if not specified, the Backtesting EndDate would be today 
        # self.SetEndDate(2017,1,1)
        
        
        self.spy = self.AddEquity("SPY", Resolution.Daily).Symbol
        # self.sma = self.SMA("SPY", 5, Resolution.Daily)
        self.UniverseSettings.Resolution = Resolution.Daily
        
        self.AddUniverse(self.CoarseSelectionFunction,self.FineSelectionFunction)
        
    # Schedule the rebalance function to execute at the begining of each month
        self.Schedule.On(self.DateRules.MonthStart(self.spy), 
        self.TimeRules.AfterMarketOpen(self.spy,5), Action(self.rebalance))
    
    def CoarseSelectionFunction(self, coarse):
    # if the rebalance flag is not 1, return null list to save time.
        if self.reb == 0:
            return Universe.Unchanged
            # return self.long + self.short
        
        selected = [x for x in coarse if (x.HasFundamentalData) 
                    and (float(x.Price) > 1)]
        sortedByDollarVolume = sorted(selected, key=lambda x: x.DollarVolume, reverse=True) 
        top = sortedByDollarVolume[:self.num_coarse]
        topSymbols = [i.Symbol for i in top]


        return topSymbols
    def FineSelectionFunction(self, fine):
    # return null list if it's not time to rebalance
        if self.reb == 0:
            return Universe.Unchanged
        self.reb = 0

        toPop = [security for security in self.mom if security not in [i.Symbol for i in fine]]
        for security in toPop:
            self.mom.pop(security)
        for security in fine:
            if security.Symbol not in self.Securities:
                self.AddEquity(security.Symbol, Resolution.Daily)
            self.mom[security.Symbol] = Momentum(self.lookback)
        addedSymbols = [k for k,v in self.mom.items() if not v.IsReady]
        history = self.History(addedSymbols, 1 + self.lookback, Resolution.Daily)
        history = history.close.unstack(level=0)
        for symbol in addedSymbols:
            ticker = str(symbol)
            if ticker in history:
                for time, value in history[ticker].items():
                    try:
                        item = IndicatorDataPoint(symbol, time, value)
                        self.mom[symbol].Update(item)
                    except:
                        self.Debug(value)
        filtered_fine = [x for x in fine if
                                        x.Symbol in self.mom and self.mom[x.Symbol].IsReady
                                        and x.OperationRatios.OperationMargin.ThreeMonths
        #                                 # and x.ValuationRatios.EVToForwardEBIT
        #                                 # and x.FinancialStatements.IncomeStatement.GrossProfit.ThreeMonths
        #                                 and x.FinancialStatements.IncomeStatement.TotalRevenue.ThreeMonths
        #                                 and x.FinancialStatements.IncomeStatement.SellingGeneralAndAdministration.ThreeMonths
        #                                 and x.OperationRatios.RevenueGrowth.ThreeMonths
        #                                 and x.ValuationRatios.CashReturn
        #                                 and x.OperationRatios.OperationIncomeGrowth.ThreeMonths
                                        # and x.OperationRatios.NetMargin.ThreeMonths
        #                                 and x.OperationRatios.ROIC.ThreeMonths
                                        ]
        self.Debug('remained to select %d'%(len(filtered_fine)))
        operationMargin = sorted(filtered_fine, key=lambda x: x.OperationRatios.OperationMargin.Value, reverse=True)
        # EVToForwardEBIT = sorted(filtered_fine, key=lambda x: x.ValuationRatios.EVToForwardEBIT, reverse=True)
        # grossMargin = sorted(filtered_fine, key=lambda x: x.FinancialStatements.IncomeStatement.GrossProfit.ThreeMonths/max(x.FinancialStatements.IncomeStatement.TotalRevenue.ThreeMonths,1), reverse=True)
        # sgaMargin = sorted(filtered_fine, key=lambda x: x.FinancialStatements.IncomeStatement.SellingGeneralAndAdministration.ThreeMonths/max(x.FinancialStatements.IncomeStatement.TotalRevenue.ThreeMonths,1), reverse=True)
        # revenueGrowth = sorted(filtered_fine, key=lambda x: x.OperationRatios.RevenueGrowth.ThreeMonths, reverse=True)
        # operationIncomeGrowth = sorted(filtered_fine, key=lambda x: x.OperationRatios.OperationIncomeGrowth.ThreeMonths, reverse=True)
        # netMargin = sorted(filtered_fine, key=lambda x: x.OperationRatios.NetMargin.ThreeMonths, reverse=True)
        # roic = sorted(filtered_fine, key=lambda x: x.OperationRatios.ROIC.ThreeMonths, reverse=True)
        stock_dict = {}
        mom = sorted(filtered_fine, key=lambda x: self.mom[x.Symbol].Current.Value, reverse=True)
        # mom = sorted([v for k,v in self.mom.items() if v.IsReady], key=lambda x: x.Current.Value, reverse=True)
        for i,ele in enumerate(operationMargin):
            
            # operationMargin_R = operationMargin.index(ele)
            # EVToForwardEBIT_R = EVToForwardEBIT.index(ele)
            # grossMargin_R = grossMargin.index(ele)
            # sgaMargin_R = sgaMargin.index(ele)
            # roic_R = roic.index(ele)
            # operationIncomeGrowth_R = operationIncomeGrowth.index(ele)
            # revenueGrowth_R = revenueGrowth.index(ele)
            #score = sum([i, operationMargin_R,roic_R])
            mom_r = mom.index(ele)
            score = sum([mom_r])
            stock_dict[ele] = score
        
        # sort the stocks by their scores
        self.sorted_stock = sorted(stock_dict.items(), key=lambda d:d[1],reverse=False)
        sorted_symbol = [x[0] for x in self.sorted_stock]

        # sotre the top stocks into the long_list and the bottom ones into the short_list
        self.long = [x.Symbol for x in sorted_symbol[:self.num_fine]]
        self.short = [x.Symbol for x in sorted_symbol[-self.num_fine:]]

        return self.long + self.short

    def OnData(self, data):
        # pass
        for symbol, mom in self.mom.items():
            mom.Update(self.Time, self.Securities[symbol].Close)
    
    def rebalance(self):
        self.Debug("rebalance")
        if self.Time.month % 4 > 0:
            return
    # if this month the stock are not going to be long/short, liquidate it.
        long_short_list = self.long + self.short
        for i in self.Portfolio.Values:
            if (i.Invested) and (i not in long_short_list):
                self.Liquidate(i.Symbol)
        
    # Assign each stock equally. Alternatively you can design your own portfolio construction method
        for i in self.long:
            self.SetHoldings(i, 1/self.num_fine)
        # always short
        # for i in self.short:
        #     self.SetHoldings(i, -1/self.num_fine)
        # tactical short
        # if self.Securities["SPY"].Price < self.sma.Current.Value:
        #     for i in self.short:
        #         self.SetHoldings(i, -1/self.num_fine)
        # else:
        #     for i in self.Portfolio.Values:
        #         if (i.Invested) and (i.IsShort):
        #             self.Liquidate(i.Symbol)
        self.reb = 1