#region imports
using System;
using System.Collections.Generic;
using System.Linq;
using QuantConnect.Indicators;
using QuantConnect.Securities;
using QuantConnect.Scheduling;
using QuantConnect.Data.UniverseSelection;
using QuantConnect.Data.Consolidators;
using QuantConnect.Orders;
using QuantConnect.Algorithm.Framework.Selection;
#endregion

namespace QuantConnect.Algorithm.CSharp
{
    /*
    *******************************************************************************************
    Connors Mean Reversion Strategy
    *******************************************************************************************
    
    ** Objective of the algorithm: **
    To capitalize on short-term mean reversion opportunities in stocks from the S&P 500.

    ** Key Trading Rules: **
    1. Long-term Trend Following Regime Filter: SPY’s total return over the last six months (126 trading days) is positive.
    2. Liquidity Filter. The stock must be one of the 500 most liquid US stocks. 
    3. The Weekly 2-period RSI of the stock must be below 20, indicating it has overreacted to the downside.
    4. Rank all qualifying stocks by their trailing 100-day historical volatility, then BUY on the close, in equal weight, the 10 stocks with the LOWEST historical volatility.
    5. SELL the stock on the close if its weekly 2-period RSI is above 80.
    6. SELL the stock on the close if the current price is more than 10% below the entry price, checked at the end of every business day.

    *******************************************************************************************
    */

    public class ConnorsMeanReversionStrategy : QCAlgorithm
    {
        private const int RsiPeriod = 2; // Rule 3: Weekly 2-period RSI
        private const int LookbackPeriod = 100; // Rule 4: 100-day historical volatility
        private const decimal MaxPortfolioRisk = 0.1m; // Rule 6: Max position size (10% risk)
        private const int PortfolioStockSize = 10; // Rule 6: Max position size (10% risk)
        private RateOfChange _spyROC126; // Rule 1: Long-term trend filter for SPY


        private Dictionary<Symbol, RelativeStrengthIndex> _RSIIndicator = new Dictionary<Symbol, RelativeStrengthIndex>();
        private Dictionary<Symbol, StandardDeviation> _volatilityIndicator = new Dictionary<Symbol, StandardDeviation>();

        public override void Initialize()
        {
            // Set the start and end dates for the backtest
            SetStartDate(1998, 1, 1);
            SetEndDate(DateTime.Now);
            SetCash(100000); // Initial capital

            // Rule 1: Add SPY for market regime filter
            var spy = AddEquity("SPY", Resolution.Daily);
            _spyROC126 = ROC(spy.Symbol, 126, Resolution.Daily);
            SetBenchmark(spy.Symbol);

            // Rule 2: The stock must be one of the 500 most liquid US stocks.
            UniverseSettings.Resolution = Resolution.Daily;
            SetUniverseSelection(new QC500UniverseSelectionModel());

            SetWarmUp(LookbackPeriod);
            Schedule.On(DateRules.Every(DayOfWeek.Monday), TimeRules.AfterMarketOpen("SPY", 0), ExecuteMondaySelling);
            Schedule.On(DateRules.Every(DayOfWeek.Monday), TimeRules.AfterMarketOpen("SPY", 1), ExecuteMondayTrading);
            Schedule.On(DateRules.EveryDay(), TimeRules.AfterMarketClose("SPY", 0), ExecuteDailySelling);
        }



        public override void OnSecuritiesChanged(SecurityChanges changes)
        {
            if (changes.AddedSecurities.Count > 0)
            {
                foreach (Security security in changes.AddedSecurities)
                {
                    // Ensure SPY is not added
                    if (!_RSIIndicator.ContainsKey(security.Symbol) && security.Symbol.Value != "SPY")
                    {
                        var symbol = security.Symbol;

                        // Rule 3: Create RSI for the stock
                        var rsi = new RelativeStrengthIndex(RsiPeriod, MovingAverageType.Wilders);
                        var consolidator = new TradeBarConsolidator(CalendarType.Weekly);
                        RegisterIndicator(symbol, rsi, consolidator);

                        // Warm up RSI with historical data
                        var rsiHistory = IndicatorHistory(rsi, symbol, LookbackPeriod, Resolution.Daily);

                        // Add RSI indicator to the dictionary
                        _RSIIndicator.Add(symbol, rsi);

                        // Rule 4: Create volatility indicator (Standard Deviation)
                        var stdDev = new StandardDeviation(LookbackPeriod);
                        RegisterIndicator(security.Symbol, stdDev, Resolution.Daily);

                        // Warm up volatility indicator with historical data
                        var stdDevHistory = IndicatorHistory(stdDev, security.Symbol, LookbackPeriod, Resolution.Daily);

                        _volatilityIndicator.Add(security.Symbol, stdDev);
                    }
                }
            }

            // Handle removed securities
            if (changes.RemovedSecurities.Count > 0)
            {
                foreach (var security in changes.RemovedSecurities)
                {
                    if (security.Invested)
                    {
                        Liquidate(security.Symbol); // Liquidate if invested
                    }
                    if (_RSIIndicator.ContainsKey(security.Symbol))
                    {
                        _RSIIndicator.Remove(security.Symbol);
                        _volatilityIndicator.Remove(security.Symbol);
                    }
                }
            }
        }

        private void ExecuteMondayTrading()
        {
            // Ensure the algorithm is not warming up
            if (IsWarmingUp) return;

            //Rule: 1 SPY’s total return over the last six months (126 trading days) is positive
            if (_spyROC126 > 0)
            {

                // List of candidate stocks and their volatility
                var candidates = new List<Symbol>();
                var volatilityDict = new Dictionary<Symbol, decimal>();

                foreach (var symbol in _RSIIndicator.Keys)
                {
                    // Check if the price data is valid (greater than zero)
                    var currentPrice = Securities[symbol].Price;
                    if (currentPrice <= 0) continue; // Ensure the price is valid

                    var rsi = _RSIIndicator[symbol]; // Access the RSI for the symbol

                    // Rule 3: Check if the RSI is below 20
                    if (rsi.Current.Value < 20)
                    {
                        candidates.Add(symbol);

                        // Rule 4: Calculate historical volatility
                        var stdDev100 = _volatilityIndicator[symbol].Current.Value; // Access volatility indicator
                        volatilityDict[symbol] = stdDev100; // Store the volatility
                    }
                }

                // Select the 10 stocks with the lowest volatility
                var selectedStocks = volatilityDict
                    .Where(x => candidates.Contains(x.Key))
                    .OrderBy(x => x.Value)
                    .Take(PortfolioStockSize)
                    .Select(x => x.Key)
                    .ToList();

                // Proceed if there are selected stocks
                if (selectedStocks.Count > 0)
                {
                    // Count current holdings
                    int currentCount = Portfolio.Values.Count(v => v.Invested);

                    // Add new stocks until the portfolio reaches 10
                    foreach (var symbol in selectedStocks)
                    {
                        if (currentCount < PortfolioStockSize) // Max of 10 new positions
                        {
                            SetHoldings(symbol, MaxPortfolioRisk); // Allocate 10% to the new stock
                            currentCount++; // Increment count
                        }
                    }
                }
            }
        }


        private void ExecuteMondaySelling()
        {
            // Check all positions for selling criteria
            foreach (var symbol in Portfolio.Keys.ToList())
            {
                // Access the RSI for the symbol
                if (_RSIIndicator.TryGetValue(symbol, out var rsi))
                {
                    // Rule 5: Sell if RSI is above 80
                    if (rsi.Current.Value > 80)
                    {
                        Liquidate(symbol);
                    }
                }
            }
        }

        public void ExecuteDailySelling()
        {
            // Rule 6: Sell if the current price is more than 10% below entry price
            foreach (var symbol in Portfolio.Keys.ToList())
            {
                var currentPrice = Securities[symbol].Price;
                var entryPrice = Portfolio[symbol].AveragePrice;

                if (currentPrice < entryPrice * 0.9m) // 10% below entry price
                {
                    Liquidate(symbol);
                }
            }
        }
    }
}