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 ExponentialMovingAverage
from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel
from itertools import chain
from math import ceil
from QuantConnect import *
from QuantConnect.Indicators import *
from QuantConnect.Algorithm.Framework.Alphas import *
from QuantConnect.Algorithm.Framework.Alphas import InsightCollection
from QuantConnect.Algorithm.Framework.Portfolio import PortfolioConstructionModel, PortfolioTarget
from QuantConnect.Data.Market import Tick, TradeBar
from QuantConnect.Orders import *
from QuantConnect.Algorithm.Framework.Execution import ExecutionModel, OrderSizing
from QuantConnect.Algorithm.Framework.Portfolio import PortfolioTargetCollection
from QuantConnect.Algorithm.Framework.Risk import RiskManagementModel
import numpy as np
from datetime import datetime

class NullRiskManagementModel(RiskManagementModel):
    '''Provides an implementation of IRiskManagementModel that does nothing'''
    def ManageRisk(self, algorithm, targets):
        return []

class QC500UniverseSelectionModel(FundamentalUniverseSelectionModel):

    def __init__(self,
                 filterFineData = True,
                 universeSettings = None, 
                 securityInitializer = None):
        '''Initializes a new default instance of the QC500UniverseSelectionModel'''
        super().__init__(filterFineData, universeSettings, securityInitializer)
        self.NumberOfSymbolsCoarse = 1000
        self.NumberOfSymbolsFine = 500
        self.lastMonth = -1
        self.dollarVolumeBySymbol = {}
        self.symbols = []

    def SelectCoarse(self, algorithm, coarse):
        '''Performs coarse selection for the QC500 constituents.
        The stocks must have fundamental data
        The stock must have positive previous-day close price
        The stock must have positive volume on the previous trading day'''
        coarse = list(coarse)

        if len(coarse) == 0:
            return self.symbols

        month = coarse[0].EndTime.month
        if month == self.lastMonth:
            return self.symbols

        self.lastMonth = month

        # The stocks must have fundamental data
        # The stock must have positive previous-day close price
        # The stock must have positive volume on the previous trading day
        filtered = [x for x in coarse if x.HasFundamentalData
                                      and x.Volume > 0
                                      and x.Price > 0]
        # sort the stocks by dollar volume and take the top 1000
        top = sorted(filtered, key=lambda x: x.DollarVolume, reverse=True)[:self.NumberOfSymbolsCoarse]

        self.dollarVolumeBySymbol = { i.Symbol: i.DollarVolume for i in top }

        self.symbols = list(self.dollarVolumeBySymbol.keys())

        return self.symbols


    def SelectFine(self, algorithm, fine):
        '''Performs fine selection for the QC500 constituents
        The company's headquarter must in the U.S.
        The stock must be traded on either the NYSE or NASDAQ
        At least half a year since its initial public offering
        The stock's market cap must be greater than 500 million'''

        # The company's headquarter must in the U.S.
        # The stock must be traded on either the NYSE or NASDAQ
        # At least half a year since its initial public offering
        # The stock's market cap must be greater than 500 million
        filteredFine = [x for x in fine if x.CompanyReference.CountryId == "USA"
                                        and (x.CompanyReference.PrimaryExchangeID == "NYS" or x.CompanyReference.PrimaryExchangeID == "NAS")
                                        and (algorithm.Time - x.SecurityReference.IPODate).days > 180
                                        and x.EarningReports.BasicAverageShares.ThreeMonths * x.EarningReports.BasicEPS.TwelveMonths * x.ValuationRatios.PERatio > 5e8]
        count = len(filteredFine)
        if count == 0: return []

        myDict = dict()
        percent = float(self.NumberOfSymbolsFine / count)

        # select stocks with top dollar volume in every single sector
        for key in ["N", "M", "U", "T", "B", "I"]:
            value = [x for x in filteredFine if x.CompanyReference.IndustryTemplateCode == key]
            value = sorted(value, key=lambda x: self.dollarVolumeBySymbol[x.Symbol], reverse = True)
            myDict[key] = value[:ceil(len(value) * percent)]

        topFine = list(chain.from_iterable(myDict.values()))[:self.NumberOfSymbolsFine]
        self.symbols = [f.Symbol for f in topFine]

        return self.symbols
        
class EmaCrossAlphaModel(AlphaModel):
    '''Alpha model that uses an EMA cross to create insights'''

    def __init__(self,
                 fastPeriod = 12,
                 slowPeriod = 26,
                 resolution = Resolution.Daily):
        '''Initializes a new instance of the EmaCrossAlphaModel class
        Args:
            fastPeriod: The fast EMA period
            slowPeriod: The slow EMA period'''
        self.fastPeriod = fastPeriod
        self.slowPeriod = slowPeriod
        self.resolution = resolution
        self.predictionInterval = Time.Multiply(Extensions.ToTimeSpan(resolution), fastPeriod)
        self.symbolDataBySymbol = {}

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


    def Update(self, algorithm, data):
        '''Updates this alpha model with the latest data from the algorithm.
        This is called each time the algorithm receives data for subscribed securities
        Args:
            algorithm: The algorithm instance
            data: The new data available
        Returns:
            The new insights generated'''
        insights = []
        for symbol, symbolData in self.symbolDataBySymbol.items():
            if symbolData.Fast.IsReady and symbolData.Slow.IsReady:

                if symbolData.FastIsOverSlow:
                    if symbolData.Slow > symbolData.Fast:
                        insights.append(Insight.Price(symbolData.Symbol, self.predictionInterval, InsightDirection.Down))
                
                elif symbolData.SlowIsOverFast:
                    if symbolData.Fast > symbolData.Slow:
                        insights.append(Insight.Price(symbolData.Symbol, self.predictionInterval, InsightDirection.Up))

            symbolData.FastIsOverSlow = symbolData.Fast > symbolData.Slow

        return insights

    def OnSecuritiesChanged(self, algorithm, changes):
        '''Event fired each time the we add/remove securities from the data feed
        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:
            symbolData = self.symbolDataBySymbol.get(added.Symbol)
            if symbolData is None:
                # create fast/slow EMAs
                algorithm.Debug(added)
                symbolData = EMASymbolData(added)
                symbolData.Fast = algorithm.EMA(added.Symbol, self.fastPeriod)
                symbolData.Slow = algorithm.EMA(added.Symbol, self.slowPeriod)
                self.symbolDataBySymbol[added.Symbol] = symbolData
            else:
                # a security that was already initialized was re-added, reset the indicators
                symbolData.Fast.Reset()
                symbolData.Slow.Reset()


class EMASymbolData:
    '''Contains data specific to a symbol required by this model'''
    def __init__(self, security):
        self.Security = security
        self.Symbol = security.Symbol
        self.Fast = None
        self.Slow = None

        # True if the fast is above the slow, otherwise false.
        # This is used to prevent emitting the same signal repeatedly
        self.FastIsOverSlow = False

    @property
    def SlowIsOverFast(self):
        return not self.FastIsOverSlow
        
class EqualWeightingPortfolioConstructionModel(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):
        self.insightCollection = InsightCollection()
        self.removedSymbols = []

    def CreateTargets(self, algorithm, insights):
        '''Create portfolio targets from the specified insights
        Args:
            algorithm: The algorithm instance
            insights: The insights to create portoflio targets from
        Returns:
            An enumerable of portoflio targets to be sent to the execution model'''
        self.insightCollection.AddRange(insights)

        targets = []

        if self.removedSymbols is not None:
            # zero out securities removes from the universe
            for symbol in self.removedSymbols:
                targets.append(PortfolioTarget(symbol, 0))
                self.removedSymbols = None

        if len(insights) == 0:
            return targets

        # Get symbols that have emit insights and still in the universe
        symbols = list(set([x.Symbol for x in self.insightCollection if x.CloseTimeUtc > algorithm.UtcTime]))

        # give equal weighting to each security
        percent = 1.0 / len(symbols)
        for symbol in symbols:
            activeInsights = [ x for x in self.insightCollection if x.Symbol == symbol ]
            direction = activeInsights[-1].Direction
            targets.append(PortfolioTarget.Percent(algorithm, symbol, direction * percent))

        return targets

    def OnSecuritiesChanged(self, algorithm, changes):
        '''Event fired each time the we add/remove securities from the data feed
        Args:
            algorithm: The algorithm instance that experienced the change in securities
            changes: The security additions and removals from the algorithm'''

        # save securities removed so we can zero out our holdings
        self.removedSymbols = [x.Symbol for x in changes.RemovedSecurities]

        # remove the insights of the removed symbol from the collection
        for removedSymbol in self.removedSymbols:
            if self.insightCollection.ContainsKey(removedSymbol):
                for insight in self.insightCollection[removedSymbol]:
                    self.insightCollection.Remove(insight)
    
class VolumeWeightedAveragePriceExecutionModel(ExecutionModel):
    '''Execution model that submits orders while the current market price is more favorable that the current volume weighted average price.'''

    def __init__(self):
        '''Initializes a new instance of the VolumeWeightedAveragePriceExecutionModel class'''
        self.targetsCollection = PortfolioTargetCollection()
        self.symbolData = {}

        # Gets or sets the maximum order quantity as a percentage of the current bar's volume.
        # This defaults to 0.01m = 1%. For example, if the current bar's volume is 100,
        # then the maximum order size would equal 1 share.
        self.MaximumOrderQuantityPercentVolume = 0.01


    def Execute(self, algorithm, targets):
        '''Executes market orders if the standard deviation of price is more
       than the configured number of deviations in the favorable direction.
       Args:
           algorithm: The algorithm instance
           targets: The portfolio targets'''

        # update the complete set of portfolio targets with the new targets
        self.targetsCollection.AddRange(targets)

        for target in self.targetsCollection.OrderByMarginImpact(algorithm):
            symbol = target.Symbol

            # calculate remaining quantity to be ordered
            unorderedQuantity = OrderSizing.GetUnorderedQuantity(algorithm, target)

            # fetch our symbol data containing our VWAP indicator
            data = self.symbolData.get(symbol, None)
            if data is None: return

            # ensure we're receiving price data before submitting orders
            if data.Security.Price == 0: return

            # check order entry conditions
            if self.PriceIsFavorable(data, unorderedQuantity):
                # get the maximum order size based on total order value
                maxOrderSize = OrderSizing.PercentVolume(data.Security, self.MaximumOrderQuantityPercentVolume)
                orderSize = np.min([maxOrderSize, np.abs(unorderedQuantity)])

                # round down to even lot size
                orderSize -= orderSize % data.Security.SymbolProperties.LotSize
                if orderSize != 0:
                    algorithm.MarketOrder(symbol, np.sign(unorderedQuantity) * orderSize)

            # check to see if we're done with this target
            unorderedQuantity = OrderSizing.GetUnorderedQuantity(algorithm, target)
            if unorderedQuantity == 0:
                self.targetsCollection.Remove(target.Symbol)


    def OnSecuritiesChanged(self, algorithm, changes):
        '''Event fired each time the we add/remove securities from the data feed
        Args:
            algorithm: The algorithm instance that experienced the change in securities
            changes: The security additions and removals from the algorithm'''
        for removed in changes.RemovedSecurities:
            # clean up removed security data
            if removed.Symbol in self.symbolData:
                if self.IsSafeToRemove(algorithm, removed.Symbol):
                    data = self.symbolData.pop(removed.Symbol)
                    algorithm.SubscriptionManager.RemoveConsolidator(removed.Symbol, data.Consolidator)

        for added in changes.AddedSecurities:
            if added.Symbol not in self.symbolData:
                self.symbolData[added.Symbol] = SymbolData(algorithm, added)


    def PriceIsFavorable(self, data, unorderedQuantity):
        '''Determines if the current price is more than the configured
       number of standard deviations away from the mean in the favorable direction.'''
        if unorderedQuantity > 0:
            if data.Security.BidPrice < data.VWAP:
                return True
        else:
            if data.Security.AskPrice > data.VWAP:
                return True

        return False

    def IsSafeToRemove(self, algorithm, symbol):
        '''Determines if it's safe to remove the associated symbol data'''
        # confirm the security isn't currently a member of any universe
        return not any([kvp.Value.ContainsMember(symbol) for kvp in algorithm.UniverseManager])

class SymbolData:
    def __init__(self, algorithm, security):
        self.Security = security
        self.Consolidator = algorithm.ResolveConsolidator(security.Symbol, security.Resolution)
        name = algorithm.CreateIndicatorName(security.Symbol, "VWAP", security.Resolution)
        self.vwap = IntradayVwap(name)
        algorithm.RegisterIndicator(security.Symbol, self.vwap, self.Consolidator)

    @property
    def VWAP(self):
       return self.vwap.Value

class IntradayVwap:
    '''Defines the canonical intraday VWAP indicator'''
    def __init__(self, name):
        self.Name = name
        self.Value = 0.0
        self.lastDate = datetime.min
        self.sumOfVolume = 0.0
        self.sumOfPriceTimesVolume = 0.0

    @property
    def IsReady(self):
        return self.sumOfVolume > 0.0

    def Update(self, input):
        '''Computes the new VWAP'''
        success, volume, averagePrice = self.GetVolumeAndAveragePrice(input)
        if not success:
            return

        # reset vwap on daily boundaries
        if self.lastDate != input.EndTime.date():
            self.sumOfVolume = 0.0
            self.sumOfPriceTimesVolume = 0.0
            self.lastDate = input.EndTime.date()

        # running totals for Σ PiVi / Σ Vi
        self.sumOfVolume += volume
        self.sumOfPriceTimesVolume += averagePrice * volume

        if self.sumOfVolume == 0.0:
           # if we have no trade volume then use the current price as VWAP
           self.Value = input.Value
           return

        self.Value = self.sumOfPriceTimesVolume / self.sumOfVolume


    def GetVolumeAndAveragePrice(self, input):
        '''Determines the volume and price to be used for the current input in the VWAP computation'''

        if type(input) is Tick:
            if input.TickType == TickType.Trade:
                return True, float(input.Quantity), float(input.LastPrice)

        if type(input) is TradeBar:
            if not input.IsFillForward:
                averagePrice = float(input.High + input.Low + input.Close) / 3
                return True, float(input.Volume), averagePrice

        return False, 0.0, 0.0

class BasicTemplateFrameworkAlgorithm(QCAlgorithmFramework):
    '''Basic template framework algorithm uses framework components to define the algorithm.'''

    def Initialize(self):
        ''' Initialise the data and resolution required, as well as the cash and start-end dates for your algorithm. All algorithms must initialized.'''

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

        self.SetStartDate(2013,10,7)   #Set Start Date
        self.SetEndDate(2013,10,12)    #Set End Date
        self.SetCash(100000)           #Set Strategy Cash

        # set algorithm framework models
        
        self.Debug("Universe")
        self.SetUniverseSelection(QC500UniverseSelectionModel())
        self.Debug("Alpha")
        self.SetAlpha(EmaCrossAlphaModel())
        self.Debug("Portfolio")
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
        self.Debug("Execution")
        self.SetExecution(VolumeWeightedAveragePriceExecutionModel())
        self.Debug("Risk")
        self.SetRiskManagement(NullRiskManagementModel())

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