#region imports
from AlgorithmImports import *
#endregion
"""
v2.5 Dual Momentum with Out Days  by Vladimir

inspired by Peter Guenther, Tentor Testivis, Dan Whitnable, Thomas Chang and T Smith.
converted to QC Algorithm Framework (Alpha Model + Portfolio Construction Model) by Joao Antunes

based on Intersection of ROC comparison using OUT_DAY approach by Vladimir
modified parameters BASE_RET = 83;

https://www.quantconnect.com/forum/discussion/10039/dual-momentum-with-out-days/p1/comment-29928
"""
import numpy as np


class DualMomentumWithOutDaysAlphaModel(AlphaModel):

    def __init__(self,
                 algorithm,
                 VOLA = 126,
                 BASE_RET = 83,
                 resolution = Resolution.Daily):
        self.VOLA = VOLA
        self.BASE_RET = BASE_RET
        self.resolution = resolution
        self.MKT = algorithm.AddEquity('SPY', resolution).Symbol
        self.SLV = algorithm.AddEquity('SLV', resolution).Symbol
        self.GLD = algorithm.AddEquity('GLD', resolution).Symbol
        self.XLI = algorithm.AddEquity('XLI', resolution).Symbol
        self.XLU = algorithm.AddEquity('XLU', resolution).Symbol
        self.DBB = algorithm.AddEquity('DBB', resolution).Symbol
        self.UUP = algorithm.AddEquity('UUP', resolution).Symbol

        self.count = 0

        pairs = [self.MKT, self.SLV, self.GLD, self.XLI, self.XLU, self.DBB, self.UUP]
        for symbol in pairs:
            self.consolidator = TradeBarConsolidator(timedelta(days=1))
            self.consolidator.DataConsolidated += self.consolidation_handler
            algorithm.SubscriptionManager.AddConsolidator(symbol, self.consolidator)

        self.history = algorithm.History(pairs, self.VOLA + 1, resolution)
        self.history = self.history['close'].unstack(level=0).dropna()

        self.predictionInterval = Time.Multiply(Extensions.ToTimeSpan(self.resolution), 30)
        resolutionString = Extensions.GetEnumString(resolution, Resolution)
        self.Name = f"{self.__class__.__name__}({resolutionString})"

        # Force alpha to only produce insights daily at 11.10am
        self.set_flag = False
        algorithm.Schedule.On(algorithm.DateRules.EveryDay(),
                              algorithm.TimeRules.AfterMarketOpen('SPY', 100),
                              self.SetFlag)

    def SetFlag(self):
        self.set_flag = True

    def consolidation_handler(self, sender, consolidated):
        self.history.loc[consolidated.EndTime, consolidated.Symbol] = consolidated.Close
        self.history = self.history.iloc[-(self.VOLA + 1):]

    def Update(self, algorithm, _data):
        if algorithm.IsWarmingUp or not self.set_flag:
            return []

        self.set_flag = False
        insights = []

        vola = self.history[self.MKT].pct_change().std() * np.sqrt(252)
        wait_days = int(vola * self.BASE_RET)
        period = int((1.0 - vola) * self.BASE_RET)
        r = self.history.pct_change(period).iloc[-1]
        exit_market = r[self.SLV] < r[self.GLD] and r[self.XLI] < r[self.XLU] and r[self.DBB] < r[self.UUP]

        direction = InsightDirection.Down
        if exit_market:
            self.count = 0
        else:
            self.count += 1
            if self.count >= wait_days:
                direction = InsightDirection.Up

        insights.append(Insight.Price(self.MKT, self.predictionInterval, direction))

        return insights

#region imports
from AlgorithmImports import *
#endregion
"""
v2.5 Dual Momentum with Out Days  by Vladimir

inspired by Peter Guenther, Tentor Testivis, Dan Whitnable, Thomas Chang and T Smith.
converted to QC Algorithm Framework (Alpha Model + Portfolio Construction Model) by Joao Antunes

based on Intersection of ROC comparison using OUT_DAY approach by Vladimir
modified parameters BASE_RET = 83;

https://www.quantconnect.com/forum/discussion/10039/dual-momentum-with-out-days/p1/comment-29928
"""
from itertools import groupby

from dual_momentum_with_out_days_alpha import DualMomentumWithOutDaysAlphaModel

VOLA = 126; BASE_RET = 83; RET = 252; EXCL = 21; LEV = 1.00;

class HorizontalQuantumCoil(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2008, 1, 1)
        self.SetEndDate(2022, 1, 1)
        self.SetCash(100_000)
        self.SetWarmUp(timedelta(350))
        # self.Portfolio.MarginCallModel = MarginCallModel.Null
        self.SetAlpha(DualMomentumWithOutDaysAlphaModel(self, VOLA, BASE_RET, Resolution.Daily))

        symbols = [
            Symbol.Create('QQQ', SecurityType.Equity, Market.USA),
            Symbol.Create('FDN', SecurityType.Equity, Market.USA),
            Symbol.Create('TLT', SecurityType.Equity, Market.USA),
            Symbol.Create('TLH', SecurityType.Equity, Market.USA),
        ]
        self.SetUniverseSelection(ManualUniverseSelectionModel(symbols))
        self.UniverseSettings.Resolution = Resolution.Hour
        self.SetPortfolioConstruction(OutDays(self, RET, EXCL, LEV, Resolution.Hour))

        self.createPlots("SPY")

    def createPlots(self, benchmark):
        self.__benchmark = benchmark

        self.__plot_every_n_days = 5
        self.__plot_every_n_days_i = 0
        plot = Chart('Performance')
        plot.AddSeries(Series(self.__benchmark, SeriesType.Line, 0, '%'))
        plot.AddSeries(Series("Algorithm", SeriesType.Line, 0, '%'))
        self.AddChart(plot)

        self.ResetPlot()

    def ResetPlot(self):
        self.year = self.Time.year
        self.__cost_portfolio = None
        self.__cost_benchmark = None

    def CalculateBenchmarkPerformance(self):
        price = self.Securities[self.__benchmark].Price
        if self.__cost_benchmark == None:
            self.__cost_benchmark = price
        return 100.0 * ((price / self.__cost_benchmark) - 1.0)

    def CalculatePortfolioPerformance(self):
        if self.__cost_portfolio == None:
            self.__cost_portfolio = self.Portfolio.TotalPortfolioValue
        return 100.0 * ((self.Portfolio.TotalPortfolioValue / self.__cost_portfolio) - 1.0)

    def OnEndOfDay(self):
        if self.IsWarmingUp or not self.Securities[self.__benchmark].HasData:
            return

        if self.Time.year != self.year:
            self.ResetPlot()
            self.__plot_every_n_days_i == -1

        self.__plot_every_n_days_i += 1
        if self.__plot_every_n_days_i % self.__plot_every_n_days != 0:
            return

        self.Plot('Performance', self.__benchmark, self.CalculateBenchmarkPerformance())
        self.Plot('Performance', "Algorithm", self.CalculatePortfolioPerformance())


class OutDays(PortfolioConstructionModel):
    def __init__(self,
                 algorithm,
                 RET=252,
                 EXCL=21,
                 LEV=1.00,
                 resolution = Resolution.Hour):
        self.algorithm = algorithm
        self.RET = RET
        self.EXCL = EXCL
        self.LEV = LEV
        self.STK1 = self.algorithm.AddEquity('QQQ', resolution).Symbol
        self.STK2 = self.algorithm.AddEquity('FDN', resolution).Symbol
        self.BND1 = self.algorithm.AddEquity('TLT', resolution).Symbol
        self.BND2 = self.algorithm.AddEquity('TLH', resolution).Symbol
        self.ASSETS = [self.STK1, self.STK2, self.BND1, self.BND2]

    def returns(self, symbol):
        prices = self.algorithm.History(symbol, TimeSpan.FromDays(self.RET + self.EXCL), Resolution.Daily).close
        return prices[-self.EXCL] / prices[0]

    def CreateTargets(self, algorithm, insights):
        if algorithm.IsWarmingUp:
            return []

        targets = []

        # We expect at most only one active insight since we only
        # generate insights for one equity.
        assert len(insights) <= 1
        if len(insights) == 1:
            insight = insights[0]
            self.bull = insight.Direction == InsightDirection.Up

            if self.bull:
                if self.returns(self.STK1) < self.returns(self.STK2):
                    selected = self.STK2
                else:
                    selected = self.STK1
            else:
                if self.returns(self.BND1) < self.returns(self.BND2):
                    selected = self.BND2
                else:
                    selected = self.BND1

            for asset in self.ASSETS:
                if asset != selected:
                    targets.append(PortfolioTarget.Percent(algorithm, asset, 0.0))
            targets.append(PortfolioTarget.Percent(algorithm, selected, 1.0))

        return targets