# QUANTCONNECT.COM - Democratizing Finance, Empowering Individuals.https://www.quantconnect.com/project/5624489#code-tab-EqualWeightingClone_py
# Lean Algorithmic Trading Engine v2.0. Copyright 2014 QuantConnect Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from AlgorithmImports import *

class PortfolioModelJGG(PortfolioConstructionModel):
    '''Provides an implementation of IPortfolioConstructionModel that gives equal weighting to all securities.
    The target percent holdings of each security is 1/N where N is the number of securities.
    For insights of direction InsightDirection.Up, long targets are returned and
    for insights of direction InsightDirection.Down, short targets are returned.'''

    def __init__(self, rebalance = Resolution.Daily, portfolioBias = PortfolioBias.LongShort):
        '''Initialize a new instance of EqualWeightingPortfolioConstructionModel
        Args:
            rebalance: Rebalancing parameter. If it is a timedelta, date rules or Resolution, it will be converted into a function.
                              If None will be ignored.
                              The function returns the next expected rebalance time for a given algorithm UTC DateTime.
                              The function returns null if unknown, in which case the function will be called again in the
                              next loop. Returning current time will trigger rebalance.
            portfolioBias: Specifies the bias of the portfolio (Short, Long/Short, Long)'''
        self.portfolioBias = portfolioBias

        # 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, activeInsights):
        '''Will determine the target percent for each insight
        Args:
            activeInsights: The active insights to generate a target for'''
        targets = []
        try:
            for insight in activeInsights:
                dollarRisk = 1000
                volatilityRisk = insight.Magnitude/(1 + insight.Magnitude)
                biasMultiplier = insight.Direction if self.RespectPortfolioBias(insight) else InsightDirection.Flat
                insightTarget = PortfolioTarget(insight.Symbol,biasMultiplier*dollarRisk/volatilityRisk/insight.ReferenceValue)
                targets.append(insightTarget)
                #algorithm.Debug(str(insight.Symbol) + ',' + str(biasMultiplier) + ',' + str(volatilityRisk))
        except:
                algorithm.Debug(str(insight.Symbol))
        return targets

    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 RsiAlphaModelJGG(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"))
        insights = []
        for symbol, symbolData in self.symbolDataBySymbol.items():
            timeDateString = format(algorithm.Time)
            timeString = timeDateString.split()            
            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)
            rsi = symbolData.RSI
            if rsi.IsReady and timeString[1] == "09:31:00":
                #try:
                    self.rsiWindows[symbol].Add(rsi.Current.Value)
                    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]:
                            highList = list(self.highWindows[symbol])
                            lowList = list(self.lowWindows[symbol])
                            recentMax = max(highList[0:periodIter-1])
                            recentMin = min(lowList[0:periodIter-1])
                            #algorithm.Log(str(symbol) + ': ' + str(lowList[0]) + ' , ' + str(lowList[1]) + ' , ' + str(lowList[2]) + ' , ' + str(lowList[3]) + ' , ' + str(lowList[4]))
                            #algorithm.Log(str(symbol) + ': ' + str(highList[0]) + ' , ' + str(highList[1]) + ' , ' + str(highList[2]) + ' , ' + str(highList[3]) + ' , ' + str(highList[4]))
                            #algorithm.Log(str(symbol) + ': ' + str(recentMin) + ' , ' + str(recentMax))
                            stockPlot = Chart('Trade Plot')
                            stockPlot.AddSeries(Series('Low', SeriesType.Line, '$', Color.Red))
                            stockPlot.AddSeries(Series('High', SeriesType.Line, '$', Color.Green))
                            stockPlot.AddSeries(Series('Close', SeriesType.Line, '$', Color.Yellow))
                            algorithm.Plot('Trade Plot','Low', data.Bars["AAPL R735QTJ8XC9X"].Low)
                            algorithm.Plot('Trade Plot','High', data.Bars["AAPL R735QTJ8XC9X"].High)
                            algorithm.Plot('Trade Plot','Close', data.Bars["AAPL R735QTJ8XC9X"].Close)
                            Mag = (recentMax - recentMin)/recentMin
                            self.insightPeriod = Time.Multiply(Extensions.ToTimeSpan(self.resolution), periodIter)
                            insights.append(Insight.Price(symbol, self.insightPeriod, InsightDirection.Up,Mag,None,None,None))
                    #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 = (recentMin - recentMax)/recentMin
                           # self.insightPeriod = Time.Multiply(Extensions.ToTimeSpan(self.resolution), periodIter)
                            #insights.append(Insight.Price(symbol, self.insightPeriod, InsightDirection.Down,Mag,None,None,None))
                #except:
                    #algorithm.Log(str(symbol)+': '+ timeDateString)
        return insights

    def OnSecuritiesChanged(self, algorithm, changes):
        
        # 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
        
        history = algorithm.History(addedSymbols, self.period + 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.openWindows[symbol] = RollingWindow[float](self.period)
            self.highWindows[symbol] = RollingWindow[float](self.period)
            self.lowWindows[symbol] = RollingWindow[float](self.period)
            self.closeWindows[symbol] = RollingWindow[float](self.period)
            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 RsiAlphaModelJGG import RsiAlphaModelJGG
from PortfolioModelJGG import PortfolioModelJGG
from AlgorithmImports import *

class SimpleRSITestQC500Universe(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2021,3,1) # Set Start Date
        self.SetEndDate(2021,3,15) # Set End Date
        self.SetCash(1e6) # Set Strategy Cash
        self.SetBenchmark("SPY")
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage)
        self.SetExecution(ImmediateExecutionModel())
        self.SetPortfolioConstruction(PortfolioModelJGG(Time.Multiply(Extensions.ToTimeSpan(Resolution.Minute), 1)))
        #self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel(Time.Multiply(Extensions.ToTimeSpan(Resolution.Minute), 1)))
        #self.SetPortfolioConstruction(EqualWeightingCloneJGG(Time.Multiply(Extensions.ToTimeSpan(Resolution.Minute), 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.Minute
        self.AddUniverse(self.Universe.QC500)
        self.AddAlpha(RsiAlphaModelJGG())
        #self.rsiOversold = float(self.GetParameter("rsiOversold"))
        #self.rsiOverbought = float(self.GetParameter("rsiOverbought"))
        #self.periodLength = int(self.GetParameter("periodLength"))
        #self.Schedule.On(self.DateRules.EveryDay(), \
                 #self.TimeRules.At(16,00), \
                 #self.EveryDayBeforeMarketClose)
                 
    #def EveryDayBeforeMarketClose(self):
        #self.Log("EveryDay, 1 min before close: Fired at: {0}".format(self.Time))