# 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, insightList):
        '''Will determine the target percent for each insight
        Args:
            activeInsights: The active insights to generate a target for'''
        targets = []
        for insight in insightList:
            dollarRisk = 100
            insightTarget = PortfolioTarget(insight.Symbol,(insight.Direction if self.RespectPortfolioBias(insight) else InsightDirection.Flat)*dollarRisk/insight.Magnitude)
            targets.append(insightTarget)
        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

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):
        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.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.blip = self.rsiWindows[symbol][0] - self.rsiWindows[symbol][1]
                    volatility=(max(self.closeWindows[symbol])-min(self.closeWindows[symbol]))/(max(self.closeWindows[symbol])+min(self.closeWindows[symbol]))*2/self.period
                    if self.rsiWindows[symbol][1] < 30 and self.blip >= 5:
                        for periodIter in [5]:
                            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
                            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] > 70 and self.blip <= -5:
                        for periodIter in [5]:
                            closeList = list(self.closeWindows[symbol])
                            recentMax = max(closeList[0:periodIter-1])
                            recentMin = min(closeList[0:periodIter-1])
                            Mag = -1*(recentMax - recentMin)/recentMax
                            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.close)
                self.highWindows[symbol].Add(tuple.close)
                self.lowWindows[symbol].Add(tuple.close)
                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(1998,3,1) # Set Start Date
        self.SetEndDate(2014,12,31) # Set End Date
        self.SetCash(10000000000) # 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.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))