from clr import AddReference
AddReference("System")
AddReference("QuantConnect.Algorithm")
AddReference("QuantConnect.Indicators")
AddReference("QuantConnect.Common")

from QuantConnect.Data.UniverseSelection import *
import numpy as np
from System import *
from QuantConnect import *
from QuantConnect.Algorithm import QCAlgorithm
from QuantConnect.Data.Market import *
from QuantConnect.Data.UniverseSelection import *

class BasicTemplateAlgorithm(QCAlgorithm):

    def Initialize(self):

        self.SetStartDate(2012, 1, 1)
        self.SetEndDate(2012, 1, 31)
        self.SetCash(10000)
        self.SetWarmUp(250)
        
        self.stock_score = {}
        self.returns = {}
        self.stock_52w = []
        self.fine = []
        self.basket = []
        self.sma_20 = None
        self.sma_50 = None
        self.sma_200 = None
        self.changes = None
        
        self.AddEquity('SPY', Resolution.Daily)
        
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
        
        self.Schedule.On(self.DateRules.WeekStart('SPY'), self.TimeRules.AfterMarketOpen('SPY', 1), Action(self.Rebalance))
        self.weekly_rebalance = False
        
    def CoarseSelectionFunction(self, coarse):
        
        if self.weekly_rebalance:
            CoarseWithFundamental = [x for x in coarse if (x.HasFundamentalData) and (int(x.Volume) > 100000)]
            return CoarseWithFundamental
        else:
            return []
    
    def FineSelectionFunction(self, fine):
        
        if self.weekly_rebalance:
            fine = [x for x in fine if (float(x.ValuationRatios.FCFRatio) < 20) and
                    (float(x.OperationRatios.ROIC.OneYear) > 0) and
                    (float(x.OperationRatios.ROE.OneYear) > 0) and
                    (float(x.OperationRatios.TotalDebtEquityRatio.OneYear) < 1)]
            sortedByROE = sorted(fine, key = lambda x: x.OperationRatios.ROE.OneYear, reverse = True)[:250]
            self.fine = [x.Symbol for x in sortedByROE]
            self.weekly_rebalance = False
            
            self.Debug(str(self.Time) + ' Fine Universe: ' + str(self.fine))
            
            self.stock_score = {}
            for stock in self.fine:
                self.sma_20 = self.SMA(stock, 20, Resolution.Daily)
                self.sma_50 = self.SMA(stock, 50, Resolution.Daily)
                self.sma_200 = self.SMA(stock, 200, Resolution.Daily)
        
            if not self.sma_200.IsReady:
                return
            else:
                score = 0
                score += np.where(self.Securities[stock].Price > self.sma_20.Current.Value, 1, 0)
                score += np.where(self.sma_20.Current.Value > self.sma_50.Current.Value, 1, 0)
                score += np.where(self.sma_50.Current.Value > self.sma_200.Current.Value, 1, 0)
                self.stock_score[stock] = score
        
            self.stocks = [key for key, value in self.stock_score.items() if value == 3]
    
            slices = self.History(self.stocks, 253, Resolution.Daily)
            self.stocks_52w = []
            if not slices.empty:
                for i in self.stocks:
                    df = slices.loc[i]
                    if self.Securities[i].Price >= df['high'].nlargest(1) * 0.98:
                        self.stocks_52w.append(i)
        
            self.returns = {}
            if not slices.empty:
                for i in self.stocks_52w:
                    df = slices.loc[i]
                    prev = df['close'][0]
                    curr = df['close'][-1]
                    self.returns[i] = ((curr - prev) / prev) * 100
        
            self.rank_returns = {key: rank for rank, key in enumerate(sorted(self.returns, key = self.returns.get, reverse = True), 1)}
            self.basket = [key for key, value in self.rank_returns.items() if value <= 10]
            
            return self.basket
        else:
            return []
    
    def Rebalance(self):
        self.weekly_rebalance = True
    
    def OnData(self, data):
        
        stock_weight = 0.099
        
        if self.changes is None:
            return
        
        for security in self.changes.RemovedSecurities:
            if security.Invested:
                self.Liquidate(security.Symbol)
        
        for security in self.changes.AddedSecurities:
            if not security.Invested:
                self.SetHoldings(security.Symbol, stock_weight)
        
        self.changes = None
    
    def OnSecuritiesChanged(self, changes):
        self.changes = changes