class CoarseFineFundamentalATRComboAlgorithm(QCAlgorithm):

    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.'''

        self.SetStartDate(2014, 1, 1)  #Set Start Date
        self.SetEndDate(  2014, 2, 1)    #Set End Date
        self.SetCash(50000)            #Set Strategy Cash

        # what resolution should the data *added* to the universe be?
        self.UniverseSettings.Resolution = Resolution.Daily
        
        # An indicator(or any rolling window) needs data(updates) to have a value
        self.atr_window = 10
        self.UniverseSettings.MinimumTimeInUniverse = self.atr_window
        self.SetWarmUp(self.atr_window)

        # this add universe method accepts two parameters:
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
        
        # Set dictionary of indicators
        self.indicators = {}

        self.__numberOfSymbols     = 100
        self.__numberOfSymbolsFine = 10
        
        
    def OnData(self, data):
    
        for symbol in self.universe:
            
            # is symbol iin Slice object? (do we even have data on this step for this asset)
            if not data.ContainsKey(symbol):
                continue
            
            # new symbol? setup indicator object. Then update
            if symbol not in self.indicators:
                self.indicators[symbol] = SymbolData(symbol, self, self.atr_window)
            # update by bar
            #self.indicators[symbol].update_bar(data[symbol])
            #update by value
            self.indicators[symbol].update_value(self.Time, data[symbol].Price)
            
            if self.IsWarmingUp: continue
            
            self.Log(str(symbol) + " : " + str(self.indicators[symbol].get_atr()))
            #self.Log("SYMBOL : ".format(symbol.Price))
            self.Log("PRICE : ".format(self.Securities[symbol].Price))
            
            # now you can use logic to trade, random example:
            atr = self.indicators[symbol].get_atr()
            if atr != 0.0: # maybe a new symbol gets added and isnt ready yet?
                if atr >= 3.0:
                    self.SetHoldings(symbol, -0.1)
                else:
                    self.Liquidate(symbol)
            

    # sort the data by daily dollar volume and take the top 'NumberOfSymbols'
    def CoarseSelectionFunction(self, coarse):
        # sort descending by daily dollar volume
        sortedByDollarVolume = sorted(coarse, key=lambda x: x.DollarVolume, reverse=True)

        # return the symbol objects of the top entries from our sorted collection
        return [ x.Symbol for x in sortedByDollarVolume[:self.__numberOfSymbols] ]


    # sort the data by P/E ratio and take the top 'NumberOfSymbolsFine'
    def FineSelectionFunction(self, fine):
        
        # sort descending by P/E ratio
        sortedByPeRatio = sorted(fine, key=lambda x: x.OperationRatios.OperationMargin.Value, reverse=False)
        
        # resulting symbols
        self.universe = [ x.Symbol for x in sortedByPeRatio[:self.__numberOfSymbolsFine] ]

        # take the top entries from our sorted collection
        return self.universe


    # this event fires whenever we have changes to our universe
    def OnSecuritiesChanged(self, changes):
        
        # liquidate removed securities
        for security in changes.RemovedSecurities:
            if security.Invested:
                self.Liquidate(security.Symbol)
                
                # clean up
                del self.indicators[security.Symbol]



class SymbolData(object):
    def __init__(self, symbol, context, window):
        self.symbol = symbol
        """
        I had to pass ATR from outside object to get it to work, could pass context and use any indica
        var atr = ATR(Symbol symbol, int period, MovingAverageType type = null, Resolution resolution = null, Func`2[Data.IBaseData,Data.Market.IBaseDataBar] selector = null)
        """
        self.window    = window
        self.indicator = context.EMA(symbol, self.window)
        #self.indicator = context.BB(symbol, self.window)
        self.atr       = 0.0

    """
    Runtime Error: Python.Runtime.PythonException: NotSupportedException : AverageTrueRange does not support Update(DateTime, decimal) method overload. Use Update(IBaseDataBar) instead.
    """
    def update_bar(self, bar):
        self.indicator.Update(bar)
        
    def update_value(self, time, value):
        self.indicator.Update(time, value)
            
    def get_atr(self):
        return self.indicator.Current.Value
        #return self.indicator.LowerBand.Current.Value, self.indicator.UpperBand.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("System.Core")
AddReference("System.Collections")
AddReference("QuantConnect.Common")
AddReference("QuantConnect.Algorithm")

from System import *
from System.Collections.Generic import List
from QuantConnect import *
from QuantConnect.Algorithm import QCAlgorithm
from QuantConnect.Data.UniverseSelection import *

### <summary>
### Demonstration of using coarse and fine universe selection together to filter down a smaller universe of stocks.
### </summary>
### <meta name="tag" content="using data" />
### <meta name="tag" content="universes" />
### <meta name="tag" content="coarse universes" />
### <meta name="tag" content="fine universes" />
class CoarseFineFundamentalComboAlgorithm(QCAlgorithm):

    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.'''

        self.SetStartDate(2014,01,06)  #Set Start Date
        self.SetEndDate(2014,01,07)    #Set End Date
        self.SetCash(50000)            #Set Strategy Cash

        # what resolution should the data *added* to the universe be?
        self.UniverseSettings.Resolution = Resolution.Daily

        # this add universe method accepts two parameters:
        # - coarse selection function: accepts an IEnumerable<CoarseFundamental> and returns an IEnumerable<Symbol>
        # - fine selection function: accepts an IEnumerable<FineFundamental> and returns an IEnumerable<Symbol>
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
        
        # Set dictionary of indicators
        self.indicator = {}

        self.__numberOfSymbols = 100
        self.__numberOfSymbolsFine = 5
        self._changes = SecurityChanges.None


    # sort the data by daily dollar volume and take the top 'NumberOfSymbols'
    def CoarseSelectionFunction(self, coarse):
        # sort descending by daily dollar volume
        sortedByDollarVolume = sorted(coarse, key=lambda x: x.DollarVolume, reverse=True)

        # return the symbol objects of the top entries from our sorted collection
        return [ x.Symbol for x in sortedByDollarVolume[:self.__numberOfSymbols] ]

    # sort the data by P/E ratio and take the top 'NumberOfSymbolsFine'
    def FineSelectionFunction(self, fine):
        
        # sort descending by P/E ratio
        sortedByPeRatio = sorted(fine, key=lambda x: x.OperationRatios.OperationMargin.Value, reverse=False)
        
        # Here we want to get our inititialized indicator
        # We are going to use a dictionary to refer the object that will keep the moving averages
        for cf in fine:
            if cf.Symbol not in self.indicator:
                self.indicator[cf.Symbol] = SymbolData(cf.Symbol)

            # Updates the SymbolData object with current EOD price
            avg = self.indicator[cf.Symbol]
            avg.update(cf.EndTime, cf.Price)

        # take the top entries from our sorted collection
        return [ x.Symbol for x in sortedByPeRatio[:self.__numberOfSymbolsFine] ]

    def OnData(self, data):
        
        # liquidate removed securities
        for security in self._changes.RemovedSecurities:
            if security.Invested:
                self.Liquidate(security.Symbol)
        
        # Set dictionary of indicators
        #self.indicator = {}
        
        self.Log("SECS : ".format(self._changes.AddedSecurities))
        # Create indicator & check Price
        for security in self._changes.AddedSecurities:
            self.indicator[security.Symbol] = self.ATR(security.Symbol, 5, Resolution.Daily)
            
            #self.Log("SECURITY : ".format(self.Securities[security.Symbol]))
            self.Log("SECURITY : ".format(security.Symbol))
            #self.Log("ATR : ".format(self.indicator[security.Symbol].AverageTrueRange.Current.Value))
            self.Log("PRICE : ".format(self.Securities[security.Symbol].Price))
            


    #def OnSecuritiesChanged(self, changes):
        #self._changes = changes
        #self._changes = SecurityChanges.None;


    # this event fires whenever we have changes to our universe
    def OnSecuritiesChanged(self, changes):
        self._changes = changes
        
class SymbolData(object):
    def __init__(self, symbol):
        self.symbol = symbol
        self.indicator = ExponentialMovingAverage(100)
        #self.indicator = AverageTrueRange(5)
        #self.indicator = BollingerBands(5)
        self.scale = 0

    def update(self, time, value):
        if self.indicator.Update(time, value):
            indicator = self.indicator.Current.Value