#region License Information
/* HeuristicLab
 * Copyright (C) 2002-2016 Heuristic and Evolutionary Algorithms Laboratory (HEAL)
 *
 * This file is part of HeuristicLab.
 *
 * HeuristicLab is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * HeuristicLab is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with HeuristicLab. If not, see <http://www.gnu.org/licenses/>.
 */
#endregion

using System;
using System.Collections.Generic;
using System.Linq;
using HeuristicLab.Common;
using HeuristicLab.Core;
using HeuristicLab.Persistence.Default.CompositeSerializers.Storable;
using HeuristicLab.Problems.DataAnalysis;
using MathNet.Numerics.Data.Matlab;
using MathNet.Numerics.LinearAlgebra;
using MathNet.Numerics.LinearAlgebra.Double;
using MathNet.Numerics.LinearAlgebra.Factorization;

namespace HeuristicLab.Algorithms.DataAnalysis.Experimental {
  // TODO: scale y
  // TODO: remove dependence of scaling and export scaling parameters
  // TODO: export / import all relevant data
  [StorableClass]
  [Item("GaussianProcessModelMKL", "Represents a Gaussian process posterior.")]
  public sealed class GaussianProcessModelMKL : RegressionModel, IGaussianProcessModel {
    public override IEnumerable<string> VariablesUsedForPrediction {
      get { return allowedInputVariables; }
    }

    [Storable]
    private double negativeLogLikelihood;
    public double NegativeLogLikelihood {
      get { return negativeLogLikelihood; }
    }

    [Storable]
    private double negativeLooPredictiveProbability;
    public double NegativeLooPredictiveProbability {
      get { return negativeLooPredictiveProbability; }
    }

    [Storable]
    private double[] hyperparameterGradients;
    public double[] HyperparameterGradients {
      get {
        var copy = new double[hyperparameterGradients.Length];
        Array.Copy(hyperparameterGradients, copy, copy.Length);
        return copy;
      }
    }

    [Storable]
    private ICovarianceFunction covarianceFunction;
    public ICovarianceFunction CovarianceFunction {
      get { return covarianceFunction; }
    }
    [Storable]
    private IMeanFunction meanFunction;
    public IMeanFunction MeanFunction {
      get { return meanFunction; }
    }

    [Storable]
    private string[] allowedInputVariables;
    public string[] AllowedInputVariables {
      get { return allowedInputVariables; }
    }

    [Storable]
    private Vector<double> alpha;
    [Storable]
    private double sqrSigmaNoise;
    public double SigmaNoise {
      get { return Math.Sqrt(sqrSigmaNoise); }
    }

    [Storable]
    private double[] meanParameter;
    [Storable]
    private double[] covarianceParameter;

    private Matrix<double> l;
    private double[,] x; // scaled training dataset, used to be storable in previous versions (is calculated lazily now)


    [Storable]
    private IDataset trainingDataset; // it is better to store the original training dataset completely because this is more efficient in persistence
    [Storable]
    private int[] trainingRows;

    [Storable]
    private Scaling inputScaling;


    [StorableConstructor]
    private GaussianProcessModelMKL(bool deserializing) : base(deserializing) { }
    private GaussianProcessModelMKL(GaussianProcessModelMKL original, Cloner cloner)
      : base(original, cloner) {
      this.meanFunction = cloner.Clone(original.meanFunction);
      this.covarianceFunction = cloner.Clone(original.covarianceFunction);
      if (original.inputScaling != null)
        this.inputScaling = cloner.Clone(original.inputScaling);
      this.trainingDataset = cloner.Clone(original.trainingDataset);
      this.negativeLogLikelihood = original.negativeLogLikelihood;
      this.negativeLooPredictiveProbability = original.negativeLooPredictiveProbability;
      this.sqrSigmaNoise = original.sqrSigmaNoise;
      if (original.meanParameter != null) {
        this.meanParameter = (double[])original.meanParameter.Clone();
      }
      if (original.covarianceParameter != null) {
        this.covarianceParameter = (double[])original.covarianceParameter.Clone();
      }

      // shallow copies of arrays because they cannot be modified
      this.trainingRows = original.trainingRows;
      this.allowedInputVariables = original.allowedInputVariables;
      this.alpha = original.alpha;
      this.l = original.l;
      this.x = original.x;
    }
    public GaussianProcessModelMKL(IDataset ds, string targetVariable, IEnumerable<string> allowedInputVariables, IEnumerable<int> rows,
      IEnumerable<double> hyp, IMeanFunction meanFunction, ICovarianceFunction covarianceFunction,
      bool scaleInputs = true)
      : base(targetVariable) {
      // MathNet.Numerics.Control.UseNativeMKL();
      // MathNet.Numerics.Control.UseSingleThread();
      // this.Description += Control.LinearAlgebraProvider.ToString();
      this.name = ItemName;
      this.description = ItemDescription;
      this.meanFunction = (IMeanFunction)meanFunction.Clone();
      this.covarianceFunction = (ICovarianceFunction)covarianceFunction.Clone();
      this.allowedInputVariables = allowedInputVariables.ToArray();


      int nVariables = this.allowedInputVariables.Length;
      meanParameter = hyp
        .Take(this.meanFunction.GetNumberOfParameters(nVariables))
        .ToArray();

      covarianceParameter = hyp.Skip(this.meanFunction.GetNumberOfParameters(nVariables))
                                             .Take(this.covarianceFunction.GetNumberOfParameters(nVariables))
                                             .ToArray();
      sqrSigmaNoise = Math.Exp(2.0 * hyp.Last());
      CalculateModel(ds, rows, scaleInputs);
    }

    private void CalculateModel(IDataset ds, IEnumerable<int> rows, bool scaleInputs = true) {
      this.trainingDataset = (IDataset)ds.Clone();
      this.trainingRows = rows.ToArray();
      this.inputScaling = scaleInputs ? new Scaling(ds, allowedInputVariables, rows) : null;

      x = GetData(ds, this.allowedInputVariables, this.trainingRows, this.inputScaling);

      IEnumerable<double> y;
      y = ds.GetDoubleValues(TargetVariable, rows);

      int n = x.GetLength(0);

      var columns = Enumerable.Range(0, x.GetLength(1)).ToArray();
      // calculate cholesky decomposed (lower triangular) covariance matrix
      var cov = covarianceFunction.GetParameterizedCovarianceFunction(covarianceParameter, columns);
      var chol = CalculateL(x, cov, sqrSigmaNoise);
      this.l = chol.Factor;

      // calculate mean
      var mean = meanFunction.GetParameterizedMeanFunction(meanParameter, columns);
      double[] m = Enumerable.Range(0, x.GetLength(0))
        .Select(r => mean.Mean(x, r))
        .ToArray();


      // solve for alpha
      Vector<double> ym = DenseVector.OfEnumerable(y.Zip(m, (a, b) => a - b));


      alpha = chol.Solve(ym);
      alpha = alpha * 1.0 / sqrSigmaNoise;

      // calculate sum of diagonal elements for likelihood
      double diagSum = chol.DeterminantLn;
      negativeLogLikelihood = 0.5 * ym.DotProduct(alpha) + 0.5 * diagSum + (n / 2.0) * Math.Log(2.0 * Math.PI * sqrSigmaNoise);

      // derivatives
      int nAllowedVariables = x.GetLength(1);

      Matrix<double> lCopy;

      lCopy = chol.Solve(DenseMatrix.CreateIdentity(n));

      // LOOCV log predictive probability (GPML page 116 and 117)
      var sumLoo = 0.0;
      var ki = new DenseVector(n);
      for (int i = 0; i < n; i++) {
        for (int j = 0; j < n; j++) ki[j] = cov.Covariance(x, i, j);
        var yi = ki.DotProduct(alpha);
        var yi_loo = yi - alpha[i] / lCopy[i, i] / sqrSigmaNoise;
        var s2_loo = sqrSigmaNoise / lCopy[i, i];
        var err = ym[i] - yi_loo;
        var nll_loo = Math.Log(s2_loo) + err * err / s2_loo;
        sumLoo += nll_loo;
      }
      sumLoo += n * Math.Log(2 * Math.PI);
      negativeLooPredictiveProbability = 0.5 * sumLoo;

      for (int i = 0; i < n; i++) {
        for (int j = 0; j <= i; j++)
          lCopy[i, j] = lCopy[i, j] / sqrSigmaNoise - alpha[i] * alpha[j];
      }

      double noiseGradient = sqrSigmaNoise * Enumerable.Range(0, n).Select(i => lCopy[i, i]).Sum();

      double[] meanGradients = new double[meanFunction.GetNumberOfParameters(nAllowedVariables)];
      for (int k = 0; k < meanGradients.Length; k++) {
        var meanGrad = new DenseVector(alpha.Count);
        for (int g = 0; g < meanGrad.Count; g++)
          meanGrad[g] = mean.Gradient(x, g, k);

        meanGradients[k] = -meanGrad.DotProduct(alpha);
      }

      double[] covGradients = new double[covarianceFunction.GetNumberOfParameters(nAllowedVariables)];
      if (covGradients.Length > 0) {
        for (int i = 0; i < n; i++) {
          for (int j = 0; j < i; j++) {
            var g = cov.CovarianceGradient(x, i, j);
            for (int k = 0; k < covGradients.Length; k++) {
              covGradients[k] += lCopy[i, j] * g[k];
            }
          }

          var gDiag = cov.CovarianceGradient(x, i, i);
          for (int k = 0; k < covGradients.Length; k++) {
            // diag
            covGradients[k] += 0.5 * lCopy[i, i] * gDiag[k];
          }
        }
      }

      hyperparameterGradients =
        meanGradients
        .Concat(covGradients)
        .Concat(new double[] { noiseGradient }).ToArray();

    }

    private static double[,] GetData(IDataset ds, IEnumerable<string> allowedInputs, IEnumerable<int> rows, Scaling scaling) {
      if (scaling != null) {
        // BackwardsCompatibility3.3
        #region Backwards compatible code, remove with 3.4
        // TODO: completely remove Scaling class
        List<string> variablesList = allowedInputs.ToList();
        List<int> rowsList = rows.ToList();

        double[,] matrix = new double[rowsList.Count, variablesList.Count];

        int col = 0;
        foreach (string column in variablesList) {
          var values = scaling.GetScaledValues(ds, column, rowsList);
          int row = 0;
          foreach (var value in values) {
            matrix[row, col] = value;
            row++;
          }
          col++;
        }
        return matrix;
        #endregion
      } else {
        return ds.ToArray(allowedInputs, rows);
      }
    }

    private static Cholesky<double> CalculateL(double[,] x, ParameterizedCovarianceFunction cov, double sqrSigmaNoise) {
      int n = x.GetLength(0);
      var l = new DenseMatrix(n, n);

      // calculate covariances
      for (int i = 0; i < n; i++) {
        for (int j = i; j < n; j++) {
          l[j, i] = l[i, j] = cov.Covariance(x, i, j) / sqrSigmaNoise;
          if (j == i) l[j, i] += 1.0;
        }
      }

      // cholesky decomposition
      return l.Cholesky();
    }


    public override IDeepCloneable Clone(Cloner cloner) {
      return new GaussianProcessModelMKL(this, cloner);
    }

    // is called by the solution creator to set all parameter values of the covariance and mean function
    // to the optimized values (necessary to make the values visible in the GUI)
    public void FixParameters() {
      covarianceFunction.SetParameter(covarianceParameter);
      meanFunction.SetParameter(meanParameter);
      covarianceParameter = new double[0];
      meanParameter = new double[0];
    }

    #region IRegressionModel Members
    public override IEnumerable<double> GetEstimatedValues(IDataset dataset, IEnumerable<int> rows) {
      return GetEstimatedValuesHelper(dataset, rows);
    }
    public override IRegressionSolution CreateRegressionSolution(IRegressionProblemData problemData) {
      return new GaussianProcessRegressionSolution(this, new RegressionProblemData(problemData));
    }
    #endregion


    private IEnumerable<double> GetEstimatedValuesHelper(IDataset dataset, IEnumerable<int> rows) {
      if (x == null) {
        x = GetData(trainingDataset, allowedInputVariables, trainingRows, inputScaling);
      }
      int n = x.GetLength(0);

      double[,] newX = GetData(dataset, allowedInputVariables, rows, inputScaling);
      int newN = newX.GetLength(0);

      var columns = Enumerable.Range(0, newX.GetLength(1)).ToArray();
      var mean = meanFunction.GetParameterizedMeanFunction(meanParameter, columns);
      var ms = Enumerable.Range(0, newX.GetLength(0))
      .Select(r => mean.Mean(newX, r))
      .ToArray();
      var cov = covarianceFunction.GetParameterizedCovarianceFunction(covarianceParameter, columns);
      for (int i = 0; i < newN; i++) {
        var Ks = new DenseVector(n);
        for (int j = 0; j < n; j++) {
          Ks[j] = cov.CrossCovariance(x, newX, j, i);
        }
        yield return ms[i] + Ks.DotProduct(alpha);
      }
    }

    public IEnumerable<double> GetEstimatedVariances(IDataset dataset, IEnumerable<int> rows) {
      if (x == null) {
        x = GetData(trainingDataset, allowedInputVariables, trainingRows, inputScaling);
      }
      int n = x.GetLength(0);

      var newX = GetData(dataset, allowedInputVariables, rows, inputScaling);
      int newN = newX.GetLength(0);
      if (newN == 0) return Enumerable.Empty<double>();

      var kss = new DenseVector(newN);
      Matrix<double> sWKs = new DenseMatrix(n, newN);
      var columns = Enumerable.Range(0, newX.GetLength(1)).ToArray();
      var cov = covarianceFunction.GetParameterizedCovarianceFunction(covarianceParameter, columns);

      if (l == null) {
        l = CalculateL(x, cov, sqrSigmaNoise).Factor;
      }

      // for stddev 
      for (int i = 0; i < newN; i++)
        kss[i] = cov.Covariance(newX, i, i);

      for (int i = 0; i < newN; i++) {
        for (int j = 0; j < n; j++) {
          sWKs[j, i] = cov.CrossCovariance(x, newX, j, i) / Math.Sqrt(sqrSigmaNoise);
        }
      }

      sWKs = l.Solve(sWKs);

      for (int i = 0; i < newN; i++) {
        var col = sWKs.Column(i);
        var sumV = col.DotProduct(col);
        kss[i] += sqrSigmaNoise; // kss is V(f), add noise variance of predictive distibution to get V(y)
        kss[i] -= sumV;
        if (kss[i] < 0) kss[i] = 0;
      }
      return kss;
    }


    public void Export(string fileName) {
      MatlabWriter.Write<double>(fileName,
        new Matrix<double>[] { DenseMatrix.OfArray(x), l, alpha.ToRowMatrix()},
        new string[] { "x", "l", "alpha", }
        );
    }
  }
}