from alphaEMA import EmaCross
from Alphas.EmaCrossAlphaModel import EmaCrossAlphaModel
from Execution.ImmediateExecutionModel import ImmediateExecutionModel
from Portfolio.EqualWeightingPortfolioConstructionModel import EqualWeightingPortfolioConstructionModel
from Selection.QC500UniverseSelectionModel import QC500UniverseSelectionModel

class FormalYellowGreenBee(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2021, 1, 2)  # Set Start Date
        self.SetEndDate(2021, 3, 15)  # Set Start Date
        self.SetCash(100000)  # Set Strategy Cash
        # self.AddEquity("SPY", Resolution.Minute)
        self.AddAlpha(EmaCross(50, 200, Resolution.Hour))

        self.SetExecution(ImmediateExecutionModel())

        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel(lambda time: None))

        self.UniverseSettings.Resolution = Resolution.Hour
        symbols = [ Symbol.Create("TSLA", SecurityType.Equity, Market.USA) , Symbol.Create("HES", SecurityType.Equity, Market.USA) ]
        self.SetUniverseSelection( ManualUniverseSelectionModel(symbols) )
        self.SetSecurityInitializer(lambda x: x.SetDataNormalizationMode(DataNormalizationMode.Adjusted))




    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)

# 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 EmaCross(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 and data.Bars.ContainsKey(symbolData.Symbol):
                if data[symbolData.Symbol] is not None:        
                    if algorithm.Portfolio[symbolData.Symbol].IsLong and symbolData.bull_insight == True:
                        if data[symbolData.Symbol].Low < symbolData.bull_SL_entry and symbolData.bull_insight_exp != data.Time: 
                            symbolData.bull_insight = False
                            insights.append(Insight.Price(symbolData.Symbol,symbolData.bull_insight_exp - data.Time , InsightDirection.Flat, 0, 0, None))
                            algorithm.Debug(  str(algorithm.Time) + "Bull Exit  " + str(symbolData.Symbol) + "   hit stoploss:  " + str(symbolData.bull_SL_entry) + " low: " +str(data[symbolData.Symbol].Low ))
                        elif data[symbolData.Symbol].Close > symbolData.bull_PT_entry and symbolData.bull_insight_exp != data.Time:
                            symbolData.bull_insight = False
                            insights.append(Insight.Price(symbolData.Symbol,symbolData.bull_insight_exp - data.Time , InsightDirection.Flat, 0, 0, None))
                            algorithm.Debug(  str(algorithm.Time) + "Bull Exit  " + str(symbolData.Symbol) + "   hit price_target:  " + str(symbolData.bull_PT_entry) + " close: " +str(data[symbolData.Symbol].Close))
                        elif symbolData.bull_insight_exp <= data.Time:
                            symbolData.bull_insight = False
                            algorithm.Debug(  str(algorithm.Time) + "Bull Exit  " + str(symbolData.Symbol) + "   insight expired:  " + str(symbolData.bull_PT_entry) + " close: " +str(data[symbolData.Symbol].Close))
    
                    if algorithm.Portfolio[symbolData.Symbol].IsShort and symbolData.bear_insight == True:
                        if data[symbolData.Symbol].High > symbolData.bear_SL_entry and symbolData.bear_insight_exp != data.Time:
                            symbolData.bear_insight = False
                            insights.append(Insight.Price(symbolData.Symbol,symbolData.bear_insight_exp - data.Time, InsightDirection.Flat, 0, 0, None))
                            algorithm.Debug(  str(algorithm.Time) + "Bear Exit " + str(symbolData.Symbol) + "   hit stoploss:  " + str(symbolData.bear_SL_entry) + " high: " +str(data[symbolData.Symbol].High ))
                        elif data[symbolData.Symbol].Close < symbolData.bear_PT_entry and symbolData.bear_insight_exp != algorithm.Time:
                            symbolData.bear_insight = False
                            insights.append(Insight.Price(symbolData.Symbol,symbolData.bear_insight_exp - data.Time, InsightDirection.Flat, 0, 0, None))
                            algorithm.Debug(  str(algorithm.Time) + "Bear Exit  " + str(symbolData.Symbol) + "   hit price_target:  " + str(symbolData.bear_PT_entry) + " close: " +str(data[symbolData.Symbol].Close))
                        elif symbolData.bear_insight_exp <= algorithm.Time:
                            symbolData.bear_insight = False
                            algorithm.Debug(  str(algorithm.Time) + "Bear Exit  " + str(symbolData.Symbol) + "   insight expired:  " + str(symbolData.bear_PT_entry) + " close: " +str(data[symbolData.Symbol].Close))



                if symbolData.FastIsOverSlow:
                    if symbolData.Slow > symbolData.Fast:
                        close = data[symbolData.Symbol].Close
                        symbolData.bear_insight_exp = data.Time + self.predictionInterval
                        symbolData.bear_PT_entry = .93*close
                        symbolData.bear_SL_entry = 1.03*close
                        symbolData.bear_insight = True
                        
                        insights.append(Insight.Price(symbolData.Symbol, self.predictionInterval, InsightDirection.Down))

                elif symbolData.SlowIsOverFast:
                    if symbolData.Fast > symbolData.Slow:
                        close = data[symbolData.Symbol].Close
                        symbolData.bull_insight_exp = data.Time + self.predictionInterval
                        symbolData.bull_PT_entry = 1.08*close
                        symbolData.bull_SL_entry = .97*close
                        symbolData.bull_insight = True
                        
                        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
                symbolData = SymbolData(added)

                symbolData.Fast = algorithm.EMA(added.Symbol, self.fastPeriod, self.resolution)
                symbolData.Slow = algorithm.EMA(added.Symbol, self.slowPeriod, self.resolution)
                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 SymbolData:
    '''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
        self.bull_PT_entry = None
        self.bear_PT_entry = None
        self.bull_SL_entry = None
        self.bear_SL_entry = None
        self.bear_insight_exp  =None
        self.bull_insight_exp  =None
        self.bull_insight = False
        self.bear_insight = False

        # 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