from UniverseSelection import UniverseSelectionModel
# from alpha_model import MacdAlphaModel
# from Alphas.EmaCrossAlphaModel import EmaCrossAlphaModel
from Alphas.MacdAlphaModel import MacdAlphaModel
from Execution.StandardDeviationExecutionModel import ImmediateExecutionModel

class PrototypeV0(QCAlgorithm):

    def Initialize(self):
        
        # Set requested data resolution - minute required for risk management
        self.UniverseSettings.Resolution = Resolution.Minute
        self.UniverseSettings.DataNormalizationMode = DataNormalizationMode.Raw
        
        # Set Start Date
        self.SetStartDate(2017, 2, 15)
        
        # Set Strategy Cash
        self.SetCash(100000)  
        
        # self.SetBenchmark("TSLA")
        
        # Manual univere selection
        symbols = [ Symbol.Create("TSLA", SecurityType.Equity, Market.USA)]
        # symbols = [ Symbol.Create("TSLA", SecurityType.Equity, Market.USA),Symbol.Create("TVIX", SecurityType.Equity, Market.USA),Symbol.Create("ADPT", SecurityType.Equity, Market.USA),Symbol.Create("DDOG", SecurityType.Equity, Market.USA),Symbol.Create("APRN", SecurityType.Equity, Market.USA),Symbol.Create("ZM", SecurityType.Equity, Market.USA),Symbol.Create("MRNA", SecurityType.Equity, Market.USA) ]
        self.AddUniverseSelection(ManualUniverseSelectionModel(symbols))
        
        # count = 10
        # self.SetUniverseSelection(UniverseSelectionModel())

        # self.AddAlpha(EmaCrossAlphaModel(50, 200, Resolution.Daily))
        self.AddAlpha(MacdAlphaModel(12, 26, 9, MovingAverageType.Exponential, Resolution.Daily))

        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel() )
        
        self.SetExecution( ImmediateExecutionModel() )
        
        self.SetRiskManagement(TrailingStopRiskManagementModel(0.05))


    def OnData(self, data):
        '''OnData event is the primary entry point for your algorithm. Each new data point will be pumped in here.
            Arguments:
                data: Slice object keyed by symbol containing the stock data
        '''

        # if not self.Portfolio.Invested:
        #    self.SetHoldings("SPY", 1)

# universe selection above of stocks going above the SMA

from clr import AddReference
AddReference("System")
AddReference("QuantConnect.Common")
AddReference("QuantConnect.Indicators")
AddReference("QuantConnect.Algorithm.Framework")

from QuantConnect.Data.UniverseSelection import *
from QuantConnect.Indicators import SimpleMovingAverage
from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel

class UniverseSelectionModel(FundamentalUniverseSelectionModel):

    def __init__(self,
                 universeCount = 20,
                 universeSettings = None,
                 securityInitializer = None):

        super().__init__(False, universeSettings, securityInitializer)
        self.universeCount = universeCount
        # self.averages = {}
        # setup state storage in initialize method
        self.stateData = {}

    def SelectCoarse(self, algorithm, coarse):
        for c in coarse:
            if c.Symbol not in self.stateData:
                self.stateData[c.Symbol] = SelectionData(c.Symbol, 10)
            avg = self.stateData[c.Symbol]
            avg.update(c.EndTime, c.AdjustedPrice, c.DollarVolume)

        # filter the values of selectionData(sd) above SMA
        values = [sd for sd in self.stateData.values() if sd.volume > sd.sma.Current.Value and sd.volume_ratio > 0]
        
        # sort sd by the largest % jump in volume.
        values.sort(key=lambda sd: sd.volume_ratio, reverse=True)

        # return the top 10 symbol objects
        return [ sd.symbol for sd in values[:10] ]

# class used to improve readability of the coarse selection function
class SelectionData(object):
    def __init__(self, symbol, period):
        self.symbol = symbol
        self.volume = 0
        self.volume_ratio = 0
        self.sma = SimpleMovingAverage(period)

    def update(self, time, price, volume):
        self.volume = volume
        if self.sma.Update(time, volume):
            # get ratio of this volume bar vs previous 10 before it.
            self.volume_ratio = volume / self.sma.Current.Value

# 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 MacdAlphaModel(AlphaModel):
    '''Defines a custom alpha model that uses MACD crossovers. The MACD signal line
    is used to generate up/down insights if it's stronger than the bounce threshold.
    If the MACD signal is within the bounce threshold then a flat price insight is returned.'''

    def __init__(self,
                 fastPeriod = 12,
                 slowPeriod = 26,
                 signalPeriod = 9,
                 movingAverageType = MovingAverageType.Exponential,
                 resolution = Resolution.Daily):
        ''' Initializes a new instance of the MacdAlphaModel class
        Args:
            fastPeriod: The MACD fast period
            slowPeriod: The MACD slow period</param>
            signalPeriod: The smoothing period for the MACD signal
            movingAverageType: The type of moving average to use in the MACD'''
        self.fastPeriod = fastPeriod
        self.slowPeriod = slowPeriod
        self.signalPeriod = signalPeriod
        self.movingAverageType = movingAverageType
        self.resolution = resolution
        self.insightPeriod = Time.Multiply(Extensions.ToTimeSpan(resolution), fastPeriod)
        self.bounceThresholdPercent = 0.01
        self.symbolData = {}

        resolutionString = Extensions.GetEnumString(resolution, Resolution)
        movingAverageTypeString = Extensions.GetEnumString(movingAverageType, MovingAverageType)
        self.Name = '{}({},{},{},{},{})'.format(self.__class__.__name__, fastPeriod, slowPeriod, signalPeriod, movingAverageTypeString, resolutionString)


    def Update(self, algorithm, data):
        ''' Determines an insight for each security based on it's current MACD signal
        Args:
            algorithm: The algorithm instance
            data: The new data available
        Returns:
            The new insights generated'''
        insights = []

        for key, sd in self.symbolData.items():
            if sd.Security.Price == 0:
                continue

            direction = InsightDirection.Flat
            normalized_signal = sd.MACD.Signal.Current.Value / sd.Security.Price

            if normalized_signal > self.bounceThresholdPercent:
                direction = InsightDirection.Up
            elif normalized_signal < -self.bounceThresholdPercent:
                direction = InsightDirection.Down

            # ignore signal for same direction as previous signal
            if direction == sd.PreviousDirection:
                continue

            insight = Insight.Price(sd.Security.Symbol, self.insightPeriod, direction)
            sd.PreviousDirection = insight.Direction
            insights.append(insight)

        return insights


    def OnSecuritiesChanged(self, algorithm, changes):
        '''Event fired each time the we add/remove securities from the data feed.
        This initializes the MACD for each added security and cleans up the indicator for each removed security.
        Args:
            algorithm: The algorithm instance that experienced the change in securities
            changes: The security additions and removals from the algorithm'''
        for added in changes.AddedSecurities:
            self.symbolData[added.Symbol] = SymbolData(algorithm, added, self.fastPeriod, self.slowPeriod, self.signalPeriod, self.movingAverageType, self.resolution)

        for removed in changes.RemovedSecurities:
            data = self.symbolData.pop(removed.Symbol, None)
            if data is not None:
                # clean up our consolidator
                algorithm.SubscriptionManager.RemoveConsolidator(removed.Symbol, data.Consolidator)

class SymbolData:
    def __init__(self, algorithm, security, fastPeriod, slowPeriod, signalPeriod, movingAverageType, resolution):
        self.Security = security
        self.MACD = MovingAverageConvergenceDivergence(fastPeriod, slowPeriod, signalPeriod, movingAverageType)

        self.Consolidator = algorithm.ResolveConsolidator(security.Symbol, resolution)
        algorithm.RegisterIndicator(security.Symbol, self.MACD, self.Consolidator)

        self.PreviousDirection = None