#self.SetUniverseSelection( ManualUniverseSelectionModel(symbl) )

# 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")
AddReference("QuantConnect.Algorithm")
AddReference("QuantConnect.Common")
AddReference("QuantConnect.Indicators")

from System import *
from QuantConnect import *
from QuantConnect.Algorithm import *
from QuantConnect.Indicators import *
from QuantConnect.Algorithm.Framework import *
from QuantConnect.Algorithm.Framework.Risk import *
from QuantConnect.Algorithm.Framework.Alphas import *
from QuantConnect.Orders.Fees import ConstantFeeModel
from QuantConnect.Algorithm.Framework.Selection import *
from QuantConnect.Algorithm.Framework.Execution import *
from QuantConnect.Algorithm.Framework.Portfolio import *
import numpy as np
import math


class PriceGapMeanReversionAlpha(QCAlgorithm):
    '''The motivating idea for this Alpha Model is that a large price gap (here we use true outliers --
    price gaps that whose absolutely values are greater than 3 * Volatility) is due to rebound
    back to an appropriate price or at least retreat from its brief extreme. Using a Coarse Universe selection
    function, the algorithm selects the top x-companies by Dollar Volume (x can be any number you choose)
    to trade with, and then uses the Standard Deviation of the 100 most-recent closing prices to determine
    which price movements are outliers that warrant emitting insights.

    This alpha is part of the Benchmark Alpha Series created by QuantConnect which are open
    sourced so the community and client funds can see an example of an alpha.'''

    def Initialize(self):

        self.SetStartDate(2019, 8, 31)  # Set Start Date
        self.SetEndDate(2019, 10, 30)  # Set Start Date
        self.SetCash(100000)           #Set Strategy Cash
        #self.SetWarmup(15)
        
        stockPlot = Chart("Supertrend Plot")
        # On the Trade Plotter Chart we want 3 series: trades and price:
        stockPlot.AddSeries(Series("Open", SeriesType.Line, 0))
        stockPlot.AddSeries(Series("High", SeriesType.Line, 0))
        stockPlot.AddSeries(Series("Low", SeriesType.Line, 0))
        stockPlot.AddSeries(Series("Close", SeriesType.Line, 0))
        stockPlot.AddSeries(Series("Ha Close", SeriesType.Line, 0))
        stockPlot.AddSeries(Series("Ha Open", SeriesType.Line, 0))
        stockPlot.AddSeries(Series("Ha High", SeriesType.Line, 0))
        stockPlot.AddSeries(Series("Ha Low", SeriesType.Line, 0))
        stockPlot.AddSeries(Series("Hahl2", SeriesType.Line, 0))
        stockPlot.AddSeries(Series("TmpUp", SeriesType.Line, 0))
        stockPlot.AddSeries(Series("TmpDn", SeriesType.Line, 0))
        stockPlot.AddSeries(Series("FinalUp", SeriesType.Line, 0))
        stockPlot.AddSeries(Series("FinalDn", SeriesType.Line, 0))
        stockPlot.AddSeries(Series("ATR", SeriesType.Line, 0))
        stockPlot.AddSeries(Series("TrendDir", SeriesType.Line, 0))
        
        self.AddChart(stockPlot)

        
    
        ## Initialize variables to be used in controlling frequency of universe selection
        self.week = -1
        symbl = [Symbol.Create(x, SecurityType.Equity, Market.USA) for x in ['IBM','AAPL']]
        #symbl = self.AddEquity("IBM", Resolution.Daily) ## Subscribe to hourly TradeBars
        

        ## Manual Universe Selection
        self.UniverseSettings.Resolution = Resolution.Daily
        #self.SetUniverseSelection( ManualUniverseSelectionModel(symbl) )
        self.AddUniverseSelection(FineFundamentalUniverseSelectionModel(self.SelectCoarse, self.SelectFine))
        ## Set trading fees to $0
        self.SetSecurityInitializer(lambda security: security.SetFeeModel(ConstantFeeModel(0)))

        ## Set custom Alpha Model
        self.SetAlpha(PriceGapMeanReversionAlphaModel())

        ## Set equal-weighting Portfolio Construction Model
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())

        ## Set Execution Model
        self.SetExecution(ImmediateExecutionModel())

        ## Set Risk Management Model
        self.SetRiskManagement(NullRiskManagementModel())


    def SelectCoarse(self, coarse):
        tickers = ["AAPL", "AIG", "IBM"]
        return [Symbol.Create(x, SecurityType.Equity, Market.USA) for x in tickers]
    
    def SelectFine(self, fine):
        return [f.Symbol for f in fine] 

class PriceGapMeanReversionAlphaModel:

    def __init__(self, *args, **kwargs):
        ''' Initialize variables and dictionary for Symbol Data to support algorithm's function '''
        self.lookback = 15
        self.window = RollingWindow[float](10)
        self.trendWindow = RollingWindow[int](10)
        self.ST_Period = 15
        self.ST_Coeff = 1
        self.Name = 'TosAlphaModel'
        self.resolution = kwargs['resolution'] if 'resolution' in kwargs else Resolution.Daily
        self.prediction_interval = Time.Multiply(Extensions.ToTimeSpan(self.resolution), 5) ## Arbitrary
        self.symbolDataBySymbol = {}
        self.symbolArr = []
        self.symbolName = ''
        
    def Update(self, algorithm, data):
        insights = []
        #algorithm.Log(self.symbolDataBySymbol.items())
        
        ## Loop through all Symbol Data objects
        for symbol, symbolData in self.symbolDataBySymbol.items():
            ## Evaluate whether or not the price jump is expected to rebound
            if not symbolData.IsTrend(data):
                continue

            ## Emit insights accordingly to the price jump sign
            #direction = flat 
            direction = InsightDirection.Down if symbolData.trendDir1 > 0 else InsightDirection.Up
            insights.append(Insight.Price(symbol, self.prediction_interval, direction, None))
            #algorithm.Log(" Symbol:" + str(symbol) + " Open:" + str(data.Bars[symbol].Open))
            algorithm.Plot("Supertrend Plot", "Open", str(symbolData.o))
            algorithm.Plot("Supertrend Plot", "High", str(symbolData.h))
            algorithm.Plot("Supertrend Plot", "Low", str(symbolData.l))
            algorithm.Plot("Supertrend Plot", "Close", str(symbolData.c))
            algorithm.Plot("Supertrend Plot", "Ha Close", str(symbolData.ha_close))
            algorithm.Plot("Supertrend Plot", "Ha Open", str(symbolData.ha_open))
            algorithm.Plot("Supertrend Plot", "Ha High", str(symbolData.ha_high))
            algorithm.Plot("Supertrend Plot", "Ha Low", str(symbolData.ha_low))
            algorithm.Plot("Supertrend Plot", "Hahl2", str(symbolData.hahl2_1))
            algorithm.Plot("Supertrend Plot", "TmpUp", str(symbolData.tmpUp1))
            algorithm.Plot("Supertrend Plot", "TmpDn", str(symbolData.tmpDn1))
            algorithm.Plot("Supertrend Plot", "FinalUp", str(symbolData.finalUp1))
            algorithm.Plot("Supertrend Plot", "FinalDn", str(symbolData.finalDn1))
            algorithm.Plot("Supertrend Plot", "ATR", str(symbolData.atr))
            algorithm.Plot("Supertrend Plot", "TrendDir", str(symbolData.trendDir1))
            
        return 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)

        symbols = [x.Symbol for x in changes.AddedSecurities
            if x.Symbol not in self.symbolDataBySymbol]

        history = algorithm.History(symbols, self.lookback, self.resolution)
        if history.empty: return

        ## Create and initialize SymbolData objects
        for symbol in symbols:
            symbolData = SymbolData(algorithm, symbol, self.lookback, self.resolution, 'TOS', self.ST_Period,self.ST_Coeff,self.window,self.trendWindow)
            symbolData.WarmUpIndicators(history.loc[symbol])
            self.symbolDataBySymbol[symbol] = symbolData
    
          
class SymbolData:
    def __init__(self, algorithm, symbol, lookback, resolution,name, period, coeff, window, trendWindow):
        self.symbol = symbol
        self.close = 0
        self.PriceJump = 0
        self.Name = name
        self.IsReady = False
        self.o = self.h = self.l = self.c = self.ha_close = self.ha_open = self.ha_high = 0
        self.pre_close = 0
        self.pre_ha_close = 0
        self.ha_low = self.hahl2_1 = self.trendDir1 = 0 
        self.trendDir = trendWindow
        self.last_price = window

        self._open = window
        self._close = window
        self.tmpUp = self.tmpDn = self.finalUp = self.finalDn = window
        self.tmpUp1 = self.tmpDn1 = self.finalUp1 = self.finalDn1 = 0
        self.super_trend = 0
        self.atr = AverageTrueRange(period, MovingAverageType.Wilders) 
        self.coeff = coeff
        
        self.consolidator = algorithm.ResolveConsolidator(symbol, resolution)
        self.volatility = StandardDeviation(f'{symbol}.STD({lookback})', lookback)
        algorithm.RegisterIndicator(symbol, self.volatility, self.consolidator)

    def RemoveConsolidators(self, algorithm):
        algorithm.SubscriptionManager.RemoveConsolidator(self.symbol, self.consolidator)

    def WarmUpIndicators(self, history):
        self.close = history.iloc[-1].close
        for tuple in history.itertuples():
            self.volatility.Update(tuple.Index, tuple.close)

    def IsTrend(self, data):

        ## Check for any data events that would return a NoneBar in the Alpha Model Update() method
        if not data.Bars.ContainsKey(self.symbol):
            return False

        self.atr.Update(data.Bars[self.symbol])
        if self.atr.IsReady: 
            self.o  =  data.Bars[self.symbol].Open         ## Open price
            self.h  =  data.Bars[self.symbol].High         ## High price
            self.l  =  data.Bars[self.symbol].Low          ## Low price
            self.c  = last_price = data.Bars[self.symbol].Close        ## Close price
            
            self.last_price.Add(last_price)
            
            #check last_price is ready or not 
            if not self.last_price.IsReady:
               self.pre_close =  0
            else:
               self.pre_close =  self.last_price[1]

            _close = (self.o + self.h + self.l + self.c) / 4.0
            self._close.Add(_close)

            #check _close is ready or not 
            if not self._close.IsReady:
               self.IsReady = False
               return
            else:
                 self.ha_close =  self._close[1]
                 
                 #self.pre_ha_close =  self._close[0]


            _open = (self.o + self.c) / 2.0
            self._open.Add(_open)
            #check _close is ready or not 
 
            if not self._open.IsReady:
               self.ha_open = (self._open[0] + self._close[0]) / 2.0
            else:
               self.ha_open = (self._open[1] + self._close[1]) / 2.0

            
            self.ha_high = Math.Max(self.o, Math.Max(self.ha_open, self.ha_close))
            self.ha_low = Math.Min(self.l, Math.Min(self.ha_open, self.ha_close))
            self.hahl2_1 = (self.ha_high  + self.ha_low) / 2
        
            tmpUp = self.hahl2_1 - (self.atr.Current.Value * self.coeff)
            self.tmpUp.Add(tmpUp)
            self.tmpUp1 = self.tmpUp[0]
            tmpDn = self.hahl2_1 + (self.atr.Current.Value * self.coeff)
            self.tmpDn.Add(tmpDn)
            self.tmpDn1 = self.tmpDn[0]

            # register final up and final down in rolling window

            self.finalUp.Add(tmpUp)
            self.finalUp1 = self.finalUp[0] 
            self.finalDn.Add(tmpDn)
            self.finalDn1 = self.finalDn[0] 
            
            if self.pre_close > self.finalUp[1]:
                self.finalUp[0] = max(self.tmpUp[0], self.finalUp[1])
            else:
                self.finalUp[0] = self.tmpUp[0]
               
            if self.pre_close < self.finalDn[1]:
                self.finalDn[0] = min(self.tmpDn[0], self.finalDn[1])
            else:
                self.finalDn[0] = self.tmpDn[0]

           
            # calculate trendDir
            if data.Bars[self.symbol].Close > self.finalDn[1]:
                trendDir = 1
            elif data.Bars[self.symbol].Close < self.finalUp[1]:
                trendDir = -1

            # adding trendDir in rolling window
            self.trendDir.Add(trendDir)
            if not self.trendDir.IsReady:
               self.IsReady = False
               return
            else:
                 self.trendDir1 =  self.trendDir[0]
                
            
            if not math.isnan(self.trendDir1):
                self.trendDir[0] = self.trendDir1
            else:
                self.trendDir[0] = 1

            self.trendDir1 = self.trendDir[0]
            
            return self.trendDir1