/*
 * 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.
*/

using QuantConnect.Algorithm.Framework.Alphas;
using QuantConnect.Algorithm.Framework.Portfolio;
using QuantConnect.Algorithm.Framework.Selection;
using System;
using System.Collections.Generic;
using System.Linq;
using QuantConnect;
using QuantConnect.Algorithm.Framework.Risk;
using QuantConnect.Algorithm;
using QuantConnect.Indicators;
using QuantConnect.Data.Consolidators;
using QuantConnect.Data.UniverseSelection;
using QuantConnect.Securities;

namespace QuantConnect.Algorithm.CSharp
{
    /// <summary>
    /// Framework algorithm that uses the <see cref="EmaCrossUniverseSelectionModel"/> to
    /// select the universe based on a moving average cross.
    /// </summary>
    public class EmaCrossUniverseSelectionFrameworkAlgorithm : QCAlgorithm
    {
        public override void Initialize()
        {
            SetStartDate(2012, 12, 19);
            SetCash(100000);

            var fastSMAPeriod = 20;
            var slowSMAPeriod= 60;
            var atrPeriod = 14;
            var atrMultiplier = 3;
            Resolution resolution = Resolution.Hour;

            UniverseSettings.Leverage = 2.0m;
            UniverseSettings.Resolution = resolution;

            SetUniverseSelection(new ManualUniverseSelectionModel(QuantConnect.Symbol.Create("SPY", SecurityType.Equity, Market.USA)));
            SetAlpha(new EmaCrossAlphaModel(fastSMAPeriod,slowSMAPeriod,resolution));
            SetPortfolioConstruction(new EqualWeightingPortfolioConstructionModel());
            AddRiskManagement(new ATRTrailingStopRiskManagementModel(atrPeriod,atrMultiplier,resolution));
        }
    }

    public class ATRTrailingStopRiskManagementModel : RiskManagementModel
    {
        //declare variables
        private readonly int _period;
        private readonly decimal _multiplier;
        private readonly Resolution _resolution;
        private readonly Dictionary<Symbol, HoldingsState> _holdingsState = new Dictionary<Symbol, HoldingsState>();
        private readonly Dictionary<Symbol, SymbolData> _symbolData = new Dictionary<Symbol, SymbolData>();

        /// <summary>
        /// Initializes a new instance of the <see cref="ATRTrailingStopRiskManagementModel"/> class
        /// </summary>
        /// <param name="period">Period of the average true range indicator</param>
        /// <param name="multiplier">The number of average true ranges away from the price to calculate stop value</param>
        /// <param name="resolution">The resolution of ATR indicator</param>
        public ATRTrailingStopRiskManagementModel(int period = 14, decimal multiplier = 3m, Resolution resolution = Resolution.Minute)
        {
            _period = period;
            _multiplier = multiplier;
            _resolution = resolution;
        }

        /// <summary>
        /// Manages the algorithm's risk at each time step
        /// </summary>
        /// <param name="algorithm">The algorithm instance</param>
        /// <param name="targets">The current portfolio targets to be assessed for risk</param>
        public override IEnumerable<IPortfolioTarget> ManageRisk(QCAlgorithm algorithm, IPortfolioTarget[] targets)
        {
            foreach (var kvp in algorithm.Securities)
            {
                var symbol = kvp.Key;
                var security = kvp.Value;

                // fetch our symbol data containing our ATR indicator
                SymbolData data;
                if (!_symbolData.TryGetValue(symbol, out data))
                {
                    //exit loop if data not available
                    continue;
                }

                // Remove holdingsState if not invested
                if (!security.Invested)
                {
                    _holdingsState.Remove(symbol);
                    //no need to manage risk if not invested
                    continue;
                }

                PositionSide position = security.Holdings.IsLong ? PositionSide.Long : PositionSide.Short;
                HoldingsState state;

                // Add newly invested security (if doesn't exist) or reset holdings state (if position changed)
                if (!_holdingsState.TryGetValue(symbol, out state) || position != state.Position)
                {  
                    decimal initialStopValue = security.Close + ((_multiplier * data.ATR) * (position == PositionSide.Long ? 1 : -1));
                    _holdingsState[symbol] = state = new HoldingsState(position, initialStopValue);  
                }
            
                // liquidate if stop loss triggered
                if (StopLossTriggered(security, state))
                {
                    _holdingsState.Remove(symbol);
                    // liquidate    
                    yield return new PortfolioTarget(security.Symbol, 0);
                }
                
                //Calculate new stop value
                state.StopValue = position == PositionSide.Long
                    ? Math.Max(state.StopValue, security.Close - (_multiplier * data.ATR))
                    : Math.Min(state.StopValue, security.Close + (_multiplier * data.ATR));
                
            }  
        }

        /// <summary>
        /// Boolean function returns true if the stop loss value is triggered by the CLOSE of the bar
        /// </summary>
        /// <param name="security">Current security in the loop</param>
        /// <param name="state">State of the current trade in the loop</param>
        protected virtual bool StopLossTriggered (Security security, HoldingsState state )
        {
            return state.Position == PositionSide.Long 
                ? state.StopValue >= security.Low
                : state.StopValue <= security.High;
        }

        /// <summary>
        /// Helper class used to store holdings state for the <see cref="TrailingStopRiskManagementModel"/>
        /// in <see cref="ManageRisk"/>
        /// </summary>
        protected class HoldingsState
        {
            public PositionSide Position;
            public decimal StopValue;

            public HoldingsState(PositionSide position, decimal stopValue)
            {
                Position = position;
                StopValue = stopValue;
            }
        }

        /// <summary>
        /// Event fired each time the we add/remove securities from the data feed
        /// </summary>
        /// <param name="algorithm">The algorithm instance that experienced the change in securities</param>
        /// <param name="changes">The security additions and removals from the algorithm</param>
        public override void OnSecuritiesChanged(QCAlgorithm algorithm, SecurityChanges changes)
        {
            foreach (var added in changes.AddedSecurities)
            {
                // initialize new securities
                if (!_symbolData.ContainsKey(added.Symbol))
                {
                    _symbolData[added.Symbol] = new SymbolData(algorithm, added, _period, _resolution);
                }
            }

            foreach (var removed in changes.RemovedSecurities)
            {
                // clean up data from removed securities
                SymbolData data;
                if (_symbolData.TryGetValue(removed.Symbol, out data))
                {
                    if (IsSafeToRemove(algorithm, removed.Symbol))
                    {
                        _symbolData.Remove(removed.Symbol);
                        algorithm.SubscriptionManager.RemoveConsolidator(removed.Symbol, data.Consolidator);
                    }
                }
            }
        }

        /// <summary>
        /// Determines if it's safe to remove the associated symbol data
        /// </summary>
        protected virtual bool IsSafeToRemove(QCAlgorithm algorithm, Symbol symbol)
        {
            // confirm the security isn't currently a member of any universe
            return !algorithm.UniverseManager.Any(kvp => kvp.Value.ContainsMember(symbol));
        }

        /// <summary>
        /// Symbol Data for this Execution Model
        /// </summary>
        protected class SymbolData
        {
            /// <summary>
            /// Security
            /// </summary>
            public Security Security { get; }

            /// <summary>
            /// Average True Range
            /// </summary>
            public AverageTrueRange ATR { get; }

            /// <summary>
            /// Data Consolidator
            /// </summary>
            public IDataConsolidator Consolidator { get; }

            /// <summary>
            /// Initialize an instance of <see cref="SymbolData"/>
            /// </summary>
            /// <param name="algorithm">Algorithm for this security</param>
            /// <param name="security">The security we are using</param>
            /// <param name="period">Period of the ATR</param>
            /// <param name="resolution">Resolution for this symbol</param>
            public SymbolData(QCAlgorithm algorithm, Security security, int period, Resolution resolution)
            {
                Security = security;
                Consolidator = algorithm.ResolveConsolidator(security.Symbol, resolution);

                var atrName = algorithm.CreateIndicatorName(security.Symbol, "ATR" + period, resolution);
                ATR = new AverageTrueRange(atrName, period);
                algorithm.RegisterIndicator(security.Symbol, ATR, Consolidator);
                algorithm.WarmUpIndicator(security.Symbol, ATR, resolution);
                
            }
        }

    }

}