#region imports
using System;
using System.Collections;
using System.Collections.Generic;
using System.Linq;
using System.Globalization;
using QuantConnect;
using QuantConnect.Algorithm.Framework;
using QuantConnect.Algorithm.Framework.Selection;
using QuantConnect.Algorithm.Framework.Alphas;
using QuantConnect.Algorithm.Framework.Portfolio;
using QuantConnect.Algorithm.Framework.Execution;
using QuantConnect.Algorithm.Framework.Risk;
using QuantConnect.Parameters;
using QuantConnect.Benchmarks;
using QuantConnect.Brokerages;
using QuantConnect.Util;
using QuantConnect.Interfaces;
using QuantConnect.Algorithm;
using QuantConnect.Indicators;
using QuantConnect.Data;
using QuantConnect.Data.Consolidators;
using QuantConnect.Data.Custom;
using QuantConnect.DataSource;
using QuantConnect.Data.Fundamental;
using QuantConnect.Data.Market;
using QuantConnect.Data.UniverseSelection;
using QuantConnect.Notifications;
using QuantConnect.Orders;
using QuantConnect.Orders.Fees;
using QuantConnect.Orders.Fills;
using QuantConnect.Orders.Slippage;
using QuantConnect.Scheduling;
using QuantConnect.Securities;
using QuantConnect.Securities.Equity;
using QuantConnect.Securities.Future;
using QuantConnect.Securities.Option;
using QuantConnect.Securities.Forex;
using QuantConnect.Securities.Crypto;
using QuantConnect.Securities.Interfaces;
using QuantConnect.Storage;
using QuantConnect.Data.Custom.AlphaStreams;
using QCAlgorithmFramework = QuantConnect.Algorithm.QCAlgorithm;
using QCAlgorithmFrameworkBridge = QuantConnect.Algorithm.QCAlgorithm;
#endregion

namespace QuantConnect.Algorithm.CSharp
{
    public class PSRMomentumStrategy : QCAlgorithm
    {
        private RateOfChange _spyROC252;  // Indicator for SPY momentum
        private RateOfChange _bilROC21;
        private RateOfChange _iefROC21;
        private const int MomentumLookbackPeriod = 252; // Lookback period of 252 trading days (~1 year)
        private const int PortfolioStockSize = 25; // Number of top stocks to select for the portfolio
        private Dictionary<Symbol, double> _psrDict = new Dictionary<Symbol, double>(); // Dictionary to store PSR (Price-to-Sales ratio) for each stock
        private Dictionary<Symbol, RateOfChange> _rocDict = new Dictionary<Symbol, RateOfChange>(); // Dictionary to store Rate of Change (ROC) for each stock (momentum)

        public override void Initialize()
        {
            // Set start and end dates for backtest
            SetStartDate(1998, 1, 1);
            SetEndDate(DateTime.Now);
            SetCash(100000); // Set initial cash to $100,000

            // Rule 1: Add SPY for market regime filter (momentum check)
            var spy = AddEquity("SPY", Resolution.Daily);
            _spyROC252 = ROC(spy.Symbol, MomentumLookbackPeriod, Resolution.Daily); // Compute the 252-day momentum for SPY

            var bil = AddEquity("BIL", Resolution.Daily);
            var ief = AddEquity("IEF", Resolution.Daily);

            _bilROC21 = ROC("BIL",21,Resolution.Daily);
            _iefROC21 = ROC("IEF",21,Resolution.Daily);
            
            var history = History("BIL", 21, Resolution.Daily);
            foreach (TradeBar tradeBar in history)
            {
                _bilROC21.Update(tradeBar.EndTime, tradeBar.Close);
            }
            history = History("IEF", 21, Resolution.Daily);
            foreach (TradeBar tradeBar in history)
            {
                _iefROC21.Update(tradeBar.EndTime, tradeBar.Close);
            }


            SetBenchmark(spy.Symbol); // Set SPY as the benchmark for performance comparison

            // Set universe selection model to filter for liquid U.S. stocks
            UniverseSettings.Resolution = Resolution.Daily;
            SetUniverseSelection(new PSRUniverseSelectionModel()); // Use the PSRUniverseSelectionModel to select stocks

            SetWarmUp(MomentumLookbackPeriod); // Set a warm-up period for the momentum calculation
            Schedule.On(DateRules.MonthEnd(), TimeRules.AfterMarketOpen("SPY", 0), ExecuteYearStartTrading); // Schedule yearly trading action
            Schedule.On(DateRules.MonthStart(), TimeRules.AfterMarketOpen("SPY", 0), ExecuteMonthly); // Schedule monthly rebalancing
        }

        public override void OnSecuritiesChanged(SecurityChanges changes)
        {
            // Handle changes in the universe (added/removed securities)
            if (changes.AddedSecurities.Count > 0)
            {
                foreach (Security security in changes.AddedSecurities)
                {
                    // Calculate PSR (Price-to-Sales ratio) for each stock
                    if (!_psrDict.ContainsKey(security.Symbol) && security.Symbol.Value != "SPY")
                    {
                        var psr = security.Fundamentals.ValuationRatios.PSRatio; // Get PSR from the fundamentals data
                        _psrDict[security.Symbol] = psr; // Store PSR for each stock
                    }

                    // Create the ROC indicator for momentum for each stock
                    if (!_rocDict.ContainsKey(security.Symbol) && security.Symbol.Value != "SPY")
                    {
                        var roc = new RateOfChange(MomentumLookbackPeriod); // Create a new ROC indicator for the stock
                        RegisterIndicator(security.Symbol, roc, Resolution.Daily); // Register ROC indicator for daily resolution
                        _rocDict[security.Symbol] = roc; // Store the ROC indicator for the stock
                    }
                }
            }

            // Remove outdated values when securities are removed from the universe
            if (changes.RemovedSecurities.Count > 0)
            {
                foreach (var security in changes.RemovedSecurities)
                {
                    _psrDict.Remove(security.Symbol); // Remove PSR for removed securities
                    _rocDict.Remove(security.Symbol); // Remove ROC for removed securities
                }
            }
        }

        private void ExecuteYearStartTrading()
        {
            // Ensure the algorithm is not warming up before executing trades
            if (IsWarmingUp) return;

            // Check if it's December and if SPY's momentum is positive, then select the top stocks
            //if ((Time.Month == 12 || Time.Month == 9 || Time.Month == 6 || Time.Month == 3) && _spyROC252 >= 0)
            if ((Time.Month == 12) && _spyROC252 >= 0)
            {
                top50PSRMomentumStocks(); // Select top 50 stocks based on PSR and momentum
            }
        }

        private void ExecuteMonthly()
        {
            // If SPY momentum is negative, liquidate the portfolio
            if (_spyROC252 < 0)
            {
                Liquidate();
                if (_iefROC21 > _bilROC21 && _iefROC21 > 0)
                {
                    SetHoldings("IEF", 1);
                }
            }
            else
            {
                // If the portfolio is not invested, select the top stocks for investment
                if (Portfolio["IEF"].Invested || !Portfolio.Invested)
                {
                    top50PSRMomentumStocks(); // Select top 50 stocks based on PSR and momentum
                }
            }
        }

        private void top50PSRMomentumStocks()
        {
            // Filter stocks with PSR ratio below 1 (value stocks) and order them by PSR
            var topQualityStocks = _psrDict.Where(x => x.Value < 1)  // Select stocks with PSR < 1
                                            .OrderBy(x => x.Value)  // Order by PSR (ascending)
                                            .Select(x => x.Key)     // Get the symbols of the selected stocks
                                            .ToList();

            // Calculate momentum for each selected stock using the ROC
            Dictionary<Symbol, decimal> _momentumDict = new Dictionary<Symbol, decimal>(); // Dictionary to store momentum values
            foreach (var symbol in topQualityStocks)
            {
                var roc = _rocDict[symbol]; // Get the ROC indicator for the stock

                // Calculate the ROC value (momentum)
                decimal momentum = roc.Current.Value;
                _momentumDict[symbol] = momentum; // Store the momentum value
            }

            // Select the top 50 stocks with the highest momentum
            var topMomentumStocks = _momentumDict.Where(x => x.Value > 0)
                                                .OrderByDescending(x => x.Value) // Order by momentum (descending)
                                                 .Take(PortfolioStockSize)  // Select the top 50 stocks
                                                 .Select(x => x.Key)       // Get the symbols of the selected stocks
                                                 .ToList();

            // Liquidate any existing positions and invest in the selected top 50 momentum stocks
            Liquidate(); // Liquidate current positions
            foreach (var symbol in topMomentumStocks)
            {
                SetHoldings(symbol, 1m / PortfolioStockSize); // Equally allocate the portfolio to each of the selected stocks
            }
        }
    }
}

#region imports
using System;
using System.Collections;
using System.Collections.Generic;
using System.Linq;
using System.Globalization;
using System.Drawing;
using QuantConnect;
using QuantConnect.Algorithm.Framework;
using QuantConnect.Algorithm.Framework.Selection;
using QuantConnect.Algorithm.Framework.Alphas;
using QuantConnect.Algorithm.Framework.Portfolio;
using QuantConnect.Algorithm.Framework.Portfolio.SignalExports;
using QuantConnect.Algorithm.Framework.Execution;
using QuantConnect.Algorithm.Framework.Risk;
using QuantConnect.Algorithm.Selection;
using QuantConnect.Api;
using QuantConnect.Parameters;
using QuantConnect.Benchmarks;
using QuantConnect.Brokerages;
using QuantConnect.Configuration;
using QuantConnect.Util;
using QuantConnect.Interfaces;
using QuantConnect.Algorithm;
using QuantConnect.Indicators;
using QuantConnect.Data;
using QuantConnect.Data.Auxiliary;
using QuantConnect.Data.Consolidators;
using QuantConnect.Data.Custom;
using QuantConnect.DataSource;
using QuantConnect.Data.Fundamental;
using QuantConnect.Data.Market;
using QuantConnect.Data.Shortable;
using QuantConnect.Data.UniverseSelection;
using QuantConnect.Notifications;
using QuantConnect.Orders;
using QuantConnect.Orders.Fees;
using QuantConnect.Orders.Fills;
using QuantConnect.Orders.OptionExercise;
using QuantConnect.Orders.Slippage;
using QuantConnect.Orders.TimeInForces;
using QuantConnect.Python;
using QuantConnect.Scheduling;
using QuantConnect.Securities;
using QuantConnect.Securities.Equity;
using QuantConnect.Securities.Future;
using QuantConnect.Securities.Option;
using QuantConnect.Securities.Positions;
using QuantConnect.Securities.Forex;
using QuantConnect.Securities.Crypto;
using QuantConnect.Securities.CryptoFuture;
using QuantConnect.Securities.Interfaces;
using QuantConnect.Securities.Volatility;
using QuantConnect.Storage;
using QuantConnect.Statistics;
using QCAlgorithmFramework = QuantConnect.Algorithm.QCAlgorithm;
using QCAlgorithmFrameworkBridge = QuantConnect.Algorithm.QCAlgorithm;
using System;
using System.Collections.Generic;
using System.Linq;
using QuantConnect.Data.Fundamental;
using QuantConnect.Data.UniverseSelection;
using QuantConnect.Securities;

#endregion

namespace QuantConnect.Algorithm.Framework.Selection
{
    /// <summary>
    /// Defines the PSR universe as a universe selection model for framework algorithm
    /// </summary>
    public class PSRUniverseSelectionModel : FundamentalUniverseSelectionModel
    {
        // Constants defining the maximum number of symbols to select in coarse and fine selections
        private const int _numberOfSymbolsCoarse = 10000; // Maximum number of symbols selected during coarse selection
        private const int _numberOfSymbolsFine = 10000;   // Maximum number of symbols selected during fine selection

        // Keeps track of the last month to prevent unnecessary re-selection within the same month
        private int _lastMonth = -1; 
        
        // Dictionary storing the dollar volume for each symbol during coarse selection
        private readonly Dictionary<Symbol, double> _dollarVolumeBySymbol = new();

        /// <summary>
        /// Initializes a new default instance of the <see cref="PSRUniverseSelectionModel"/>
        /// </summary>
        public PSRUniverseSelectionModel()
            : base(true)
        {
        }

        /// <summary>
        /// Initializes a new instance of the <see cref="PSRUniverseSelectionModel"/>
        /// with custom universe settings.
        /// </summary>
        /// <param name="universeSettings">Universe settings define the properties to be applied to selected securities</param>
        public PSRUniverseSelectionModel(UniverseSettings universeSettings)
            : base(true, universeSettings)
        {
        }

        /// <summary>
        /// Performs coarse selection for the universe.
        /// Stocks must meet certain criteria such as having fundamental data, positive price and volume,
        /// and the highest dollar volume to be selected.
        /// </summary>
        public override IEnumerable<Symbol> SelectCoarse(QCAlgorithm algorithm, IEnumerable<CoarseFundamental> coarse)
        {
            // Avoid reselecting universe during the same month
            if (algorithm.Time.Month == _lastMonth)
            {
                return Universe.Unchanged;
            }

            // Sort the stocks by descending dollar volume and filter out invalid ones (missing fundamental data, negative volume or price)
            var sortedByDollarVolume =
                (from x in coarse
                 where x.HasFundamentalData && x.Volume > 0 && x.Price > 0
                 orderby x.DollarVolume descending
                 select x).Take(_numberOfSymbolsCoarse).ToList();

            // Clear the dollar volume data for the current selection
            _dollarVolumeBySymbol.Clear();
            foreach (var x in sortedByDollarVolume)
            {
                _dollarVolumeBySymbol[x.Symbol] = x.DollarVolume; // Store the dollar volume of each symbol
            }

            // If no symbols meet the coarse selection criteria, return unchanged universe
            if (_dollarVolumeBySymbol.Count == 0)
            {
                return Universe.Unchanged;
            }

            // Return the list of selected symbols for the coarse universe
            return _dollarVolumeBySymbol.Keys;
        }

        /// <summary>
        /// Performs fine selection for the universe.
        /// Additional filters are applied, such as ensuring the company is U.S.-based,
        /// traded on NYSE or NASDAQ, has been publicly listed for at least 6 months,
        /// and meets the price-to-sales ratio (PSR) criteria.
        /// </summary>
        public override IEnumerable<Symbol> SelectFine(QCAlgorithm algorithm, IEnumerable<FineFundamental> fine)
        {
            // Filter stocks that meet specific criteria
            var filteredFine =
                (from x in fine
                 where x.CompanyReference.CountryId == "USA" &&  // Company must be based in the USA
                       (x.CompanyReference.PrimaryExchangeID == "NYS" || x.CompanyReference.PrimaryExchangeID == "NAS") && // Traded on NYSE or NASDAQ
                       (algorithm.Time - x.SecurityReference.IPODate).Days > 180 &&  // Must have been listed for at least 6 months
                       x.ValuationRatios.PSRatio < 1 &&  // Price-to-sales ratio must be below 1
                        x.MarketCap > 5000000m
                 select x).ToList();

            var count = filteredFine.Count;

            // If no stocks meet the fine selection criteria, return unchanged universe
            if (count == 0)
            {
                return Universe.Unchanged;
            }

            // Update the last month variable to ensure universe is reselected for the next month
            _lastMonth = algorithm.Time.Month;

            // Calculate the percentage of stocks to be selected from each sector
            var percent = _numberOfSymbolsFine / (double)count;

            // Group stocks by sector and select the top stocks with the highest dollar volume within each sector
            var topFineBySector =
                (from x in filteredFine
                     // Group by industry sector
                 group x by x.CompanyReference.IndustryTemplateCode into g
                 let y = from item in g
                         orderby _dollarVolumeBySymbol[item.Symbol] descending
                         select item
                 let c = (int)Math.Ceiling(y.Count() * percent) // Select a percentage of stocks per sector
                 select new { g.Key, Value = y.Take(c) }
                 ).ToDictionary(x => x.Key, x => x.Value);

            // Return the selected stocks sorted by dollar volume, limiting to the top fine symbols
            return topFineBySector.SelectMany(x => x.Value)
                .OrderByDescending(x => _dollarVolumeBySymbol[x.Symbol]) // Sort by dollar volume
                .Take(_numberOfSymbolsFine) // Limit the number of fine symbols
                .Select(x => x.Symbol); // Return the symbols of the selected stocks
        }
    }
}