import pandas as pd
class WellDressedSkyBlueSardine(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 12, 21)
        self.SetEndDate(2020, 12, 31)
        self.SetCash(1000000)
        
        self.SPY = self.AddEquity('SPY', Resolution.Daily).Symbol
        
        self.AddUniverse(self.CoarseFilter, self.FineFilter)
        self.UniverseSettings.Resolution = Resolution.Daily
        self.UniverseSettings.SetDataNormalizationMode = DataNormalizationMode.SplitAdjusted
        
        # self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.AfterMarketOpen('SPY', 30), self.daily_check)
        
        self.Data = {}
        self.SCTRuniverse = []
        
        self.percentagebuy = 0.05
        
    def CoarseFilter(self, coarse):
        sortedByDollarVolume = sorted(coarse, key=lambda x: x.DollarVolume, reverse=True)
        return [x.Symbol for x in sortedByDollarVolume if x.Price > 10
                                                and x.HasFundamentalData][:1000]

    def FineFilter(self, fine):
        sortedByPE = sorted(fine, key=lambda x: x.MarketCap)
        universe =  [x.Symbol for x in sortedByPE if x.MarketCap > 10e9][:500]
        return universe

    def OnSecuritiesChanged(self, changes):
        # close positions in removed securities
        for x in changes.RemovedSecurities:
            self.Liquidate()
            if x.Symbol in self.Data:
                del self.Data[x.Symbol]
        
        # can't open positions here since data might not be added correctly yet
        for x in changes.AddedSecurities:
            self.Data[x.Symbol] = SymbolData(self, x.Symbol)
        
    def OnData(self, data):
        
        if self.IsWarmingUp: return
            
        SCTR = {}
        
        for symboldata in self.Data.values():
            if symboldata.IsReady():
                SCTR[symboldata.symbol] = symboldata.SCTR()

        self.SCTRuniverse = [pair[0] for pair in sorted(SCTR.items(), key=lambda kv: kv[1], reverse=True)[:int(len(SCTR)*0.15)]]
        
        for symbol in self.SCTRuniverse:
            if not data.ContainsKey(symbol): continue
            
            self.Data[symbol].Update()
            
            # won't make a difference in setting specific time to trade if you use daily resolution
            symbolClose = float(self.Securities[symbol].Close)
            
            if (not self.Securities[symbol].Invested) & (self.Portfolio.MarginRemaining > 0.9*self.percentagebuy*self.Portfolio.TotalPortfolioValue):
                sma = self.Data[symbol].get_sma()
                sto = self.Data[symbol].get_sto()
                
                if (symbolClose > sma) & (sto <= 30):
                    self.SetHoldings(symbol, self.percentagebuy)
                    self.Log(str(symbol) + " -sto: " + str(sto))
                    self.Log(str(symbol) + " -sma: " + str(sma) + " -price: " + str(symbolClose))

class SymbolData:
    
    def __init__(self, algo, symbol):
        
        self.symbol = symbol
        
        self.EMA200 = algo.EMA(symbol, 200, Resolution.Daily)
        self.EMA50 = algo.EMA(symbol, 50, Resolution.Daily)
        self.ROC125 = algo.ROC(symbol, 125, Resolution.Daily)
        self.ROC20 = algo.ROC(symbol, 20, Resolution.Daily)
        self.PPO = algo.PPO(symbol, 12, 26, MovingAverageType.Exponential, Resolution.Daily)
        self.RSI14 = algo.RSI(symbol, 14, MovingAverageType.Simple, Resolution.Daily)
        self.SMA200 = algo.SMA(symbol, 200, Resolution.Daily)
        self.Slow_Stoch5 = algo.STO(symbol, 5, Resolution.Daily)

        self.PPOWindow = RollingWindow[float](3)
        
        history = algo.History(symbol, 200, Resolution.Daily)
        for index, bar in history.loc[symbol].iterrows():
            self.EMA200.Update(index, bar.close)
            self.EMA50.Update(index, bar.close)
            self.ROC125.Update(index, bar.close)
            self.ROC20.Update(index, bar.close)
            self.PPO.Update(index, bar.close)
            self.RSI14.Update(index, bar.close)
            self.SMA200.Update(index, bar.close)
            
            self.PPOWindow.Add(self.PPO.Current.Value)
            
        for bar in history.iloc[-5:].itertuples():
            tradeBar = TradeBar(bar.Index[1], bar.Index[0], bar.open, bar.high, bar.low, bar.close, bar.volume, timedelta(1))
            self.Slow_Stoch5.Update(tradeBar)
            
    def Update(self):
        self.PPOWindow.Add(self.PPO.Current.Value)
        
    def IsReady(self):
        return self.EMA200.IsReady and self.EMA50.IsReady \
                    and self.ROC125.IsReady and self.ROC20.IsReady\
                    and self.PPO.IsReady and self.RSI14.IsReady and self.PPOWindow.IsReady
                    
    def SCTR(self):
        return self.EMA200.Current.Value*0.3 + self.EMA50.Current.Value*0.15 + self.ROC125.Current.Value*0.3 + self.ROC20.Current.Value*0.15 \
                + (self.PPOWindow[0] - self.PPOWindow[2])/(3*self.PPOWindow[2]) *0.05 + self.RSI14.Current.Value*0.05
        
    def get_sma(self):
        return self.SMA200.Current.Value
        
    def get_sto(self):
        return self.Slow_Stoch5.Current.Value