# region imports
from AlgorithmImports import *
from math import floor
# endregion

class RiskOnRiskOff(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2016, 1, 1)
        self.SetCash(10000)

        dd_period:int = 10
        self.market:Symbol = self.AddEquity("SPY", Resolution.Daily).Symbol
        self.spy_prices:RollingWindow = RollingWindow[float](dd_period)
        self.dd_threshold:float = 0.05
 
        # Add ETFs to universe
        etfs:List[str] = ["SHY", "TECL", "TQQQ", "UPRO", "TMF", "USDU", "QID", "TBF", "IEI", "GLD", "TIP", "BSV"]
        for etf in etfs:
            data = self.AddEquity(etf, Resolution.Daily)
            data.SetLeverage(10)

        # Add ETFs for RSI calculation
        self.vixm:Symbol = self.AddEquity("VIXM", Resolution.Daily).Symbol
        self.vixm_rsi:RelativeStrengthIndex = self.RSI(self.vixm, 40, Resolution.Daily)
        self.rsi_threshold:float = 0.69

        rsi_etfs:List[str] = ["TECL", "TQQQ", "UPRO", "TMF", "QID", "TBF"]
        self.rsi_symbols:Dict[str, RateOfChange] = {}
        for etf in rsi_etfs:
            self.rsi_symbols[etf] = self.RSI(etf, 10 if etf in ["TECL", "TQQQ", "UPRO", "TMF"] else 20, Resolution.Daily)

        # Add ETFs for ROC calculation
        self.roc_symbols:Dict[str, Tuple] = {}
        for etf in ["BND", "BIL", "TLT"]:
            self.AddEquity(etf, Resolution.Daily)
            self.roc_symbols[etf] = (self.ROC(etf, 20, Resolution.Daily), self.ROC(etf, 60, Resolution.Daily))
        
        self.SetWarmup(60, Resolution.Daily)
        self.recent_day:int = -1

    def OnData(self, data: Slice) -> None:
        # if not (self.Time.hour == 16 and self.Time.minute == 0): return
        
        # Store daily market prices
        if self.market in data and data[self.market]:
            self.spy_prices.Add(data[self.market].Close)

        if self.IsWarmingUp: return

        # Trading logic
        if self.vixm_rsi.IsReady and \
            self.spy_prices.IsReady and \
            all(x.IsReady for x in self.rsi_symbols.values()) and \
            all(x[0].IsReady for x in self.roc_symbols.values()) and \
            all(x[1].IsReady for x in self.roc_symbols.values()):

            if self.recent_day != self.Time.day:
                self.recent_day = self.Time.day

                if self.vixm_rsi.Current.Value / 100. > self.rsi_threshold:
                    should_rebalance:bool = self.liquidate(["SHY"])
                    if should_rebalance:
                        # quantity:float = floor(self.Portfolio.TotalPortfolioValue / data["SHY"].Value)
                        # self.MarketOnOpenOrder("SHY", quantity)

                        self.SetHoldings("SHY", 1)
                else:
                    if self.roc_symbols["BND"][1].Current.Value > self.roc_symbols["BIL"][1].Current.Value:
                        rsi_values:Dict[str, float] = { x : self.rsi_symbols[x].Current.Value for x in ["TECL", "TQQQ", "UPRO", "TMF"] }
                        bottom_by_rsi:List[str] = sorted(rsi_values, key=rsi_values.get, reverse=True)[-3:]

                        should_rebalance:bool = self.liquidate(bottom_by_rsi)
                        if should_rebalance:
                            for etf in bottom_by_rsi:
                                # quantity:float = floor(self.Portfolio.TotalPortfolioValue / len(bottom_by_rsi) / data[etf].Value)
                                # self.MarketOnOpenOrder(etf, quantity)

                                self.SetHoldings(etf, 1. / len(bottom_by_rsi))
                    else:
                        if self.roc_symbols["TLT"][0].Current.Value < self.roc_symbols["BIL"][0].Current.Value:
                            
                            rsi_values:Dict[str, float] = { x : self.rsi_symbols[x].Current.Value for x in ["QID", "TBF"] }
                            bottom_by_rsi:List[str] = sorted(rsi_values, key=rsi_values.get, reverse=True)[-1:]
                            should_rebalance:bool = self.liquidate(bottom_by_rsi + ["USDU"])
                            
                            if should_rebalance:
                                # quantity:float = floor(self.Portfolio.TotalPortfolioValue / 2 / data["USDU"].Value)
                                # self.MarketOnOpenOrder("USDU", quantity)
                                # quantity:float = floor(self.Portfolio.TotalPortfolioValue / 2 / data[bottom_by_rsi[0]].Value)
                                # self.MarketOnOpenOrder(bottom_by_rsi[0], quantity)

                                self.SetHoldings("USDU", 0.5)
                                self.SetHoldings(bottom_by_rsi[0], 0.5)
                        else:
                            max_spy_drawdown:float = self.calculate_max_drawdown(np.array(list(self.spy_prices)[::-1]))
                            if max_spy_drawdown > -self.dd_threshold:
                                should_rebalance:bool = self.liquidate(["UPRO", "TMF"])
        
                                if should_rebalance:
                                    # quantity:float = floor(self.Portfolio.TotalPortfolioValue * 0.55 / data["UPRO"].Value)
                                    # self.MarketOnOpenOrder("UPRO", quantity)
                                    # quantity:float = floor(self.Portfolio.TotalPortfolioValue * 0.45 / data["TMF"].Value)
                                    # self.MarketOnOpenOrder("TMF", quantity)

                                    self.SetHoldings("UPRO", 0.55)
                                    self.SetHoldings("TMF", 0.45)
                            else:
                                tickers_to_hold:List[str] = ["IEI", "GLD", "TIP", "BSV"]
                                should_rebalance:bool = self.liquidate(tickers_to_hold)
                                
                                if should_rebalance:
                                    for etf in tickers_to_hold:
                                        # quantity:float = floor(self.Portfolio.TotalPortfolioValue / len(tickers_to_hold) / data[etf].Value)
                                        # self.MarketOnOpenOrder(etf, quantity)

                                        self.SetHoldings(etf, 1 / len(tickers_to_hold))

    def liquidate(self, tickers_to_hold:List[str]) -> bool:
        should_rebalance:bool = False

        invested:List[str] = [x.Key.Value for x in self.Portfolio if x.Value.Invested]
        if len(invested) == 0:
            should_rebalance = True

        for ticker in invested:
            if ticker not in tickers_to_hold:
                symbol = self.Symbol(ticker)
                # self.MarketOnOpenOrder(symbol, -self.Portfolio[symbol].Quantity)
                self.Liquidate(ticker)
                should_rebalance = True
        
        return should_rebalance
        
    def calculate_max_drawdown(self, prices:np.ndarray) -> float:
        prices:pd.Series = pd.Series(prices)
        roll_max:pd.Series = prices.cummax()
        daily_dd:pd.Series = prices / roll_max - 1.0
        max_daily_dd:float = daily_dd.min()

        return max_daily_dd