from datetime import timedelta, datetime
from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel
from QuantConnect.Data.UniverseSelection import * 

class HorizontalQuantumCompensator(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2019, 8, 10)  # Set Start Date
        self.SetCash(100000)
        self.AddUniverseSelection
        self.UniverseSettings.Resolution = Resolution.Daily
        self.UniverseSettings.DataNormalizationMode = DataNormalizationMode.Raw
        self.SetAlpha(MOMAlphaModel())
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
        self.SetExecution(ImmediateExecutionModel())
        self.AddUniverseSelection(LiquidValueUniverseSelectionModel())
        
class LiquidValueUniverseSelectionModel(FundamentalUniverseSelectionModel):
    def __init__(self):
        super().__init__(True, None, None)
        
    def SelectCoarse(self, algorithm, coarse):
        sortedByDollarVolume = sorted([x for x in coarse if x.HasFundamentalData],key=lambda x: x.DollarVolume, reverse=True)
        return [x.Symbol for x in sortedByDollarVolume[:200]]
        
class MOMAlphaModel(AlphaModel): 
    
    def __init__(self):
         self.data = {}
    
    def OnSecuritiesChanged(self, algorithm, changes):
        for security in changes.AddedSecurities:
            symbol = security.Symbol
            if symbol not in self.data:
               self.data[symbol] = SymbolData(algorithm, symbol)
    
    def Update(self, algorithm, data):
        insights = []
        
        for symbol, symbolData in self.data.items():
            SMA20 = symbolData.SMAS.Current.Value 
            SMA50 = symbolData.SMAL.Current.Value
            RSI = symbolData.RSI.Current.Value
            
            if SMA20>SMA50 and  RSI < 30 :
                insights += [Insight.Price(symbol,timedelta(1),InsightDirection.Up)]
            if SMA20<SMA50 and  RSI > 70 :
                insights += [Insight.Price(symbol,timedelta(1),InsightDirection.Down)]
                
        return insights
        
class SymbolData:
    
    def __init__(self, algorithm, symbol):
        self.algorithm = algorithm
        self.symbol = symbol
        self.SMAS = algorithm.SMA(symbol, 20, Resolution.Daily)
        self.SMAL = algorithm.SMA(symbol, 50, Resolution.Daily)
        self.RSI = algorithm.RSI(symbol, 14, Resolution.Daily)
        
        history = algorithm.History(symbol, 50, Resolution.Daily)
        
        if not history.empty:
            history = history.close.unstack(0)
            if not history.empty:
                df = history[symbol]
                for time, close in df.iteritems():
                    self.SMAS.Update(time, close)
                    self.SMAL.Update(time, close)
                    self.RSI.Update(time, close)