# region imports
from AlgorithmImports import *
# endregion

class DancingAsparagusChimpanzee(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2021,8, 1)  # Set Start Date
        self.SetEndDate(2023, 5, 5)  # Set Start Date
        self.SetCash(100000)  # Set Strategy Cash
        self.leverage = 0.99  # Set Leverage
        # stocks
        self.STK1 = self.AddEquity("QQQ", Resolution.Daily).Symbol
        self.STK2 = self.AddEquity("SPY", Resolution.Daily).Symbol
        self.Schedule.On(self.DateRules.MonthStart("SPY"), self.TimeRules.AfterMarketOpen("SPY"), self.rebalance)
        self.look_back_period = 63  # LookBack period
        self.SetWarmUp(self.look_back_period)
        self.weight = 0.99

    def returns(self, symbol, period):
        closing_bars = self.History(symbol, TimeSpan.FromDays(period), Resolution.Daily).close
        self.Debug('{}'.format(closing_bars.size))
        return (closing_bars[-1] - closing_bars[0]) / closing_bars[-1]

    # Check if any assets from symbol list has negative return for lookBack period
    def negative_returns_check(self, symbol_list):
        for symbol in symbol_list:
            if self.returns(symbol, self.look_back_period) < 0:
                return True
        return False

    def liquidate_all_pos(self, symbol_list):
        for symbol in symbol_list:
            if self.Portfolio[symbol].IsLong:
                self.Liquidate(symbol)

    # choose best performer
    def find_best_performer(self, symbol_list):
        max_return = self.returns(symbol_list[0], self.look_back_period)
        max_return_symbol = symbol_list[0]
        for i in range(2, len(symbol_list)):
            if self.returns(symbol_list[i], self.look_back_period) > max_return:
                max_return = self.returns(symbol_list[i], self.look_back_period)
                max_return_symbol = symbol_list[i]
        return max_return_symbol

    def allocate_safely(self, symbol):
        # check if invested
        stock_invested = [x.Symbol.Value for x in self.Portfolio.Values if x.Invested]     
        if symbol.Value not in stock_invested:
            self.SetHoldings(symbol, self.weight)

    def check_group(self, symbol_list):
        # check if return is positive
        if self.negative_returns_check(symbol_list):
            # if negative
            # liquidate positions if any
            self.liquidate_all_pos(symbol_list)
        else:
            # and find best performer
            buy_symbol = self.find_best_performer(symbol_list)
            self.allocate_safely(buy_symbol)

    def rebalance(self):
        # check stocks
        self.check_group([self.STK1, self.STK2])