from AlgorithmImports import *
import math 

class PortfolioModelJGG(PortfolioConstructionModel):

    def __init__(self, rebalance = Resolution.Daily, portfolioBias = PortfolioBias.LongShort):
        self.portfolioBias = portfolioBias
        self.insightTargets = {}
        self.percentRisk = 1.0
        self.marginCallPercentBuffer = 25
        # If the argument is an instance of Resolution or Timedelta
        # Redefine rebalancingFunc
        rebalancingFunc = rebalance
        if isinstance(rebalance, int):
            rebalance = Extensions.ToTimeSpan(rebalance)
        if isinstance(rebalance, timedelta):
            rebalancingFunc = lambda dt: dt + rebalance
        if rebalancingFunc:
            self.SetRebalancingFunc(rebalancingFunc)

    def CreateTargets(self, algorithm, newInsights):
        '''Will determine the target percent for each insight
        Args:
            activeInsights: The active insights to generate a target for'''
        timeDateString = format(algorithm.Time)
        timeString = timeDateString.split()
        unbufferedTargets = {}
        portfolioTargets = []
        equity=algorithm.Portfolio.TotalPortfolioValue
        dollarRisk = equity*self.percentRisk/100
        #algorithm.Log(timeDateString)
        
        #try:
        #algorithm.Log('Prior Insight Targets before deleting expired')     
        for existingSymbol in self.insightTargets:
            for insightIndex in reversed(range(len(self.insightTargets[existingSymbol][1]))):
                algorithm.Log(str(existingSymbol)+' '+str(self.insightTargets[existingSymbol][0][insightIndex])+' '+str(self.insightTargets[existingSymbol][1][insightIndex]))
                if self.insightTargets[existingSymbol][1][insightIndex] < algorithm.UtcTime: 
                    del(self.insightTargets[existingSymbol][0][insightIndex])
                    del(self.insightTargets[existingSymbol][1][insightIndex])
                        
        #algorithm.Log('Prior Insight Targets after deleting expired')            
        #for existingSymbol in self.insightTargets:
            #for insightIndex in reversed(range(len(self.insightTargets[existingSymbol][1]))):            
                #algorithm.Log(str(existingSymbol)+' '+str(self.insightTargets[existingSymbol][0][insightIndex])+' '+str(self.insightTargets[existingSymbol][1][insightIndex]))
                        
        #algorithm.Log('New Insights')
        for insight in newInsights:
            #algorithm.Log(str(insight.Symbol)+' '+str(insight.CloseTimeUtc))
            volatilityRisk = insight.Magnitude/(1 + insight.Magnitude)                
            biasMultiplier = 1 if self.RespectPortfolioBias(insight) else 0
            # The block of code below only works for long positions!
            insightTarget = math.floor(biasMultiplier*dollarRisk/volatilityRisk/algorithm.Securities[insight.Symbol].Price) #simplify this calculation
            #algorithm.Log(str(biasMultiplier)+' '+str(volatilityRisk)+' '+str(algorithm.Securities[insight.Symbol].Price)+' '+str(insightTarget))
            if insight.Symbol not in self.insightTargets:
                self.insightTargets[insight.Symbol] = [[insightTarget],[insight.CloseTimeUtc]]
            elif len(self.insightTargets[insight.Symbol]) == 0:
                self.insightTargets[insight.Symbol] = [[insightTarget],[insight.CloseTimeUtc]]
            else:    
                self.insightTargets[insight.Symbol][0].append(insightTarget)
                self.insightTargets[insight.Symbol][1].append(insight.CloseTimeUtc) # add if statement in case insight.CloseTimeUtc doesn't exist
                    
        #algorithm.Log('Updated Insight Targets')
        #algorithm.Log('length: '+str(len(self.insightTargets)))
        #for symbol in self.insightTargets:
            #algorithm.Log(str(symbol)+' '+str(self.insightTargets[symbol][0])+' '+str(self.insightTargets[symbol][1]))
                    
        predictedMarginUsed = 0
        #algorithm.Log('unbuffered Portfolio Targets')
        for portfolioSymbol in self.insightTargets:
            if len(self.insightTargets[portfolioSymbol][0]) == 0:
                unbufferedTargets[portfolioSymbol] = 0
            else:
                unbufferedTargets[portfolioSymbol] = max(self.insightTargets[portfolioSymbol][0])
            #algorithm.Log(str(portfolioSymbol)+' '+str(unbufferedTargets[portfolioSymbol]))
            predictedMarginUsed += unbufferedTargets[portfolioSymbol]*algorithm.Securities[portfolioSymbol].Price/2
                
        #algorithm.Log('Buffered Portfolio Targets')
        bufferedMarginLimit = (100 - self.marginCallPercentBuffer)/100*(algorithm.Portfolio.TotalMarginUsed + algorithm.Portfolio.MarginRemaining)        
        for unbufferedSymbol in unbufferedTargets:
            if predictedMarginUsed > bufferedMarginLimit: 
                bufferedPositionSize = bufferedMarginLimit/predictedMarginUsed*unbufferedTargets[unbufferedSymbol]
            else:
                bufferedPositionSize = unbufferedTargets[unbufferedSymbol]
            portfolioTargets.append(PortfolioTarget(unbufferedSymbol,bufferedPositionSize))
            #algorithm.Log(str(unbufferedSymbol)+' '+str(bufferedPositionSize))
                
        for key in [key for key in self.insightTargets if self.insightTargets[key] == [[],[]]]: del self.insightTargets[key]
        #algorithm.Log('Insight Targets after clearing empty symbols')
        #for symbol in self.insightTargets:
            #algorithm.Log(str(symbol)+' '+str(self.insightTargets[symbol][0])+' '+str(self.insightTargets[symbol][1]))
        
        #except:
            #algorithm.Debug('Portfolio model error') #for '+str(insight.Symbol))
        return portfolioTargets

    def RespectPortfolioBias(self, insight):
        '''Method that will determine if a given insight respects the portfolio bias
        Args:
            insight: The insight to create a target for
        '''
        return self.portfolioBias == PortfolioBias.LongShort or insight.Direction == self.portfolioBias

from AlgorithmImports import *

class RsiDivergenceAlphaModelJGG(AlphaModel):

    def __init__(self, period = 14, resolution = Resolution.Daily):
        self.period = period
        self.resolution = resolution
        self.symbolDataBySymbol = {}
        self.openWindows = {}
        self.highWindows = {}
        self.lowWindows = {}
        self.closeWindows = {}
        self.rsiWindows = {}
        resolutionString = Extensions.GetEnumString(resolution, Resolution)
        self.Name = '{}({},{})'.format(self.__class__.__name__, period, resolutionString)

    def Update(self, algorithm, data):
        self.rsiOversold = float(algorithm.GetParameter("rsiOversold"))
        self.rsiOverbought = float(algorithm.GetParameter("rsiOverbought"))
        self.minHiccup = float(algorithm.GetParameter("minHiccup"))
        self.periodLength = int(algorithm.GetParameter("periodLength"))
        self.minRsiRetrace = float(algorithm.GetParameter("minRsiRetrace"))
        self.minRetracementRatio = float(algorithm.GetParameter("minRetracementRatio"))
        self.rsiPivotDelta = float(algorithm.GetParameter("rsiPivotDelta"))
        insights = []
        for symbol, symbolData in self.symbolDataBySymbol.items():
            timeDateString = format(algorithm.Time)
            timeString = timeDateString.split()            
            rsi = symbolData.RSI
            if data.ContainsKey(symbol) and data[symbol] is not None: # and timeString[1] == "16:00:00":
                self.openWindows[symbol].Add(data[symbol].Open)
                self.highWindows[symbol].Add(data[symbol].High)
                self.lowWindows[symbol].Add(data[symbol].Low)
                self.closeWindows[symbol].Add(data[symbol].Close)
                self.rsiWindows[symbol].Add(rsi.Current.Value)
            if rsi.IsReady: #and timeString[1] == "16:00:00":
                try:
                    self.rsiHiccup = self.rsiWindows[symbol][0] - self.rsiWindows[symbol][1]
                    if self.rsiWindows[symbol][1] < self.rsiOversold and self.rsiHiccup >= self.minHiccup:
                        for periodIter in [self.periodLength]:     #move this line!!!!!
                            highList = list(self.highWindows[symbol])
                            lowList = list(self.lowWindows[symbol])
                            closeList = list(self.closeWindows[symbol])
                            rsiList = list(self.rsiWindows[symbol])
                            recentMax = max(highList[0:periodIter-1])
                            recentMin = min(lowList[0:periodIter-1])
                            Mag = (recentMax - recentMin)/recentMin
                            #
                            Pivot1RSI = min(rsiList[3:self.divergenceLookback-1])
                            Pivot1Index = rsiList.index(Pivot1RSI,3,self.divergenceLookback-1)
                            Pivot2Index = 1  
                            if Pivot1Index - Pivot2Index >= self.minLength:
                                Pivot1Price = closeList[Pivot1Index]
                                rsiRetrace = max(rsiList[2:Pivot1Index-1]) - Pivot1RSI
                                rsiRetraceIndex = rsiList.index(rsiRetrace+Pivot1RSI,2,Pivot1Index-1)
                                retracePrice = closeList[rsiRetraceIndex]
                                Pivot2RSI = rsiList[Pivot2Index]
                                Pivot2Price = closeList[Pivot2Index]
                                RetracementRatio = (retracePrice-Pivot1Price)/(retracePrice-Pivot2Price)
                                #if symbol == "BEN R735QTJ8XC9X":
                                    #algorithm.Debug(str(Pivot1RSI) + ' ' + str(Pivot1Index) + ' ' + str(Pivot1Price) + ' ' + str(rsiRetrace) + ' ' + str(Pivot2RSI) + ' ' + str(Pivot2Price) + ' ' + str(RetracementRatio))
                                    #algorithm.Debug(str(len(rsiList)))
                                    #algorithm.Debug(rsiList[46])
                                    #algorithm.Debug(rsiList[56])
                                    #algorithm.Debug(closeList[46])
                                    #algorithm.Debug(closeList[56])
                                    #for testIndex in range(len(rsiList)):
                                        #algorithm.Debug(rsiList[testIndex])
                                if Pivot1RSI < Pivot2RSI - self.rsiPivotDelta and rsiRetrace > self.minRsiRetrace and Pivot2Price < Pivot1Price and RetracementRatio > self.minRetracementRatio:
                                    self.insightPeriod = Time.Multiply(Extensions.ToTimeSpan(self.resolution), periodIter)
                                    insights.append(Insight.Price(symbol, self.insightPeriod, InsightDirection.Up,Mag,None,None,None))
                                highList = []
                                lowList = []
                                closeList = []
                                rsiList = []
                    #if self.rsiWindows[symbol][1] > self.rsiOverbought and self.rsiHiccup <= -1*self.minHiccup:
                     #   for periodIter in [self.periodLength]:
                      #      highList = list(self.highWindows[symbol])
                       #     lowList = list(self.lowWindows[symbol])
                        #    recentMax = max(highList[0:periodIter-1])
                         #   recentMin = min(lowList[0:periodIter-1])
                          #  Mag = (recentMax - recentMin)/recentMin 
                           # #Mag = (recentMin - recentMax)/recentMax
                            #self.insightPeriod = Time.Multiply(Extensions.ToTimeSpan(self.resolution), periodIter)
                            #insights.append(Insight.Price(symbol, self.insightPeriod, InsightDirection.Up,Mag,None,None,None))
                            ##insights.append(Insight.Price(symbol, self.insightPeriod, InsightDirection.Down,Mag,None,None,None))
                            #highList = []
                            #lowList = []    
                except:
                    algorithm.Log('Alpha model error for '+str(symbol))
        return insights

    def OnSecuritiesChanged(self, algorithm, changes):
        
        self.divergenceLookback = 90 #int(algorithm.GetParameter("divergenceLookback"))
        self.minLength = int(algorithm.GetParameter("minLength"))
        # clean up data for removed securities
        symbols = [ x.Symbol for x in changes.RemovedSecurities ]
        if len(symbols) > 0:
            for subscription in algorithm.SubscriptionManager.Subscriptions:
                if subscription.Symbol in symbols:
                    self.symbolDataBySymbol.pop(subscription.Symbol, None)
                    subscription.Consolidators.Clear()
                
        # initialize data for added securities
        
        addedSymbols = [ x.Symbol for x in changes.AddedSecurities if x.Symbol not in self.symbolDataBySymbol]
        if len(addedSymbols) == 0: return
        
        #self.windowLength = max(self.divergenceLookback,self.period)
        history = algorithm.History(addedSymbols, self.period + 20, self.resolution) #self.windowLength + 20, self.resolution)
        
        for symbol in addedSymbols:
            algorithm.Securities[symbol].FeeModel = ConstantFeeModel(0)
            #algorithm.Securities[symbol].SetSlippageModel(ConstantSlippageModel(0))
            rsi = algorithm.RSI(symbol, self.period, MovingAverageType.Wilders, self.resolution)
            self.rsiWindows[symbol] = RollingWindow[float](20+self.divergenceLookback) #self.windowLength)
            self.openWindows[symbol] = RollingWindow[float](20+self.divergenceLookback) #self.windowLength)
            self.highWindows[symbol] = RollingWindow[float](20+self.divergenceLookback) #self.windowLength)
            self.lowWindows[symbol] = RollingWindow[float](20+self.divergenceLookback) #self.windowLength)
            self.closeWindows[symbol] = RollingWindow[float](20+self.divergenceLookback) #self.windowLength)
            for tuple in history.loc[symbol].itertuples():
                self.openWindows[symbol].Add(tuple.open)
                self.highWindows[symbol].Add(tuple.high)
                self.lowWindows[symbol].Add(tuple.low)
                self.closeWindows[symbol].Add(tuple.close)
                rsi.Update(tuple.Index, tuple.close)
                if rsi.IsReady:
                    self.rsiWindows[symbol].Add(rsi.Current.Value)
            self.symbolDataBySymbol[symbol] = SymbolData(symbol, rsi)

class SymbolData:
    def __init__(self, symbol, rsi):
        self.Symbol = symbol
        self.RSI = rsi

from Execution.ImmediateExecutionModel import ImmediateExecutionModel
from Execution.NullExecutionModel import NullExecutionModel
from Portfolio.EqualWeightingPortfolioConstructionModel import EqualWeightingPortfolioConstructionModel
from Portfolio.NullPortfolioConstructionModel import NullPortfolioConstructionModel
from Risk.MaximumDrawdownPercentPerSecurity import MaximumDrawdownPercentPerSecurity
from Selection.QC500UniverseSelectionModel import QC500UniverseSelectionModel
from RsiDivergenceAlphaModelJGG import RsiDivergenceAlphaModelJGG
from PortfolioModelJGG import PortfolioModelJGG
from AlgorithmImports import *

class SimpleRSITestQC500Universe(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(1998,3,1) # Set Start Date
        self.SetEndDate(2014,12,31) # Set End Date
        self.SetCash(100000) # Set Strategy Cash
        self.SetBenchmark("SPY")
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage)
        self.SetExecution(ImmediateExecutionModel())
        self.SetPortfolioConstruction(PortfolioModelJGG(Time.Multiply(Extensions.ToTimeSpan(Resolution.Daily), 1)))
        self.SetRiskManagement(NullRiskManagementModel())
        #symbols = [ Symbol.Create("SPY", SecurityType.Equity, Market.USA), Symbol.Create("GE", SecurityType.Equity, Market.USA), Symbol.Create("BA", SecurityType.Equity, Market.USA) ]
        #self.SetUniverseSelection(ManualUniverseSelectionModel(symbols))
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.Universe.QC500)
        self.AddAlpha(RsiDivergenceAlphaModelJGG(resolution = Resolution.Daily))