#region imports
from AlgorithmImports import *
#endregion
#region imports
from QuantConnect import Resolution, Extensions
from QuantConnect.Algorithm.Framework.Alphas import *
from QuantConnect.Algorithm.Framework.Portfolio import *
from itertools import groupby
from datetime import datetime, timedelta
from pytz import utc
UTCMIN = datetime.min.replace(tzinfo=utc)
#endregion
class EqualWeightingPortfolio(PortfolioConstructionModel):


    def __init__(self, initialAllocationPerSecurity = 0.1):
        

        # portfolio exposure per security (as a % of total equity)
        self.initialAllocationPerSecurity = initialAllocationPerSecurity
        
        self.insightCollection = InsightCollection()
        self.removedSymbols = []
        
        self.nextRebalance = None


    def CreateTargets(self, algorithm, insights):

        targets = []

          
        if len(insights) == 0 and len([x.Symbol.Value for x in algorithm.Portfolio.Values if x.Invested]) != 0:
            return targets

        # here we get the new insights and add them to our insight collection
        for insight in insights:
            self.insightCollection.Add(insight)
            
        # create flatten target for each security that was removed from the universe
        if len(self.removedSymbols) > 0:
            universeDeselectionTargets = [ PortfolioTarget(symbol, 0) for symbol in self.removedSymbols ]
            targets.extend(universeDeselectionTargets)
            algorithm.Log('(Portfolio module) liquidating: ' + str([x.Value for x in self.removedSymbols]) + ' if they are active, due to not being in the universe')
            self.removedSymbols = []

        expiredInsights = self.insightCollection.RemoveExpiredInsights(algorithm.UtcTime)

        expiredTargetsLog = []
        expiredTargets = []
        for symbol, f in groupby(expiredInsights, lambda x: x.Symbol):
            if not self.insightCollection.HasActiveInsights(symbol, algorithm.UtcTime):
                expiredTargets.append(PortfolioTarget(symbol, 0))
                expiredTargetsLog.append(symbol)
                continue
        
        algorithm.Log(f'(Portfolio module) sold {expiredTargetsLog} due to insight being expired')
        targets.extend(expiredTargets)

        # get insight that have not expired of each symbol that is still in the universe
        activeInsights = self.insightCollection.GetActiveInsights(algorithm.UtcTime)

        # get the last generated active insight for each symbol
        lastActiveInsights = []
        for symbol, g in groupby(activeInsights, lambda x: x.Symbol):
            lastActiveInsights.append(sorted(g, key = lambda x: x.GeneratedTimeUtc)[-1])
        
        # determine target percent for the given insights
        boughtTargetsLog = []
        for insight in lastActiveInsights:
            allocationPercent = self.initialAllocationPerSecurity * insight.Direction
            target = PortfolioTarget.Percent(algorithm, insight.Symbol, allocationPercent)
            boughtTargetsLog.append(insight.Symbol)
            targets.append(target)
        
        algorithm.Log(f'(Portfolio module) Bought {boughtTargetsLog} stocks, that expires at {Expiry.EndOfMonth}')

        return targets

        
    def OnSecuritiesChanged(self, algorithm, changes):
        
        
        newRemovedSymbols = [x.Symbol for x in changes.RemovedSecurities if x.Symbol not in self.removedSymbols]
        
        # get removed symbol and invalidate them in the insight collection
        self.removedSymbols.extend(newRemovedSymbols)
        self.insightCollection.Clear(self.removedSymbols)
            
        removedList = [x.Value for x in self.removedSymbols]
        algorithm.Log('(Portfolio module) securities removed from Universe: ' + str(removedList))

#region imports
from AlgorithmImports import *
#endregion


class MarketOrderModel(ExecutionModel):
    '''Provides an implementation of IExecutionModel that immediately submits market orders to achieve the desired portfolio targets'''

    def __init__(self):
        '''Initializes a new instance of the ImmediateExecutionModel class'''
        self.targetsCollection = PortfolioTargetCollection()

    def Execute(self, algorithm, targets):

        # for performance we check count value, OrderByMarginImpact and ClearFulfilled are expensive to call
        self.targetsCollection.AddRange(targets)
        if self.targetsCollection.Count > 0:
            for target in self.targetsCollection.OrderByMarginImpact(algorithm):
                security = algorithm.Securities[target.Symbol]
                # calculate remaining quantity to be ordered
                quantity = OrderSizing.GetUnorderedQuantity(algorithm, target, security)
                if quantity != 0:
                    aboveMinimumPortfolio = BuyingPowerModelExtensions.AboveMinimumOrderMarginPortfolioPercentage(security.BuyingPowerModel, security, quantity, algorithm.Portfolio, algorithm.Settings.MinimumOrderMarginPortfolioPercentage)
                    if aboveMinimumPortfolio:
                        algorithm.MarketOrder(security, quantity)

                    if quantity > 0:
                        algorithm.Plot('EMA value and SPY value', 'Buy', 1)
                    if quantity < 0:
                        algorithm.Plot('EMA value and SPY value', 'Sell', 1)

            self.targetsCollection.ClearFulfilled(algorithm)

#region imports
from AlgorithmImports import *
#endregion
class MomentumAlphaModel(AlphaModel):
    def __init__(self, lookback, resolution):
        self.lookback = lookback
        self.resolution = resolution
        self.predictionInterval = Expiry.EndOfMonth
        self.symbolDataBySymbol = {}
        
        self.num_insights = 10
        self.lastMonth = -1

    def Update(self, algorithm, data):
        
        for symbol, symbolData in self.symbolDataBySymbol.items():
            if not algorithm.IsMarketOpen(symbol):
                return []

        if algorithm.Time.month == self.lastMonth:    
            return []
        self.lastMonth = algorithm.Time.month

        insights = []

        for symbol, symbolData in self.symbolDataBySymbol.items():

            if symbolData.CanEmit:
                direction = InsightDirection.Up
                magnitude = symbolData.Return

                insights.append(Insight.Price(symbol, self.predictionInterval, direction, magnitude = magnitude))

        insights1 = sorted([x for x in insights], key = lambda x: x.Magnitude, reverse=True)

        algorithm.Log(f'(Alpha module) Sent {len([x for x in insights1[:self.num_insights]])} insights to the pcm module')

        return [x for x in insights1[:self.num_insights]]

    def OnSecuritiesChanged(self, algorithm, changes):
        
        # clean up data for removed securities
        for removed in changes.RemovedSecurities:
            symbolData = self.symbolDataBySymbol.pop(removed.Symbol, None)
            if symbolData is not None:
                symbolData.RemoveConsolidators(algorithm)

        # initialize data for added securities
        symbols = [ x.Symbol for x in changes.AddedSecurities ]
        history = algorithm.History(symbols, self.lookback, self.resolution)
        if history.empty: return

        tickers = history.index.levels[0]
        for ticker in tickers:
            symbol = SymbolCache.GetSymbol(ticker)
            
            if symbol == "SPY":
                return

            if symbol not in self.symbolDataBySymbol:
                symbolData = SymbolData(symbol, self.lookback)
                self.symbolDataBySymbol[symbol] = symbolData
                symbolData.RegisterIndicators(algorithm, self.resolution)
                symbolData.WarmUpIndicators(history.loc[ticker])


class SymbolData:

    def __init__(self, symbol, lookback):
        
        self.Symbol = symbol
        self.ROC = RateOfChange('{}.ROC({})'.format(symbol, lookback), lookback)
        self.Consolidator = None
        self.previous = 0

    def RegisterIndicators(self, algorithm, resolution):
        self.Consolidator = algorithm.ResolveConsolidator(self.Symbol, resolution)
        algorithm.RegisterIndicator(self.Symbol, self.ROC, self.Consolidator)

    def RemoveConsolidators(self, algorithm):
        if self.Consolidator is not None:
            algorithm.SubscriptionManager.RemoveConsolidator(self.Symbol, self.Consolidator)

    def WarmUpIndicators(self, history):
        for tuple in history.itertuples():
            self.ROC.Update(tuple.Index, tuple.close)

    @property
    def Return(self):
        return float(self.ROC.Current.Value)

    @property
    def CanEmit(self):
        if self.previous == self.ROC.Samples:
            return False

        self.previous = self.ROC.Samples
        return self.ROC.IsReady

    def __str__(self, **kwargs):
        return '{}: {:.2%}'.format(self.ROC.Name, (1 + self.Return)**252 - 1)

#region imports
from AlgorithmImports import *
from datetime import timedelta, time, datetime
#endregion
class RiskModelWithSpy(RiskManagementModel):
    
    def __init__(self, algorithm, spy, lookback,  resolution):
        self.spy = spy
        
        self.lookback = lookback
        self.resolution = resolution
        
        self.symboldata = {}

        self.lastDay = -1
        
        #Flag so we only instanciate it once
        self.symboldata[self.spy.Symbol] = EMASymbolData(algorithm, self.spy, self.lookback, self.resolution)
        
    def ManageRisk(self, algorithm, targets):
        
        targets = []
        
        for symbol, symboldata in self.symboldata.items():
            #logic. If price is below the current value for EMA, we send a portfoliotarget of 0
            spyValue = self.spy.Price
            AlmaValue = symboldata.EMA.Current.Value
            
            for kvp in algorithm.Securities:
                security = kvp.Value
                
                if spyValue <= AlmaValue:
                    targets.append(PortfolioTarget(security.Symbol, 0))

        
        if len(targets) >= 1:
            algorithm.Log('(Risk module) sold the entire portfolio due to SPY being below EMA')
        
        self.PlotCharts(algorithm, symboldata.EMA.Current.Value)

        return targets



    def PlotCharts(self, algorithm, EMAValue):
        if algorithm.Time.day != self.lastDay:
            algorithm.Plot('EMA value and SPY value', 'EMA value', EMAValue)
            algorithm.Plot('EMA value and SPY value', 'SPY value', self.spy.Price)
            self.lastDay = algorithm.Time.day

            
     
class EMASymbolData:
    
    def __init__(self, algorithm, security, lookback, resolution):
        symbol = security.Symbol
        self.Security = symbol
        self.Consolidator = algorithm.ResolveConsolidator(symbol, resolution)
        
        smaName = algorithm.CreateIndicatorName(symbol, f"SMA{lookback}", resolution)
        self.EMA = ExponentialMovingAverage(smaName, lookback)
        algorithm.RegisterIndicator(symbol, self.EMA, self.Consolidator)
        
        history = algorithm.History(symbol, lookback, resolution)
        if 'close' in history:
            history = history.close.unstack(0).squeeze()
            for time, value in history.iteritems():
                self.EMA.Update(time, value)

from AlgorithmImports import *
from datetime import timedelta, time, datetime
from MomentumAlphaModel import MomentumAlphaModel
from EqualWeightingPortfolio import EqualWeightingPortfolio
from RiskModelWithSpy import RiskModelWithSpy
from MarketOrderExecution import MarketOrderModel
from System.Drawing import Color


class MomentumFrameworkAlgo(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2019, 1, 1)
        self.SetEndDate(2020, 1, 1)
        self.SetCash(100000)  # Set Strategy Cash
        self.UniverseSettings.Resolution = Resolution.Hour
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)

        seeder = FuncSecuritySeeder(self.GetLastKnownPrices)
        self.SetSecurityInitializer(lambda security: seeder.SeedSecurity(security))

        self.SetWarmup(timedelta(10))
        self.SetBenchmark('SPY')
        
        self.spy = self.AddEquity('SPY', Resolution.Hour)
        
        self.AddUniverse(self.CoarseUniverse)
        self.AddAlpha(MomentumAlphaModel(lookback=203, resolution=Resolution.Daily)) 
        pcm = EqualWeightingPortfolio()
        self.SetPortfolioConstruction(pcm)
        self.SetExecution(MarketOrderModel())
        self.AddRiskManagement(RiskModelWithSpy(self, self.spy, 200, Resolution.Daily))
        
        self.num_coarse = 45
        self.lastMonth = -1


        EMAplot = Chart('EMA value and SPY value')
        EMAplot.AddSeries(Series('EMA value', SeriesType.Line, '$', Color.Blue))
        EMAplot.AddSeries(Series('SPY value', SeriesType.Line, '$', Color.White))
        EMAplot.AddSeries(Series('Buy', SeriesType.Scatter, '$', Color.Green, ScatterMarkerSymbol.Triangle))
        EMAplot.AddSeries(Series('Sell', SeriesType.Scatter, '$', Color.Red, ScatterMarkerSymbol.TriangleDown))
        self.AddChart(EMAplot)


    def CoarseUniverse(self, coarse):
        if self.Time.month == self.lastMonth: 
            return Universe.Unchanged
        self.lastMonth = self.Time.month
        
        selected = sorted([x for x in coarse if x.Price > 10 and x.Price < 5000 and x.HasFundamentalData], key = lambda x: x.DollarVolume, reverse=True)
        
        self.Log(f'(Universe module)Sent {len([x.Symbol for x in selected[:self.num_coarse]])} symbols to the alpha module')

        return [x.Symbol for x in selected[:self.num_coarse]]

    def OnEndOfDay(self):
        self.Plot("Positions", "Num", len([x.Symbol for x in self.Portfolio.Values if self.Portfolio[x.Symbol].Invested]))
        self.Plot(f"Margin", "Used", self.Portfolio.TotalMarginUsed)
        self.Plot(f"Margin", "Remaning", self.Portfolio.MarginRemaining)
        self.Plot(f"Cash", "Remaining", self.Portfolio.Cash)