import math
import numpy as np
from RiskManager import StopLimit


class DeterminedLightBrownAlpaca(QCAlgorithm):

    def Initialize(self):

        self.SetCash(10000)
        self.SetStartDate(2012, 1, 1)
        self.AddRiskManagement( StopLimit(self, UpPct = 0.15, DownPct = 1.0) )

        self.initials = {}
        self.baseline_set = False
        self.SymbolData = []

        self.Keeps = int(self.GetParameter("Keeps"))
        self.CashReservePct = float(self.GetParameter("CashReservePct"))
        self.WeightRa = float(self.GetParameter("WeightRa"))
        self.WeightRb = float(self.GetParameter("WeightRb"))
        self.WeightV = float(self.GetParameter("WeightV"))
        self.FilterSMA = int(self.GetParameter("FirstSMA"))
        self.LookbackRa = self.GetParameter("LookbackRa")
        self.LookbackRb = self.GetParameter("LookbackRb")
        self.LookbackV = self.GetParameter("LookbackV")
        self.Url = ""
        self.RebalancePeriod = self.GetParameter("RebalancePeriod")
        self.GrowthSymbols = [
            #"HYG", # corporate
            "IEF", # 7-10
            #"IGIB", # 5-10
            #"JNK",
            "LQD",
            "TBF",
            #"TBX",
            "TLT"
        ]

        self.Inverses = [
            "TYO",
            #"DFVS", # data issue
            #"JNK",
            "SJB",
            "TMV",
            #"TBX", # data issue
            "TTT"
        ]
        
        for i in self.Inverses:
            self.AddEquity(i)
        
        self.TargetDict = dict(zip(self.GrowthSymbols, self.Inverses))
        
        #TODO: not in love with this
        self.SafetySymbol = self.AddEquity("SHY", Resolution.Daily).Symbol
        # List<string> GrowthSymbols = new List<String>{
        #     // safety symbol
        #     "SHY"
        # };
        self.rebalance = True
        self.LastRotationTime = self.Time

        self.Debug("WeightRa: " + str(self.WeightRa))
        self.Debug("WeightRb: " + str(self.WeightRb))
        self.Debug("WeightV: " + str(self.WeightV))
        self.Debug("SymbolUrl: " + str(self.Url))
        self.Debug("RebalancePeriod: " + str(self.RebalancePeriod))

        periodADays = self.DaysFromLookback(self.LookbackRa)
        periodBDays = self.DaysFromLookback(self.LookbackRb)
        periodVolDays = self.DaysFromLookback(self.LookbackV)

        for symbol in self.GrowthSymbols:
            self.Debug("adding symbol to universe: " + symbol)

            symbolObject = self.AddSecurity(SecurityType.Equity, symbol, Resolution.Daily).Symbol

            sData = SymbolData(self, symbolObject, periodADays, periodBDays, periodVolDays, self.FilterSMA)
            self.SymbolData.append(sData)

        # we don't look back more than a year
        self.SetWarmup(timedelta(days=365))
        self.lastDay = -1
        self.newWeek = True
        self.newMonth = True
        self.counter = 2
        self.newBiWeek = True

        self.spy = self.AddEquity("SPY", Resolution.Daily).Symbol
        self.SetBenchmark('SPY')

        self.Schedule.On(self.DateRules.MonthStart("SPY"),
                         self.TimeRules.At(0, 0),
                         self.TrackMonth)

        self.Schedule.On(self.DateRules.WeekStart("SPY"),
                         self.TimeRules.At(0, 0),
                         self.TrackWeek)

        self.Schedule.On(self.DateRules.WeekStart("SPY"),
                         self.TimeRules.At(0, 0),
                         self.TrackBiWeekly)
                         
        self.Schedule.On(self.DateRules.EveryDay("SPY"),
                         self.TimeRules.AfterMarketOpen('SPY', 0), # plot performance from prev day
                         self.PlotStuff)

    def TrackMonth(self):
        self.newMonth = True

    def TrackWeek(self):
        self.newWeek = True

    def TrackBiWeekly(self):
        if self.counter == 2:
            self.newBiWeek = True
            self.counter = 1
        else:
            self.counter += 1

    def DaysFromLookback(self, period):
        strAmount = period[0:len(period) - 1]
        # strAmount = period.substr(0, len(period)-1)
        amount = np.int32(strAmount)
        # amount = int32.Parse(strAmount)
        multiplier = period.endswith("M")
        if not multiplier:
            multiplier = 1
        else:
            multiplier = 21

        return multiplier * amount

    def GetEtfReplayRankings(self):

        # orderedReturnA = self.SymbolData.OrderByDescending(x.ReturnA)
        orderedReturnA = sorted(self.SymbolData, key=lambda x: x.ReturnA.Current.Value, reverse=True)
        # orderedReturnB = self.SymbolData.OrderByDescending(x.ReturnB)
        orderedReturnB = sorted(self.SymbolData, key=lambda x: x.ReturnB.Current.Value, reverse=True)
        # orderedVol = self.SymbolData.OrderBy(x.Volatility)
        orderedVol = sorted(self.SymbolData, key=lambda x: x.Volatility.Current.Value, reverse=False)
        rankings = {}

        # // add all the symbols
        for sd in self.SymbolData:
            rankings[sd.Symbol] = 0.0

        for i in range(0, len(orderedReturnA), 1):
            current = orderedReturnA[i]
            rankings[current.Symbol] += i * self.WeightRa

        for i in range(0, len(orderedReturnB), 1):
            current = orderedReturnB[i]
            rankings[current.Symbol] += i * self.WeightRb

        for i in range(0, len(orderedVol), 1):
            current = orderedVol[i]
            rankings[current.Symbol] += i * self.WeightV

        return rankings

    def OnData(self, data):

        if self.IsWarmingUp: return

        self.Log("OnData slice: date=" + str(self.Time))  # + ", cnt=" + data.Count);

        if not self.rebalance:
            # // var delta = Time.Subtract(LastRotationTime);
            # // rebalance = delta > RotationInterval;

            self.rebalance = self.LastRotationTime.month != self.Time.month
            if (self.RebalancePeriod == "Quarterly"):
                # // January - 1
                # // April - 4
                # // July - 7
                # // Oct - 10
                self.rebalance = (self.rebalance and
                                  ((self.Time.month == 1) or
                                   (self.Time.month == 4) or
                                   (self.Time.month == 7) or
                                   (self.Time.month == 10)))

                if self.rebalance:
                    self.Debug("Rebalancing Quarterly on " + str(self.Time))

            elif (self.RebalancePeriod == "Daily"):
                if self.Time.day != self.lastDay:
                    self.rebalance = True
                    self.lastDay = self.Time.day
                    self.Debug("Rebalancing Daily on " + str(self.Time))
                else:
                    self.rebalance = False

            elif (self.RebalancePeriod == "Weekly"):
                if self.newWeek:
                    self.rebalance = True
                    self.newWeek = False
                    self.Debug("Rebalancing Weekly on " + str(self.Time))
                else:
                    self.rebalance = False

            elif (self.RebalancePeriod == "Bi-Monthly"):
                if self.newBiWeek:
                    self.rebalance = True
                    self.newBiWeek = False
                    self.Debug("Rebalancing Bi Weekly on " + str(self.Time))
                else:
                    self.rebalance = False

            elif (self.RebalancePeriod == "Monthly"):
                if self.newMonth:
                    self.rebalance = True
                    self.newMonth = False
                    self.Debug("Rebalancing Monthly on " + str(self.Time))
                else:
                    self.rebalance = False

        if self.rebalance:
            self.rebalance = False
            self.LastRotationTime = self.Time

            # pick which one is best from growth and safety symbols
            rankings = self.GetEtfReplayRankings()
            # orderedObjScores = self.SymbolData.OrderBy(rankings[x.Symbol]).ToList()
            orderedObjScores = sorted(self.SymbolData, key=lambda x: rankings[x.Symbol], reverse=False)

            nextHoldings = []

            self.Debug("OrderedObjScores Count: " + str(len(orderedObjScores)))
            for i in range(0, self.Keeps, 1):
                currentSymbolData = orderedObjScores[i]
                # if currentSymbolData.Symbol not in data.Bars:continue
                lastClose = self.Securities[currentSymbolData.Symbol].Close
                maFilter = lastClose > currentSymbolData.FilterSMA.Current.Value
                if maFilter:
                    # // this meets our investment criteria
                    if self.Portfolio[self.SafetySymbol].Invested:
                        # exit shy
                        self.Liquidate(self.SafetySymbol)
                        
                    # short inverse mapped to this ticker
                    k = str(currentSymbolData.Symbol).split(' ')[0]
                    v = self.TargetDict[k]
                    nextHoldings.append(v)
                else:
                    if self.SafetySymbol not in nextHoldings:
                        nextHoldings.append(self.SafetySymbol)

                # FIXME - allocate to "Safety" symbol, right now saftey symbol is part of the growth symbols so it should bubble up but that doesn't match replay
                # else:
                # if not self.Portfolio[self.SafetySymbol].Invested:
                #     self.Liquidate()
                #     self.SetHoldings(self.SafetySymbol, 1)
                #     break
                # if self.Portfolio[self.SafetySymbol].Invested:
                #     break

            if self.Portfolio[self.SafetySymbol].Invested:
                return

            self.Log(">>NextPositions<<")
            for position in nextHoldings:
                self.Log("\t>>" + str(position))

            if len(nextHoldings) == 0:
                # go to 100% cash
                self.Log("LIQUIDATE>>CASH")
                self.Liquidate()

            else:
                for kvp in self.Portfolio.Securities.Values:
                    # liquidate anything we are currently invested in but not part of our next portfolio state
                    if kvp.Invested and not (kvp.Symbol in nextHoldings):
                        self.Log("LIQUIDATE>>" + str(kvp.Symbol))
                        self.Liquidate(kvp.Symbol)

                allocationPercentage = (1.0 - self.CashReservePct) / self.Keeps
                for symbol in nextHoldings:
                    self.Log("BUY>>" + str(symbol))
                    self.SetHoldings([PortfolioTarget(symbol, -allocationPercentage)])

    def PlotStuff(self):
        # normalize data for plotting
        if not self.baseline_set:
            self.initials['SPY'] = self.Securities['SPY'].Price
            if self.initials['SPY'] != 0:
                self.baseline_set = True
        if not self.baseline_set:
            return

        self.Plot("Performance", 'Portfolio', (self.Portfolio.TotalPortfolioValue - 10000) / 10000 * 100)
        self.Plot("Performance", 'SPY',
                  (self.Securities['SPY'].Price - self.initials['SPY']) / self.initials['SPY'] * 100)

class SymbolData:
    def __init__(self, algo, symbol, periodADays, periodBDays, periodVolDays, FilterSMA):
        self.Symbol = symbol
        self.algo = algo
        self.ReturnA = self.algo.MOMP(symbol, periodADays, Resolution.Daily)
        self.ReturnB = self.algo.MOMP(symbol, periodBDays, Resolution.Daily)
        # self.Volatility = Volatility

        self.FilterSMA = SimpleMovingAverage(FilterSMA)
        consolidator = TradeBarConsolidator(CalendarType.Monthly)
        self.algo.RegisterIndicator(self.Symbol, self.FilterSMA, consolidator)

        std = StandardDeviation(periodVolDays)
        annualizationFactor = float(math.sqrt(252))
        logReturns = self.algo.LOGR(symbol, periodVolDays, Resolution.Daily)

        # CompositeIndicator<IndicatorDataPoint> volatility = std.Of(logReturns).Times(annualizationFactor);
        volatility_ = IndicatorExtensions.Of(std, logReturns)
        self.Volatility = IndicatorExtensions.Times(volatility_, annualizationFactor)

from clr import AddReference
AddReference("System")
AddReference("QuantConnect.Common")
AddReference("QuantConnect.Algorithm")
AddReference("QuantConnect.Algorithm.Framework")

from QuantConnect import *
from QuantConnect.Algorithm import *
from QuantConnect.Algorithm.Framework import *
from QuantConnect.Algorithm.Framework.Portfolio import PortfolioTarget
from QuantConnect.Algorithm.Framework.Risk import RiskManagementModel

class StopLimit(RiskManagementModel):
    '''Provides an implementation of IRiskManagementModel that limits the drawdown per holding to the specified percentage'''

    def __init__(self, algo, UpPct = 0.05, DownPct = 0.05):
        '''Initializes a new instance of the MaximumDrawdownPercentPerSecurity class
        Args:
            maximumDrawdownPercent: The maximum percentage drawdown allowed for any single security holding'''
        self.maxDown = -abs(DownPct)
        self.maxUp = abs(UpPct)
        self.algo = algo

    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'''
        targets = []
        for kvp in algorithm.Securities:
            security = kvp.Value

            if not security.Invested:
                continue

            pnl = security.Holdings.UnrealizedProfitPercent
            if pnl < self.maxDown:
                # liquidate
                self.algo.Debug(f'Cut Loss {security.Symbol}. Loss: {self.algo.Portfolio[security.Symbol].UnrealizedProfitPercent*100:.2f}%')
                targets.append(PortfolioTarget(security.Symbol, 0))
            elif pnl > self.maxUp:
                # liquidate
                self.algo.Debug(f'Take Profit {security.Symbol}. Profit: {self.algo.Portfolio[security.Symbol].UnrealizedProfitPercent*100:.2f}%')
                targets.append(PortfolioTarget(security.Symbol, 0))

        return targets