import pandas as pd
import numpy as np
from AlgorithmImports import *
# endregion
from QuantConnect.Indicators import ExponentialMovingAverage

class QEMAAlphaModel(AlphaModel):

    openingBar = None
    currentBar = None
    """
    This class emits insights to hold a long (short) position after the chikou line of a
    security's Ichimoku Cloud crosses over (under) the top (bottom) of the cloud.
    """
    def __init__(self, vfastPeriod = 5, fastPeriod = 9, midPeriod = 21, slowPeriod = 55, resolution = Resolution.Minute, portfolioBias = PortfolioBias.Long):

        '''Initializes a new instance of the EmaCrossAlphaModel class
        Args:
            fastPeriod: The fast EMA period
            slowPeriod: The slow EMA period'''
        self.vfastPeriod = vfastPeriod
        self.fastPeriod = fastPeriod
        self.midPeriod = midPeriod
        self.slowPeriod = slowPeriod
        self.resolution = resolution
    
        #consFiveMin = TradeBarConsolidator(5)
        #consFiveMin.DataConsolidated += self.OnDataConsolidated
        #self.SubscriptionManager.AddConsolidator(symbol, consFiveMin)
        
        self.predictionInterval = Time.Multiply(Extensions.ToTimeSpan(resolution), slowPeriod)
        self.symbolDataBySymbol = {}

        resolutionString = Extensions.GetEnumString(resolution, Resolution)
        self.Name = '{}({},{},{})'.format(self.__class__.__name__, vfastPeriod, fastPeriod, midPeriod, 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.vFast.IsReady and symbolData.Slow.IsReady:

                if symbolData.vFast == symbolData.Fast:
                    insights.append(Insight.Price(symbolData.Symbol, self.predictionInterval, InsightDirection.Down))

                if (symbolData.Fast > symbolData.Mid and symbolData.Mid > symbolData.Slow):
                    insights.append(Insight.Price(symbolData.Symbol, self.predictionInterval, InsightDirection.Up))


        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:
                symbolData = SymbolData(added, self.vfastPeriod, self.fastPeriod, self.midPeriod, self.slowPeriod, algorithm, self.resolution)
                self.symbolDataBySymbol[added.Symbol] = symbolData
            else:
                # a security that was already initialized was re-added, reset the indicators
                symbolData.vFast.Reset()
                symbolData.Fast.Reset()
                symbolData.Mid.Reset()
                symbolData.Slow.Reset()

        for removed in changes.RemovedSecurities:
            data = self.symbolDataBySymbol.pop(removed.Symbol, None)
            if data is not None:
                # clean up our consolidators
                data.RemoveConsolidators()

    def OnDataConsolidated(self, bar):
        self.currentBar = bar


class SymbolData:
    '''Contains data specific to a symbol required by this model'''
    def __init__(self, security, vfastPeriod, fastPeriod, midPeriod, slowPeriod, algorithm, resolution):
        self.Security = security
        self.Symbol = security.Symbol
        self.algorithm = algorithm
        self.Resolution = resolution
        resolution = timedelta(minutes=5)

        self.vFastConsolidator = algorithm.ResolveConsolidator(security.Symbol, resolution)
        self.FastConsolidator = algorithm.ResolveConsolidator(security.Symbol, resolution)
        self.MidConsolidator = algorithm.ResolveConsolidator(security.Symbol, resolution)
        self.SlowConsolidator = algorithm.ResolveConsolidator(security.Symbol, resolution)

        algorithm.SubscriptionManager.AddConsolidator(security.Symbol, self.vFastConsolidator)
        algorithm.SubscriptionManager.AddConsolidator(security.Symbol, self.FastConsolidator)
        algorithm.SubscriptionManager.AddConsolidator(security.Symbol, self.MidConsolidator)
        algorithm.SubscriptionManager.AddConsolidator(security.Symbol, self.SlowConsolidator)

        # create fast/slow EMAs
        self.vFast = ExponentialMovingAverage(security.Symbol, vfastPeriod, ExponentialMovingAverage.SmoothingFactorDefault(vfastPeriod))
        self.Fast = ExponentialMovingAverage(security.Symbol, fastPeriod, ExponentialMovingAverage.SmoothingFactorDefault(fastPeriod))
        self.Mid = ExponentialMovingAverage(security.Symbol, midPeriod, ExponentialMovingAverage.SmoothingFactorDefault(midPeriod))
        self.Slow = ExponentialMovingAverage(security.Symbol, slowPeriod, ExponentialMovingAverage.SmoothingFactorDefault(slowPeriod))

        algorithm.RegisterIndicator(security.Symbol, self.vFast, self.vFastConsolidator);
        algorithm.RegisterIndicator(security.Symbol, self.Fast, self.FastConsolidator);
        algorithm.RegisterIndicator(security.Symbol, self.Mid, self.MidConsolidator);
        algorithm.RegisterIndicator(security.Symbol, self.Slow, self.SlowConsolidator);

        algorithm.WarmUpIndicator(security.Symbol, self.vFast, resolution);
        algorithm.WarmUpIndicator(security.Symbol, self.Fast, resolution);
        algorithm.WarmUpIndicator(security.Symbol, self.Mid, resolution);
        algorithm.WarmUpIndicator(security.Symbol, self.Slow, resolution);

        

        # True if the fast is above the slow, otherwise false.
        # This is used to prevent emitting the same signal repeatedly
        self.FastIsOverSlow = True
        
    def RemoveConsolidators(self):
        self.algorithm.SubscriptionManager.RemoveConsolidator(self.Security.Symbol, self.vFastConsolidator)
        self.algorithm.SubscriptionManager.RemoveConsolidator(self.Security.Symbol, self.FastConsolidator)
        self.algorithm.SubscriptionManager.RemoveConsolidator(self.Security.Symbol, self.MidConsolidator)
        self.algorithm.SubscriptionManager.RemoveConsolidator(self.Security.Symbol, self.SlowConsolidator)

    @property
    def SlowIsOverFast(self):
        return not self.FastIsOverSlow

from AlgorithmImports import *

from clr import AddReference

AddReference("System")
AddReference("QuantConnect.Algorithm")
AddReference("QuantConnect.Common")

from System import *
from QuantConnect import *
from QuantConnect.Algorithm import *

from alphamodel import QEMAAlphaModel
from risk import MaxDrawdownPercentModel

class QEMA_AlgorithmUniverse(QCAlgorithm):

    def Initialize(self):
        self.SetTimeZone(TimeZones.NewYork)
        self.SetStartDate(2021, 9, 20)
        self.SetEndDate(2022, 10, 20)
        self.SetCash(1000)
        
        self.AddUniverse(self.SelectCoarse)

        self.SetAlpha(QEMAAlphaModel(portfolioBias = PortfolioBias.Long))
        self.SetWarmup(55)
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel(portfolioBias = PortfolioBias.Long))
        self.SetRiskManagement(MaxDrawdownPercentModel())
        self.SetExecution(ImmediateExecutionModel())
        
        self.averages = {}

        self.SetWarmup(200)


        
    def SelectCoarse(self, universe):
        """
        Coarse universe selection is called each day at midnight.
        Returns the symbols that have fundamental data.
        """

        selected = []
        universe = sorted(universe, key=lambda c: c.Volume, reverse=True)
        universe = [c for c in universe if c.Price < 50][:100]

        # Create a loop to use all the coarse data:
        for coarse in universe:
            symbol = coarse.Symbol

            if symbol not in self.averages:
                history = self.History(symbol, 200, Resolution.Daily)
                self.averages[symbol] = SelectionData(history)
                
            self.averages[symbol].update(self.Time, coarse.AdjustedPrice)

            # Check if EMA50 > EMA200 and if so, append the symbol to the selected list
            if self.averages[symbol].unifast > self.averages[symbol].unislow:
                if self.averages[symbol].is_ready():
                    selected.append(symbol)

        return selected[:20]


class SelectionData():

    def __init__(self, history):
        self.unislow = SimpleMovingAverage(200)
        self.unifast = SimpleMovingAverage(50)
        

        for bar in history.itertuples():
            self.unifast.Update(bar.Index[1], bar.close)
            self.unislow.Update(bar.Index[1], bar.close)

    def is_ready(self):
        return self.unislow.IsReady

    def update(self, time, price):
        self.unifast.Update(time, price)
        self.unislow.Update(time, price)



    

#region imports
from AlgorithmImports import *

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

class MaxDrawdownPercentModel(RiskManagementModel):
    '''Provides an implementation of IRiskManagementModel that limits the drawdown of the portfolio to the specified percentage.'''

    def __init__(self, maximumDrawdownPercent = 0.025, isTrailing = False):
        '''Initializes a new instance of the MaximumDrawdownPercentPortfolio class
        Args:
            maximumDrawdownPercent: The maximum percentage drawdown allowed for algorithm portfolio compared with starting value, defaults to 5% drawdown</param>
            isTrailing: If "false", the drawdown will be relative to the starting value of the portfolio.
                        If "true", the drawdown will be relative the last maximum portfolio value'''
        self.maximumDrawdownPercent = -abs(maximumDrawdownPercent)
        self.isTrailing = isTrailing
        self.initialised = False
        self.portfolioHigh = 0

    def ManageRisk(self, algorithm, targets):
        '''Manages the algorithm's risk at each time step
        Args:
            algorithm: The algorithm instance
            targets: The current portfolio targets to be assessed for risk'''
        currentValue = algorithm.Portfolio.TotalPortfolioValue

        if not self.initialised:
            self.portfolioHigh = currentValue   # Set initial portfolio value
            self.initialised = True

        # Update trailing high value if in trailing mode
        if self.isTrailing and self.portfolioHigh < currentValue:
            self.portfolioHigh = currentValue
            return []   # return if new high reached

        pnl = self.GetTotalDrawdownPercent(currentValue)
        if pnl < self.maximumDrawdownPercent and len(targets) != 0:
            self.initialised = False # reset the trailing high value for restart investing on next rebalcing period
            return [ PortfolioTarget(target.Symbol, 0) for target in targets ]

        return []

    def GetTotalDrawdownPercent(self, currentValue):
        return (float(currentValue) / float(self.portfolioHigh)) - 1.0