# https://seekingalpha.com/article/4299701-leveraged-etfs-for-long-term-investing
# Momentum addaped verision of:
# https://www.quantconnect.com/forum/discussion/7708/using-levered-etfs-in-ira-10-years-24-cagr-1-56-sharpe/p1
# with some extra symbols.

import pandas as pd
import numpy as np

class MultidimensionalTransdimensionalPrism(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 7, 7)  # Set Start Date
        self.SetCash(10000)  # Set Strategy Cash
        self.SetEndDate(2020, 7, 16)
        
        # create a dictionary to store momentum indicators for all symbols 
        self.mom = {}
        self.rsi = {}
        self.momp = {}
        
        #MOM period
        period = 21*4
        #rsi period
        period2 = 10
        #MOMP period
        period3 = 44
        
        self.SetRiskManagement(TrailingStopRiskManagementModel(0.25))
        
        # warm up the MOM indicator
        self.SetWarmUp(period)
        self.symbols = []
        tickers = ["TQQQ", "UBT", "TECL","UGLD","TYD"]
        for ticker in tickers:
            symbol = self.AddEquity(ticker, Resolution.Hour).Symbol
            self.symbols.append(symbol)
            self.mom[symbol] = self.MOM(symbol, period, Resolution.Hour)
            self.rsi[symbol] = self.RSI(symbol, period2, Resolution.Daily)
            self.momp[symbol] = self.MOMP(symbol, period3, Resolution.Daily)

        self.Schedule.On(self.DateRules.Every(DayOfWeek.Wednesday), \
                 self.TimeRules.At(15, 30), \
                 self.Rebalance)
                
    def Rebalance(self):
        if self.IsWarmingUp: return
        
        alldata = pd.DataFrame({'momp':[self.mom[symbol].Current.Value for symbol in self.symbols],
                                'rsi': [self.rsi[symbol].Current.Value for symbol in self.symbols]})
        highest_rsi = alldata[alldata['rsi'] > 80]
        top3 = highest_rsi.sort_values(by='momp', ascending = False)[:3]
        top3_symbols = [self.symbols[i] for i in top3.index]
        
        for kvp in self.Portfolio:
            security_hold = kvp.Value
            if security_hold.Invested and (security_hold.Symbol not in top3_symbols):
                self.Liquidate(security_hold.Symbol)
        
        added_symbols = []        
        for symbol in top3_symbols:
            if not self.Portfolio[symbol].Invested:
                added_symbols.append(symbol)
        for added in added_symbols:
            self.SetHoldings(added, 1/len(added_symbols))