# QUANTCONNECT.COM - Democratizing Finance, Empowering Individuals.
# Lean Algorithmic Trading Engine v2.0. Copyright 2014 QuantConnect Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from AlgorithmImports import *
from datetime import timedelta

class SpyWeeklyRotation(QCAlgorithm):

    _selection_data_by_symbol = {}
    
    def initialize(self):
        self.set_start_date(2020, 1, 1)
        self.set_end_date(2025, 6, 1)

        self.set_cash(100000)

        self.set_benchmark("SPY")

        # Seed the price of each asset that enters the universe with its last known price so you can trade it 
        # on the same morning it enters the universe without getting warnings.
        self.settings.seed_initial_prices = True

        # Add a chained universe that selects the SPY constituents trading above their 200-day SMA.
        spy = Symbol.create("SPY", SecurityType.EQUITY, Market.USA)
        self.universe_settings.resolution = Resolution.DAILY
        self._universe = self.add_universe(self.universe.etf(spy), self._fundamental_selection)

        # Trade daily at market open since the trading signal is generated on a daily resolution.
        self.schedule.on(self.date_rules.every_day(spy), self.time_rules.after_market_open(spy, 1), self._rebalance)
    
    def _fundamental_selection(self, fundamental: List[Fundamental]) -> List[Symbol]:
        # Create/Update an SMA indicator for each asset that enters the ETF.
        universe_symbols = []
        for f in fundamental:
            # Ignore stocks trading below $1 to avoid illiquid penny stocks.
            if f.price < 1:
                continue
            universe_symbols.append(f.symbol)
            if f.symbol not in self._selection_data_by_symbol:
                self._selection_data_by_symbol[f.symbol] = SelectionData(self, f.symbol)
            self._selection_data_by_symbol[f.symbol].update_price(self.time, f.adjusted_price)
            self._selection_data_by_symbol[f.symbol].update_volume(self.time, f.volume)
    
        # Remove indicators for assets that are no longer in the ETF to release the computational resources.
        symbols_to_remove = [s for s in self._selection_data_by_symbol.keys() if s not in universe_symbols]
        for symbol in symbols_to_remove:
            self._selection_data_by_symbol.pop(symbol)

        # Select only stocks with 20 days average volume > 1,000,000 to ensure sufficient liquidity.
        # Ignore stocks with RSI < 50 
        # Calculate RateOfChange(200) and select only those with ROC > 0
        # Sort by highest ROC and pick top 10 stocks. 

        # Select the Equities where rsi is not overbought and volume is above threshold.
        selected = [
            selection_data for _, selection_data in self._selection_data_by_symbol.items() if selection_data.is_volume_above_threshold()
            and selection_data.is_rsi_not_overbought()
        ]
        
        # Sort by RateOfChange(200) descending and take top 10
        selected = sorted(selected, key=lambda x: x.roc_value() or float('-inf'), reverse=True)[:10]
        selected = [selection_data.symbol for selection_data in selected]

        # Plot the results.
        self.plot("Universe", "Possible", len(list(fundamental)))
        self.plot("Universe", "Selected", len(selected))
    
        return selected

    def _rebalance(self) -> None:
        # If current day is not Monday return 
        if self.time.weekday() != 0:
            return
        
        self.log(f"Rebalancing Portfolio on {self.time}")

        # Form an equal-weighted portfolio with the Equities that are selected in the universe.
        symbols = [symbol for symbol in self._universe.selected if self.securities[symbol].price]
        if not symbols:
            return
        weight = 1 / len(symbols)
        self.set_holdings([PortfolioTarget(symbol, weight) for symbol in symbols], liquidate_existing_holdings=True)


# Define a separate class to contain the SMA indicator.
class SelectionData(object):

    symbol = None
    def __init__(self, algorithm:QCAlgorithm, symbol, period=200):
        
        self.symbol = symbol
        #  Create the SMA indicator for trend detection and filtering.
        self._sma = SimpleMovingAverage(period)
        algorithm.warm_up_indicator(symbol, self._sma, Resolution.DAILY)

        self._rsi = RelativeStrengthIndex(14)
        algorithm.warm_up_indicator(symbol, self._rsi, Resolution.DAILY)

        self._roc = RateOfChange(200)
        algorithm.warm_up_indicator(symbol, self._roc, Resolution.DAILY)

        self._volumeSMA = SimpleMovingAverage(20)
        algorithm.warm_up_indicator(symbol, self._volumeSMA, Resolution.DAILY, selector=lambda x: x.volume)
        
    def update_price(self,time, value):
        self._sma.update(time, value)
        self._rsi.update(time, value)
        self._roc.update(time, value)

    def update_volume(self,time, value):
        self._volumeSMA.update(time, value)

    def is_rsi_not_overbought(self):
        return self._rsi.is_ready and self._rsi.current.value < 50
    
    def is_volume_above_threshold(self, threshold=1_000_000):
        return self._volumeSMA.is_ready and self._volumeSMA.current.value > threshold
    
    def roc_value(self):
        return self._roc.current.value if self._roc.is_ready else None