from Alphas.RsiAlphaModel import RsiAlphaModel
from Execution.ImmediateExecutionModel import ImmediateExecutionModel
from System import *
from System.Collections.Generic import List
from QuantConnect import *
from QuantConnect.Algorithm import QCAlgorithm
from QuantConnect.Data.UniverseSelection import *

'''
Simple testing strategy so that find all stocks for which the RSI is more than 20
Buy. - if RSI of today is more than 25, but yesterday was less than 25
Liquidate - if RSI of open positions goes above 75

SELL - if RSI today is less than 75 but yesterday was more than 75
Liquidate - if RSI of open positions goes below 25
'''
class ResistanceVentralThrustAssembly(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2018, 11, 20) # Set Start Date
        self.SetEndDate(2018, 12, 25)
        self.SetCash(10000)  # Set Strategy Cash
        
        # Strategy is for Daily data
        self.UniverseSettings.Resolution = Resolution.Daily
        self.SetWarmUp(timedelta(20))
        self.AddUniverse(self.MyCoarseFilterFunction)
        self.AddAlpha(RsiAlphaModel(60, Resolution.Daily))
        self.SetExecution(ImmediateExecutionModel())
        self.SetPortfolioConstruction(InsightWeightingPortfolioConstructionModel())
        self.SetRiskManagement(TrailingStopRiskManagementModel(0.03))
        self.__numberOfSymbols = 10
        self.stateData = {}

    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
        '''

        # if not self.Portfolio.Invested:
        #    self.SetHoldings("SPY", 1)
        pass
        
    def MyCoarseFilterFunction(self, coarse):
        
        # Get all symbols by dollar volume
        sortedByDollarVolume = sorted(coarse, key=lambda x: x.DollarVolume, reverse=True)
        filtered = [x for x in sortedByDollarVolume 
                      if x.Price > 10 ]
                      
        for sec in filtered:
        
            if sec.Symbol not in self.stateData:
                self.stateData[sec.Symbol] = SelectionData(self, sec.Symbol, 14)
                
            # rsiSelectionData = 
            self.stateData[sec.Symbol].update(sec.EndTime, sec.AdjustedPrice, sec.DollarVolume)
            
        validBuyValues = list(filter(lambda x : x.isBuyZone == True, self.stateData.values()))[:10]
        self.Debug('final buy signals')
        self.Debug(len(validBuyValues))
            
        return [x.symbol for x in validBuyValues]
    
class SelectionData(object):
    def __init__(self, algo, symbol, period):
        self.symbol = symbol
        self.rsi = RelativeStrengthIndex(period, MovingAverageType.Wilders)
        self.isBuyZone = False
        self.volume = 0
        self.algo = algo
        self.buySecurity = []

    def update(self, time, price, volume):
        self.volume = volume
        self.rsi.Update(time, price)
        
        if self.rsi.Current.Value < 30 and self.rsi.Current.Value > 10:
            self.isBuyZone = True
            if self.symbol not in self.buySecurity:
                self.buySecurity.append(self.symbol)