from Alphas.EmaCrossAlphaModel import EmaCrossAlphaModel

from Execution.ImmediateExecutionModel import ImmediateExecutionModel

from Risk.NullRiskManagementModel import NullRiskManagementModel

from Selection.QC500UniverseSelectionModel import QC500UniverseSelectionModel

from SlopeBasedEquityMomentumAlphaModel import SlopeBasedEquityMomentumAlphaModel

class BasicTemplateFrameworkAlgorithm(QCAlgorithmFramework):

    def Initialize(self):

        # Set requested data resolution
        self.UniverseSettings.Resolution = Resolution.Daily

        self.SetStartDate(2018, 9, 4)   #Set Start Date
        self.SetEndDate(2019, 3, 7)    #Set End Date
        self.SetCash(100000)           #Set Strategy Cash

        # selection will run on mon/tues/thurs at 00:00/06:00/12:00/18:00
        self.SetUniverseSelection(QC500UniverseSelectionModel())
        
        self.SetAlpha(SlopeBasedEquityMomentumAlphaModel())
        
        self.SetPortfolioConstruction(NullPortfolioConstructionModel()) 
        
        self.SetExecution(ImmediateExecutionModel())
        
        self.SetRiskManagement(NullRiskManagementModel())
        
    

    def OnOrderEvent(self, orderEvent):
        if orderEvent.Status == OrderStatus.Filled:
            # self.Debug("Purchased Stock: {0}".format(orderEvent.Symbol))
            pass

# QUANTCONNECT.COM - Democratizing Finance, Empowering Individuals.
# 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 clr import AddReference
AddReference("QuantConnect.Common")
AddReference("QuantConnect.Algorithm")
AddReference("QuantConnect.Algorithm.Framework")
AddReference("QuantConnect.Indicators")

from QuantConnect import *
from QuantConnect.Indicators import *
from QuantConnect.Algorithm import *
from QuantConnect.Algorithm.Framework import *
from QuantConnect.Algorithm.Framework.Alphas import *


class SlopeBasedEquityMomentumAlphaModel(AlphaModel):

    def __init__(self,
                 shortTermMomentumWindow = 60,
                 longTermMomentumWindow = 90,
                 minimumMomentum = 60,
                 indexAverageWindow = 100,
                 resolution = Resolution.Daily):

        self.shortTermMomentumWindow = shortTermMomentumWindow
        self.longTermMomentumWindow = longTermMomentumWindow
        self.minimumMomentum = minimumMomentum
        self.indexAverageWindow = indexAverageWindow
        self.resolution = resolution
        self.symbolData = {}
        resolutionString = Extensions.GetEnumString(resolution, Resolution)
        self.symbolDataBySymbol = {}
        self.month = None
        
    def slope(ts):

        x = np.arange(len(ts))
        log_ts = np.log(ts)
        slope, intercept, r_value, p_value, std_err = stats.linregress(x, log_ts)
        annualized_slope = (np.power(np.exp(slope), 250) - 1) * 100

        return annualized_slope * (r_value ** 2)
    
    def Update(self, algorithm, data):
        insights = []
        
        ## If it has already run Update this month, then return nothing
        ## You can replace month with week if you want weekly
        if algorithm.Time.month == self.month:
            return []
        algorithm.Log('Update() called: ' + str(algorithm.Time))
        ## Update self.month so that it won't do anything in Update until a month has gone by
        self.month = algorithm.Time.month

        
        ## Local dictionary of rolling windows
        shortTermBars = {}
        longTermBars = {}
        ## Add current bar to rolling window
        
        for symbol, symbolData in self.symbolDataBySymbol.items():
            if not data.Bars.ContainsKey(symbol):
                ## If there is no new bar, i.e. dividend being paid, then forward fill rolling window
                if symbolData.LongTermBars.Samples > 0:
                    symbolData.UpdateRollingWindows(symbolData.LongTermBars[0])
                elif symbolData.ShortTermBars.Samples > 0:
                    symbolData.UpdateRollingWindows(symbolData.ShortTermBars[0])
                else:
                    continue

            else:
                symbolData.UpdateRollingWindows(data.Bars[symbol])
        
            shortTermBars[symbol] = symbolData.ShortTermBars
            longTermBars[symbol] = symbolData.LongTermBars
            
        ## Insight code here

        
        return insights
        
    
    def OnSecuritiesChanged(self, algorithm, changes):
        
        for security in changes.AddedSecurities:
            if security.Symbol not in self.symbolDataBySymbol.keys():
                symbolData = SymbolData(security.Symbol, self.shortTermMomentumWindow, self.longTermMomentumWindow)
                self.symbolDataBySymbol[security.Symbol] = symbolData
    
        for security in changes.RemovedSecurities:
            self.symbolDataBySymbol.pop(security.Symbol)
            
class SymbolData:
    
    def __init__(self, symbol, shortTermMomentumWindow, longTermMomentumWindow):
        self.Symbol = symbol
        self.shortTermRollingWindow = RollingWindow[TradeBar](shortTermMomentumWindow)
        self.longTermRollingWindow = RollingWindow[TradeBar](longTermMomentumWindow)
        
    def UpdateRollingWindows(self, bar):
        self.shortTermRollingWindow.Add(bar)
        self.longTermRollingWindow.Add(bar)

    @property
    def ShortTermBars(self):
        return self.shortTermRollingWindow
    
    @property
    def LongTermBars(self):
        return self.longTermRollingWindow