#region imports
using System;
using System.Collections.Generic;
using System.Linq;
using QuantConnect.Algorithm;
using QuantConnect.Data;
using QuantConnect.Data.Market;
using QuantConnect.Data.UniverseSelection;
using QuantConnect.Indicators;
using QuantConnect.Orders;
using QuantConnect.Securities;
#endregion

namespace QuantConnect.Algorithm.CSharp
{

public class ConnorsCrashAlgorithm : QCAlgorithm
{
    private readonly int _rsiPeriod = 3, _streakPeriod = 2, _pctRankPeriod = 100, _volaPeriod = 100, _maxShorts = 40;
    private readonly decimal _crsiEntry = 90m, _crsiExit = 30m;
    private Universe _universe;
    private readonly Dictionary<Symbol, (ConnorsRelativeStrengthIndex Crsi, CustomVolatility Vola)> _indicators = [];

    public override void Initialize()
    {
        SetStartDate(2024, 9, 1);
        SetEndDate(2024, 12, 31);
        SetCash(100_000);

        Settings.AutomaticIndicatorWarmUp = true;
        Settings.SeedInitialPrices = true;
        UniverseSettings.Resolution = Resolution.Daily;

        // Select liquid, tradeable-priced equities for the universe.
        _universe = AddUniverse(fundamentals =>
            fundamentals.Where(f => f.Price > 5 && f.DollarVolume > 1e6).Select(f => f.Symbol));

        Schedule.On(
            DateRules.EveryDay("SPY"),
            TimeRules.At(8, 0),
            Rebalance
        );
    }

    public override void OnSecuritiesChanged(SecurityChanges changes)
    {
        foreach (var security in changes.AddedSecurities)
        {
            // Attach ConnorsRSI indicator to each security.
            var crsi = CRSI(security.Symbol, _rsiPeriod, _streakPeriod, _pctRankPeriod);
            // Initialize and warm up the custom volatility indicator.
            var vola = new CustomVolatility(_volaPeriod);
            foreach (var bar in History<TradeBar>(security.Symbol, _volaPeriod + 1))
            {
                vola.Update(bar);
            }
            RegisterIndicator(security.Symbol, vola, Resolution.Daily);
            _indicators[security.Symbol] = (crsi, vola);
        }
        foreach (var security in changes.RemovedSecurities)
        {
            if (_indicators.TryGetValue(security.Symbol, out var pair))
            {
                DeregisterIndicator(pair.Crsi);
                DeregisterIndicator(pair.Vola);
                _indicators.Remove(security.Symbol);
            }
            Liquidate(security.Symbol);
        }
    }

    private void Rebalance()
    {
        if (!_universe.Selected.Any()) return;

        // Filter to securities with ready indicators.
        var securities = _universe.Selected
            .Where(s => Securities.ContainsKey(s) && _indicators.ContainsKey(s)
                        && _indicators[s].Crsi.IsReady && _indicators[s].Vola.IsReady)
            .Select(s => Securities[s])
            .ToList();

        // Filter to securities with above-threshold annualized volatility.
        var filterVola = securities.Where(s => _indicators[s.Symbol].Vola.Current.Value > 100).ToList();

        // Find currently invested short positions.
        var shortPositions = securities.Where(s => s.Holdings.IsShort).ToList();

        // Liquidate short positions when ConnorsRSI falls below exit threshold.
        foreach (var security in shortPositions)
        {
            if (_indicators[security.Symbol].Crsi.Current.Value < _crsiExit)
            {
                Liquidate(security.Symbol);
            }
        }

        // Exclude symbols with pending open orders.
        var pendingSymbols = Transactions.GetOpenOrderTickets()
            .Select(t => t.Symbol)
            .ToHashSet();

        // Find short entry candidates that are not currently invested (high volatility and high CRSI).
        var shortCandidates = filterVola
            .Where(s => _indicators[s.Symbol].Crsi.Current.Value > _crsiEntry
                        && !s.Holdings.Invested
                        && !pendingSymbols.Contains(s.Symbol))
            .OrderBy(s => _indicators[s.Symbol].Crsi.Current.Value)
            .ThenBy(s => _indicators[s.Symbol].Vola.Current.Value)
            .ToList();

        // Set union: liquidated positions are only counted once.
        var occupied = shortPositions.Select(s => s.Symbol).ToHashSet();
        occupied.UnionWith(pendingSymbols);
        var availableSlots = _maxShorts - occupied.Count;

        if (availableSlots <= 0 || !shortCandidates.Any()) return;

        var nOrders = Math.Min(shortCandidates.Count, availableSlots);
        var targetWeight = (decimal)(-1.0 / _maxShorts);

        foreach (var security in shortCandidates.TakeLast(nOrders))
        {
            var quantity = (int)CalculateOrderQuantity(security.Symbol, targetWeight);
            if (quantity != 0)
            {
                LimitOrder(security.Symbol, quantity, Math.Round(1.03m * security.Price, 3));
            }
        }
    }
}

public class CustomVolatility : TradeBarIndicator, IIndicatorWarmUpPeriodProvider
{
    private readonly RollingWindow<double> _window;

    public override bool IsReady => _window.IsReady;
    public int WarmUpPeriod => _window.Size;

    public CustomVolatility(int period) : base("CustomVolatility")
    {
        _window = new RollingWindow<double>(period);
    }

    protected override decimal ComputeNextValue(TradeBar input)
    {
        // Annualized log-return volatility.
        var price = (double)input.Value;
        if (price <= 0) return Current.Value;

        _window.Add(price);
        if (!_window.IsReady) return 0m;

        // Collect prices from oldest to newest to compute chronological log returns.
        var prices = _window.ToArray().Reverse().ToArray();
        var logDiffs = new double[prices.Length - 1];
        for (var i = 0; i < logDiffs.Length; i++)
        {
            logDiffs[i] = Math.Log(prices[i + 1] / prices[i]);
        }

        // Annualized standard deviation of log returns.
        var mean = logDiffs.Average();
        var variance = logDiffs.Select(d => (d - mean) * (d - mean)).Average();
        return (decimal)(Math.Sqrt(variance) * Math.Sqrt(252) * 100.0);
    }
}


}