namespace QuantConnect
{
    internal enum MarketRegime
    {
        Bearish = -1,
        Uncertain = 0,
        Bullish = 1
    }

    public class RsiMaStrategy : QCAlgorithm
    {
        private Security _appl;
        private SimpleMovingAverage _smoothedRSI;
        private RelativeStrengthIndex _rsi;
        private RollingWindow<IndicatorDataPoint> _rsiLag;

        private CompositeIndicator<IndicatorDataPoint> _marketRegimeIndicator;

        private const decimal _tol = 0.0001m;

        public override void Initialize()
        {
            SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage);

            // backtest parameters
            SetStartDate(2015, 01, 01);
            SetEndDate(DateTime.Today);

            // cash allocation
            SetCash(10000);

            _appl = AddEquity("AAPL");
            _rsi = new RelativeStrengthIndex(14);
            _smoothedRSI = new SimpleMovingAverage(4).Of(_rsi);
            _rsiLag = new RollingWindow<IndicatorDataPoint>(2);

            var fiveMinuteConsolidator = new TradeBarConsolidator(TimeSpan.FromMinutes(5));
            fiveMinuteConsolidator.DataConsolidated += FiveMinuteConsolidator_DataConsolidated;
            SubscriptionManager.AddConsolidator(_appl.Symbol, fiveMinuteConsolidator);

            _marketRegimeIndicator = SMA(_appl.Symbol, 5, Resolution.Daily).Minus(SMA(_appl.Symbol, 200, Resolution.Daily));

            Schedule.On(DateRules.EveryDay(), TimeRules.BeforeMarketClose(_appl.Symbol, 4), () =>
            {
                CloseOpenPositions(_appl);
            });

            SetWarmUp(TimeSpan.FromDays(100));
        }

        private void CloseOpenPositions(Security security)
        {
            if (security.Invested)
            {
                Liquidate(security.Symbol);
            }
        }

        private MarketRegime GetMarketRegime()
        {
            MarketRegime marketregime = MarketRegime.Uncertain;
            if (Math.Abs(_marketRegimeIndicator) > _tol)
            {
                marketregime = (MarketRegime)Math.Sign(_marketRegimeIndicator);
            }
            return marketregime;
        }

        private void FiveMinuteConsolidator_DataConsolidated(object sender, TradeBar e)
        {
            _rsi.Update(new IndicatorDataPoint
            {
                Time = e.Time,
                Value = e.Value
            });
            _rsiLag.Add(_rsi.Current);

            if (!_marketRegimeIndicator.IsReady || !_rsiLag.IsReady) return; 

            if (_appl.Invested)
            {
                ExitLogic();
            }
            else
            {
                EntryLogic();
            }
        }

        private void ExitLogic()
        {
            var bullishRSIExit = _rsiLag[1] > 55 && _rsiLag[0] < 55;
            var bearishRSIExit = _rsiLag[1] < 45 && _rsiLag[0] > 45;
            if (   (GetMarketRegime() == MarketRegime.Bullish && bullishRSIExit)
                || (GetMarketRegime() == MarketRegime.Bearish && bearishRSIExit))
            {
                Liquidate(_appl.Symbol);
            }
        }

        private void EntryLogic()
        {
            var bullishRSIEntry = _rsiLag[1] < 30 && _rsiLag[0] > 30;
            var bearishRSIEntry = _rsiLag[1] > 70 && _rsiLag[0] < 70;
            if (GetMarketRegime() == MarketRegime.Bullish && bullishRSIEntry)
            {
                SetHoldings(_appl.Symbol, 1);
            }
            else if (GetMarketRegime() == MarketRegime.Bearish && bearishRSIEntry)
            {
                SetHoldings(_appl.Symbol, -1);
            }
        }

        public override void OnData(Slice data)
        {
        }
    }
}