# region imports
from datetime import timedelta
from turtle import up
from AlgorithmImports import *
from QuantConnect.Indicators import *
# endregion


class FORM(QCAlgorithm):

  def Initialize(self):
    self.SetStartDate(2009, 7, 1)
    self.SetCash(100000)
    self.UniverseSettings.Resolution = Resolution.Daily
    self.EnableAutomaticIndicatorWarmUp = True

    self.AddEquity("SPY", Resolution.Daily)
    self.Schedule.On(self.DateRules.MonthEnd(
        "SPY"), self.TimeRules.BeforeMarketClose("SPY", 10), self.UpdateTrend)
    self.Schedule.On(self.DateRules.MonthStart("SPY"),
                     self.TimeRules.AfterMarketOpen("SPY", 0), self.Rebalance)

    symbol = Symbol.Create("QQQ", SecurityType.Equity, Market.USA)
    universe_model = ETFConstituentsUniverseSelectionModel(
        symbol, self.UniverseSettings, self.ETFConstituentsFilter)
    self.AddUniverseSelection(universe_model)

    self.spy_sma = self.SMA("SPY", 200, Resolution.Daily)
    self.symbol_data = {}
    self.is_uptrend = False
    self.last_spy_close = 0

  def ETFConstituentsFilter(self, constituents: List[ETFConstituentData]) -> List[Symbol]:
    return [c.Symbol for c in constituents]

  def OnSecuritiesChanged(self, changes):
    for security in changes.AddedSecurities:
      if security.Symbol != "SPY":
        symbol = security.Symbol
        new_symbol_data = SymbolData(symbol)
        self.WarmUpIndicator(new_symbol_data.symbol, new_symbol_data.sma)
        self.WarmUpIndicator(new_symbol_data.symbol, new_symbol_data.rsi)
        self.symbol_data[symbol] = new_symbol_data
    for security in changes.RemovedSecurities:
      if security.Symbol in self.symbol_data:
        self.symbol_data.pop(security.Symbol)

  def OnData(self, slice):
    for symbol, trade_bar in slice.Bars.items():
      if symbol == "SPY":
        self.last_spy_close = trade_bar.Close
      elif symbol in self.symbol_data:
        self.symbol_data[symbol].update(trade_bar.EndTime, trade_bar.Close)
        history = self.History(symbol, trade_bar.EndTime - timedelta(days=90), trade_bar.EndTime - timedelta(days=89), Resolution.Daily, True, True)
        if 'close' in history and history['close'].any:
          self.symbol_data[symbol].update_performance(trade_bar.Close, history['close'][0])

  def UpdateTrend(self):
    if not self.spy_sma.IsReady:
      return
    self.is_uptrend = self.last_spy_close > self.spy_sma.Current.Value
    # if not self.is_uptrend:
    #   self.Liquidate()

  def Rebalance(self):
    if self.is_uptrend:
      uptrend_securities = [symbol for symbol,
                            dat in self.symbol_data.items() if dat.is_uptrend]
      uptrend_securities.sort(
          key=lambda x: self.symbol_data[x].performance, reverse=True)
      uptrend_securities = set(uptrend_securities[:10])

      uptrend_symbols = [symbol.Value for symbol in uptrend_securities]
      self.Log(f"Uptrend securities: {uptrend_symbols}")

      for kvp in self.Portfolio:
        holding = kvp.Value
        symbol = holding.Symbol
        if holding.Invested and symbol not in uptrend_securities and symbol.Value != "SPY":
          self.Log("Liquidating " + str(symbol))
          self.Liquidate(symbol)
      for symbol in uptrend_securities:
        self.Log("Rebalancing " + str(symbol))
        self.SetHoldings(symbol, 0.1)


class SymbolData(object):

  def __init__(self, symbol: str):
    self.symbol = symbol
    self.sma = SimpleMovingAverage(200)
    self.rsi = RelativeStrengthIndex(90)
    self.is_uptrend = False
    self.performance = 0

  def update(self, time, value):
    if self.sma.Update(time, value):
      self.is_uptrend = value > self.sma.Current.Value
    self.rsi.Update(time, value)

  def update_performance(self, cur_price, prev_price):
    self.performance = (cur_price - prev_price) / prev_price