from AlgorithmImports import *
import datetime as dt

class RsiMacdAlphaModel(AlphaModel):
    '''Alpha model that uses an STO_RSI and MACD values to create insights'''
    last_consolidation_time = None
    minute = 0

    def __init__(self,
                 algo,
                 fastPeriod = 12,
                 slowPeriod = 26,
                 signalPeriod = 9,
                 rsiPeriod = 14,
                 movingAverageType = MovingAverageType.Exponential,
                 resolution = Resolution.Minute):
        ''' Initializes a new instance of the MacdAlphaModel class
        Args:
            fastPeriod: The MACD fast period
            slowPeriod: The MACD slow period</param>
            signalPeriod: The smoothing period for the MACD signal
            rsiPeriod: period for STO_RSI calculation
            movingAverageType: The type of moving average to use in the MACD'''

        self.algo = algo
        self.fastPeriod = fastPeriod
        self.slowPeriod = slowPeriod
        self.signalPeriod = signalPeriod
        self.rsiPeriod = rsiPeriod
        self.movingAverageType = movingAverageType
        self.resolution = resolution
        self.insightPeriod = timedelta(days=3)
        self.bounceThresholdPercent = 0
        self.symbolData = {}
        self.Fminute = None

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


    def Update(self, algorithm, data):
        ''' Determines an insight for each security based on it's current MACD signal
        Args:
            algorithm: The algorithm instance
            data: The new data available
        Returns:
            The new insights generated'''

        insights = []

        if algorithm.Portfolio.Invested:
            if self.last_fconsolidation_time is not None and self.Fminute != self.last_fconsolidation_time.minute:
                self.Fminute = self.last_fconsolidation_time.minute
                for key, sd in self.symbolData.items():
                    normalized_signal = sd.MACD_15.Signal.Current.Value / sd.Security.Price
                    is_macd_bullish = normalized_signal > self.bounceThresholdPercent
                    is_rsi_bullish = sd.STO_RSI_15.Current.Value >= 80
                    is_rsi_bearish = sd.STO_RSI_15.Current.Value <= 20

                    if algorithm.Portfolio[self.algo.long_symbol].Invested:
                        if not is_macd_bullish and is_rsi_bearish:
                            self.algo.rm_model.maximumDrawdownPercent = 0.12
                    elif algorithm.Portfolio[self.algo.short_symbol].Invested:
                        if is_macd_bullish and is_rsi_bullish:
                            self.algo.rm_model.maximumDrawdownPercent = 0.12

        if self.last_consolidation_time is not None and self.minute != self.last_consolidation_time.minute:
            self.minute = self.last_consolidation_time.minute

            for key, sd in self.symbolData.items():
                self.key = key
                if sd.Security.Price == 0:
                    continue

                normalized_signal = sd.MACD.Signal.Current.Value / sd.Security.Price
                is_macd_bullish = normalized_signal > self.bounceThresholdPercent
                is_rsi_bullish = sd.STO_RSI.Current.Value >= 80
                is_rsi_bearish = sd.STO_RSI.Current.Value <= 20

                if is_macd_bullish and is_rsi_bullish:
                    insights.append(Insight.Price(self.algo.long_symbol, self.insightPeriod, InsightDirection.Up))
                    insights.append(Insight.Price(self.algo.short_symbol, self.insightPeriod, InsightDirection.Flat))
                    self.algo.rm_model.maximumDrawdownPercent = 0.05

                elif not is_macd_bullish and is_rsi_bearish:
                    insights.append(Insight.Price(self.algo.short_symbol, self.insightPeriod, InsightDirection.Up))
                    insights.append(Insight.Price(self.algo.long_symbol, self.insightPeriod, InsightDirection.Flat))
                    self.algo.rm_model.maximumDrawdownPercent = 0.05
                    

        elif ((algorithm.Time.weekday() == 4) and (algorithm.Time.time()>dt.time(15,45))) or self.algo.rm_liquidation:
            insights.append(Insight.Price(self.algo.short_symbol, self.insightPeriod, InsightDirection.Flat))
            insights.append(Insight.Price(self.algo.long_symbol, self.insightPeriod, InsightDirection.Flat))
            self.algo.rm_liquidation = False       
            
        return insights

    def consolidation_handler(self, sender, bar):
        self.last_consolidation_time = bar.Time

    def fifteenconsolidation_handler(self, sender, bar):
        self.last_fconsolidation_time = bar.Time

    def OnSecuritiesChanged(self, algorithm, changes):
        '''Event fired each time the we add/remove securities from the data feed.
        This initializes the MACD for each added security and cleans up the indicator for each removed security.
        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:
            if added.Symbol.Value == 'QQQ':
                self.symbolData[added.Symbol] = SymbolData(algorithm, added, 
                self.fastPeriod, self.slowPeriod, self.signalPeriod, self.rsiPeriod, self.movingAverageType, self.resolution)

                self.consolidator = TradeBarConsolidator(timedelta(minutes=30))
                self.consolidator.DataConsolidated += self.consolidation_handler
                algorithm.SubscriptionManager.AddConsolidator(added.Symbol, self.consolidator)

                self.fifteenconsolidator = TradeBarConsolidator(timedelta(minutes=15))
                self.fifteenconsolidator.DataConsolidated += self.fifteenconsolidation_handler
                algorithm.SubscriptionManager.AddConsolidator(added.Symbol, self.fifteenconsolidator)


        for removed in changes.RemovedSecurities:
            if added.Symbol.Value == 'QQQ':
                data = self.symbolData.pop(removed.Symbol, None)
                if data is not None:
                    # clean up our consolidator
                    algorithm.SubscriptionManager.RemoveConsolidator(removed.Symbol, data.Consolidator)


class SymbolData:
    def __init__(self, algorithm, security, fastPeriod, slowPeriod, signalPeriod, rsiPeriod, movingAverageType, resolution):
        self.Security = security
        self.MACD = MovingAverageConvergenceDivergence(fastPeriod, slowPeriod, signalPeriod, movingAverageType)
        self.MACD_15 = MovingAverageConvergenceDivergence(fastPeriod, slowPeriod, signalPeriod, movingAverageType)

        self.RSI = RelativeStrengthIndex(rsiPeriod, MovingAverageType.Wilders)
        self.Stochastic = Stochastic(period=14, kPeriod=3, dPeriod=3)
        self.STO_RSI = IndicatorExtensions.Of(self.Stochastic,self.RSI)
        self.STO_RSI_15 = IndicatorExtensions.Of(self.Stochastic,self.RSI)

        self.Consolidator = algorithm.ResolveConsolidator(security.Symbol, timedelta(minutes=30))
        self.FifteenConsolidator = algorithm.ResolveConsolidator(security.Symbol, timedelta(minutes=15))

        algorithm.RegisterIndicator(security.Symbol, self.MACD, self.Consolidator)
        algorithm.WarmUpIndicator(security.Symbol, self.MACD, timedelta(days=5))

        algorithm.RegisterIndicator(security.Symbol, self.STO_RSI, self.Consolidator)
        algorithm.WarmUpIndicator(security.Symbol, self.STO_RSI, timedelta(days=5))

        algorithm.RegisterIndicator(security.Symbol, self.MACD_15, self.FifteenConsolidator)
        algorithm.WarmUpIndicator(security.Symbol, self.MACD_15, timedelta(days=5))

        algorithm.RegisterIndicator(security.Symbol, self.STO_RSI_15, self.FifteenConsolidator)
        algorithm.WarmUpIndicator(security.Symbol, self.STO_RSI_15, timedelta(days=5))


        algorithm.PlotIndicator('STORSI',self.STO_RSI)
        algorithm.PlotIndicator('MACD',self.MACD)
        algorithm.PlotIndicator('STORSI15',self.STO_RSI_15)
        algorithm.PlotIndicator('MACD_15',self.MACD_15)

        self.PreviousDirection = None

#region imports
from AlgorithmImports import *
import datetime as dt
#endregion


CASH = 100000
START_DATE = '01-01-2019' #'DD-MM-YYYY'
END_DATE = None #'05-01-2019'
FREE_PORTFOLIO_VALUE_PCT = 0.025


############### Parsing Logic #############
######## Do not edit anything below ######
START_DATE = dt.datetime.strptime(START_DATE, '%d-%m-%Y')
if END_DATE:
    END_DATE = dt.datetime.strptime(END_DATE, '%d-%m-%Y')

from AlgorithmImports import *

# Import from files
from constants import *
from alpha import *
from riskmanagement import *

################################################################################
class MACDRSIQQQ(QCAlgorithm):
    def Initialize(self):
        """Initialize algorithm."""
        # Set backtest details
        self.SetBacktestDetails()
        self.AddInstruments()
        self.AddModels()
        self.AddOperationalParameters()
        # self.AddSchedules()
        # self.AddPlots()

#-------------------------------------------------------------------------------
    def SetBacktestDetails(self):
        """Set the backtest details."""
        self.SetStartDate(START_DATE)
        if END_DATE:
            self.SetEndDate(END_DATE)
        self.SetCash(CASH)
        
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, 
            AccountType.Margin)

        # Adjust the cash buffer from the default 2.5% to custom setting
        self.Settings.FreePortfolioValuePercentage = FREE_PORTFOLIO_VALUE_PCT 
        self.Settings.DataSubscriptionLimit = 500 

    def AddInstruments(self):
        self.underlying = self.AddEquity('QQQ',Resolution.Minute).Symbol
        self.long_symbol = self.AddEquity('TQQQ',Resolution.Minute).Symbol
        self.short_symbol = self.AddEquity('SQQQ',Resolution.Minute).Symbol

    def AddModels(self):
        self.SetAlpha(RsiMacdAlphaModel(self))
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel(rebalancingFunc=self.Rebalance()))
        self.rm_model = TrailingStopRiskManagementModel(self,maximumDrawdownPercent=0.05)
        self.AddRiskManagement(self.rm_model)

    def Rebalance(self):
        return None

    def AddOperationalParameters(self):
        self.rm_liquidation = False

#region imports
from AlgorithmImports import *
#endregion


class TrailingStopRiskManagementModel(RiskManagementModel):
    '''Provides an implementation of IRiskManagementModel that limits the maximum possible loss
    measured from the highest unrealized profit'''
    def __init__(self, algo, maximumDrawdownPercent = 0.05):
        '''Initializes a new instance of the TrailingStopRiskManagementModel class
        Args:
            maximumDrawdownPercent: The maximum percentage drawdown allowed for algorithm portfolio 
	       compared with the highest unrealized profit, defaults to 5% drawdown'''
        self.algo = algo
        self.maximumDrawdownPercent = abs(maximumDrawdownPercent)
        self.trailing = dict()

    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'''
        riskAdjustedTargets = list()

        for kvp in algorithm.Securities:
            symbol = kvp.Key
            security = kvp.Value

            # Remove if not invested
            if not security.Invested:
                self.trailing.pop(symbol, None)
                continue

            profitPercent = security.Holdings.UnrealizedProfitPercent

            # Add newly invested securities
            value = self.trailing.get(symbol)
            if value == None:
                newValue = profitPercent if profitPercent > 0 else 0
                self.trailing[symbol] = newValue
                continue

            # Check for new high and update
            if value < profitPercent:
                self.trailing[symbol] = profitPercent
                continue

            # If unrealized profit percent deviates from local max for more than affordable percentage
            if profitPercent < value - self.maximumDrawdownPercent:
                # liquidate
                algorithm.Log(f'Liquidating at {algorithm.Time} as stop-loss has been breached')
                self.algo.rm_liquidation = True
                riskAdjustedTargets.append(PortfolioTarget(symbol, 0))

        return riskAdjustedTargets