using System;
using System.Linq;
using MathNet.Numerics;

namespace QuantConnect
{
    public class GoldStatArbAlgorithm : QCAlgorithm
    {
        private PairCalculation _pair;

        public override void Initialize()
        {
            SetStartDate(2006, 05, 23);  //Set Start Date
            SetEndDate(2007, 11, 30);    //Set End Date
            SetCash(100000);             //Set Strategy Cash
    
            var gld = AddEquity("GLD", Resolution.Daily);
            var gdx = AddEquity("GDX", Resolution.Daily);

            _pair = new PairCalculation();
        }
        
        public override void OnData(Slice data)
        {
        	if (!data.ContainsKey("GLD") || !data.ContainsKey("GDX")) return;
        	
            if (!_pair.Update(data["GLD"], data["GDX"])) return;

            var zScore = _pair.ZScore;
            if (Math.Abs(zScore) > .5m && !Portfolio.Invested)
            {
                MarketOrder("GLD", -1000 * Math.Sign(zScore));
                MarketOrder("GDX", (int)Math.Round(1000 * _pair.Beta));
            }

            if (Portfolio.Invested)
            {
                if (zScore * Math.Sign(Portfolio["GLD"].Quantity) < 0)
                {
                    Liquidate();
                }
            }
        }
    }

    public class PairCalculation
    {
        private int _tradingSize = 30;
        private int _trainingSize = 90;
        private RollingWindow<IBaseData> _xWindow;
        private RollingWindow<IBaseData> _yWindow;

        private int _period = 5;
        private SimpleMovingAverage _sma;
        private StandardDeviation _std;

        public decimal ZScore { get; private set; }
		public decimal Beta {  get; private set; }

        public PairCalculation()
        {
            _xWindow = new RollingWindow<IBaseData>(_trainingSize);
            _yWindow = new RollingWindow<IBaseData>(_trainingSize);

            _sma = new SimpleMovingAverage(_period);
            _std = new StandardDeviation(_period);
        }

        public bool Update(IBaseData dataOne, IBaseData dataTwo)
        {
            var time = dataOne.EndTime;
            var spread = dataOne.Price - Beta * dataTwo.Price;

            _xWindow.Add(dataOne);
            _yWindow.Add(dataTwo);

            if (!_xWindow.IsReady || !_yWindow.IsReady) return false;

            //
            if (_xWindow.Samples == 90)//_xWindow.Samples % _tradingSize == 0)
            {
                var x = _xWindow.OrderBy(t => t.EndTime).Select(a => (double)a.Price).ToArray();
                var y = _yWindow.OrderBy(t => t.EndTime).Select(a => (double)a.Price).ToArray();

                var ols = Fit.Line(x, y);
                Beta = 1 / (decimal)ols.Item2;

                var spreadWindow = Enumerable.Zip(_xWindow, _yWindow,
                    (a, b) => new IndicatorDataPoint(a.EndTime, a.Price - Beta * b.Price))
                    .OrderBy(t => t.EndTime);

                _sma.Reset();
                _std.Reset();

                foreach (var input in spreadWindow)
                {
                    _sma.Update(input);
                    _std.Update(input);
                }

                ZScore = (spreadWindow.Last() - _sma) / _std;
            }
            else
            {
                _sma.Update(time, spread);
                _std.Update(time, spread);

                ZScore = (spread - _sma) / _std;
            }
            return true;
        }
    }
}