//                                               -*- C++ -*-
/**
 *  @file  KrigingAlgorithm.cxx
 *  @brief The class building gaussian process regression
 *
 *  Copyright 2005-2015 Airbus-EDF-IMACS-Phimeca
 *
 *  This library is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU Lesser General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This library 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 Lesser General Public License for more details.
 *
 *  You should have received a copy of the GNU Lesser General Public
 *  along with this library.  If not, see <http://www.gnu.org/licenses/>.
 *
 *  @author schueller
 */

#include "KrigingAlgorithm.hxx"
#include "PersistentObjectFactory.hxx"
#include "LinearNumericalMathFunction.hxx"
#include "SpecFunc.hxx"
#include "ProductCovarianceModel.hxx"
#include "KrigingEvaluation.hxx"
#include "KrigingGradient.hxx"
#include "CenteredFiniteDifferenceHessian.hxx"
#include "MethodBoundNumericalMathEvaluationImplementation.hxx"

BEGIN_NAMESPACE_OPENTURNS

CLASSNAMEINIT(KrigingAlgorithm);

static Factory<KrigingAlgorithm> RegisteredFactory("KrigingAlgorithm");


/* Default constructor */
KrigingAlgorithm::KrigingAlgorithm()
  : MetaModelAlgorithm()
  , inputSample_(0, 0)
  , normalizedInputSample_(0, 0)
  , inputTransformation_()
  , normalize_(false)
  , outputSample_(0, 0)
  , basis_(0)
  , covarianceModel_()
  , optimizer_()
  , optimizerProvided_(false)
  , outputIndex_(0)
  , beta_(0)
  , gamma_(0)
  , F_(0, 0)
  , result_()
{
  // Nothing to do
}


/* Constructor */
KrigingAlgorithm::KrigingAlgorithm(const NumericalSample & inputSample,
                                   const NumericalSample & outputSample,
                                   const Basis & basis,
                                   const CovarianceModel & covarianceModel,
                                   const Bool normalize)
  : MetaModelAlgorithm()
  , inputSample_(inputSample)
  , normalizedInputSample_(0, inputSample.getDimension())
  , inputTransformation_()
  , normalize_(normalize)
  , outputSample_(outputSample)
  , basis_(basis)
  , covarianceModel_(covarianceModel)
  , optimizer_()
  , optimizerProvided_(false)
  , outputIndex_(0)
  , beta_(0)
  , gamma_(0)
  , F_(0, 0)
  , result_()
{
  // Check the sample sizes
  if (inputSample_.getSize() != outputSample_.getSize())
    throw InvalidArgumentException(HERE) << "Input sample size (" << inputSample_.getSize() << ") does not match output sample size (" << outputSample_.getSize() << ").";
  // Check the covariance model
  // if (covarianceModel.getDimension() != 1)
  //   throw InvalidArgumentException(HERE) << "The covariance model dimension=" << covarianceModel.getDimension() << " is different from 1.";
  const UnsignedInteger dimension(inputSample.getDimension());
  if (dimension != covarianceModel.getSpatialDimension())
  {
    if (covarianceModel.getSpatialDimension() != 1) throw InvalidArgumentException(HERE) << "Input sample dimension (" << dimension << ") does not match covariance model spatial dimension (" << covarianceModel.getSpatialDimension() << ").";
    else covarianceModel_ = ProductCovarianceModel(ProductCovarianceModel::CovarianceModelCollection(dimension, covarianceModel));
  }
  else covarianceModel_ = covarianceModel;
  if (basis.getSize() > 0)
    if (basis[0].getOutputDimension() > 1) LOGWARN(OSS() << "Expected a basis of functions with output dimension=1, got output dimension=" << basis[0].getOutputDimension() << ". Only the first output component will be taken into account.");

  // Build a normalization function if needed
  if (normalize_)
  {
    const NumericalPoint mean(inputSample.computeMean());
    const NumericalPoint stdev(inputSample.computeStandardDeviationPerComponent());
    SquareMatrix linear(dimension);
    for (UnsignedInteger j = 0; j < dimension; ++ j)
    {
      linear(j, j) = 1.0;
      if (fabs(stdev[j]) > SpecFunc::MinNumericalScalar) linear(j, j) /= stdev[j];
    }
    const NumericalPoint zero(dimension);
    setInputTransformation(LinearNumericalMathFunction(mean, zero, linear));
  }
}


/* Constructor */
KrigingAlgorithm::KrigingAlgorithm(const NumericalSample & inputSample,
                                   const Distribution & inputDistribution,
                                   const NumericalSample & outputSample,
                                   const Basis & basis,
                                   const CovarianceModel & covarianceModel)
  : MetaModelAlgorithm()
  , inputSample_(inputSample)
  , normalizedInputSample_(0, inputSample.getDimension())
  , inputTransformation_()
  , normalize_(true)
  , outputSample_(outputSample)
  , basis_(basis)
  , covarianceModel_(covarianceModel)
  , optimizer_()
  , optimizerProvided_(false)
  , outputIndex_(0)
  , beta_(0)
  , gamma_(0)
  , F_(0, 0)
  , result_()
{
  // Check the sample sizes
  if (inputSample_.getSize() != outputSample_.getSize())
    throw InvalidArgumentException(HERE) << "Input sample size (" << inputSample_.getSize() << ") does not match output sample size (" << outputSample_.getSize() << ").";
  const UnsignedInteger dimension(inputSample.getDimension());
  if (dimension != covarianceModel.getSpatialDimension())
  {
    if (covarianceModel.getSpatialDimension() != 1) throw InvalidArgumentException(HERE) << "Input sample dimension (" << dimension << ") does not match covariance model spatial dimension (" << covarianceModel.getSpatialDimension() << ").";
    else covarianceModel_ = ProductCovarianceModel(ProductCovarianceModel::CovarianceModelCollection(dimension, covarianceModel));
  }
  else covarianceModel_ = covarianceModel;

  if (basis.getSize() > 0)
    if (basis[0].getOutputDimension() > 1) LOGWARN(OSS() << "Expected a basis of functions with output dimension=1, got output dimension=" << basis[0].getOutputDimension() << ". Only the first output component will be taken into account.");

  // use the isoprobabilistic transformation
  setInputTransformation(inputDistribution.getIsoProbabilisticTransformation());
}

/* Normalize the input sample */
void KrigingAlgorithm::normalizeInputSample()
{
  // Nothing to do if the sample has alredy been normalized
  if (normalizedInputSample_.getSize() != 0) return;
  // If we don't want to normalize the data
  if (!normalize_)
  {
    normalizedInputSample_ = inputSample_;
    return;
  }
  normalizedInputSample_ = inputTransformation_(inputSample_);
}

/* Compute the design matrix */
void KrigingAlgorithm::computeF()
{
  // Nothing to do if the design matrix has already been computed
  if (F_.getNbRows() != 0) return;
  // No early exit based on the sample/basis size as F_ must be initialized with the correct dimensions
  const UnsignedInteger size(normalizedInputSample_.getSize());
  const UnsignedInteger basisSize(basis_.getSize());
  // Compute F
  F_ = Matrix(size, basisSize);
  for (UnsignedInteger j = 0; j < basisSize; ++ j )
  {
    // Here we use potential parallelism in the evaluation of the basis functions
    const NumericalSample basisSample(basis_[j](normalizedInputSample_));
    for (UnsignedInteger i = 0; i < size; ++i) F_(i, j) = basisSample[i][0];
  }
}

/* Perform regression */
void KrigingAlgorithm::run()
{
  LOGINFO("normalize the data");
  normalizeInputSample();
  LOGINFO("Compute the design matrix");
  computeF();
  const UnsignedInteger outputDimension(outputSample_.getDimension());
  NumericalMathFunction::NumericalMathFunctionCollection processEvaluationCollection(outputDimension);
  // For each output component, build the associated kriging model
  LOGINFO("Loop over the dimensions");
  NumericalSample trendCoefficients(outputDimension, basis_.getSize());
  KrigingResult::CovarianceModelCollection conditionalCovarianceModels(outputDimension);
  NumericalSample covarianceCoefficients(outputDimension, inputSample_.getSize());
  for (outputIndex_ = 0; outputIndex_ < outputDimension; ++ outputIndex_)
  {
    LOGDEBUG(OSS() << "Output index=" << outputIndex_);
    NumericalPointWithDescription covarianceModelScale;
    // Here we compute both the meta-parameters and the kriging coefficients
    if (optimizerProvided_)
    {
      LOGINFO("Optimize the parameter of the marginal covariance model");
      covarianceModelScale = optimizeLogLikelihood();
    }
    else
    {
      LOGINFO("Extract the parameter of the marginal covariance model and compute the log-likelihood");
      // Here we simply read the meta-scales
      covarianceModelScale = covarianceModel_.getScale();
      // And we compute the kriging coefficients
      computeLogLikelihood(covarianceModelScale);
    }
    LOGINFO("Store the estimates");
    trendCoefficients[outputIndex_] = beta_;
    conditionalCovarianceModels[outputIndex_] = covarianceModel_;
    conditionalCovarianceModels[outputIndex_].setScale(covarianceModelScale);
    covarianceCoefficients[outputIndex_] = gamma_;
    LOGINFO("Build the output meta-model");
    NumericalMathFunction processEvaluation;
    processEvaluation.setEvaluation(new KrigingEvaluation(basis_, normalizedInputSample_, conditionalCovarianceModels[outputIndex_], beta_, gamma_));
    processEvaluation.setGradient(new KrigingGradient(basis_, normalizedInputSample_, conditionalCovarianceModels[outputIndex_], beta_, gamma_));
    processEvaluation.setHessian(new CenteredFiniteDifferenceHessian(ResourceMap::GetAsNumericalScalar( "CenteredFiniteDifferenceGradient-DefaultEpsilon" ), processEvaluation.getEvaluation()));
    processEvaluationCollection[outputIndex_] = processEvaluation;
  }
  NumericalMathFunction metaModel;
  // First build the meta-model on the transformed data
  if (outputDimension > 1) metaModel = NumericalMathFunction(processEvaluationCollection);
  else metaModel = NumericalMathFunction(processEvaluationCollection[0]);
  // Then add the transformation if needed
  if (normalize_) metaModel = NumericalMathFunction(metaModel, inputTransformation_);
  // compute residual, relative error
  const NumericalPoint outputVariance(outputSample_.computeVariance());
  const NumericalSample mY(metaModel(inputSample_));
  const NumericalPoint squaredResiduals((outputSample_ - mY).computeRawMoment(2));
  NumericalPoint residuals(outputDimension);
  NumericalPoint relativeErrors(outputDimension);

  const UnsignedInteger size(inputSample_.getSize());
  for ( UnsignedInteger outputIndex = 0; outputIndex < outputDimension; ++ outputIndex )
  {
    residuals[outputIndex] = sqrt(squaredResiduals[outputIndex] / size);
    relativeErrors[outputIndex] = squaredResiduals[outputIndex] / outputVariance[outputIndex];
  }
  result_ = KrigingResult(inputSample_, outputSample_, metaModel, residuals, relativeErrors, basis_, trendCoefficients, conditionalCovarianceModels, covarianceCoefficients);
}


NumericalScalar KrigingAlgorithm::computeLogLikelihood(const NumericalPoint & theta) const
{
  if (theta.getSize() != covarianceModel_.getScale().getSize())
    throw InvalidArgumentException(HERE) << "Error; Could not compute likelihood"
                                         << " the covariance model requires arguments of size" << covarianceModel_.getScale().getSize()
                                         << " here we got " << theta.getSize();
  LOGINFO(OSS(false) << "Extract marginal covariance model index=" << outputIndex_);
  CovarianceModel cov(covarianceModel_);
  LOGINFO(OSS(false) << "Replace the parameter value of the covariance model by " << theta);
  cov.setScale(theta);

  const UnsignedInteger size(inputSample_.getSize());

  LOGINFO("Discretize the covariance model into R");
  CovarianceMatrix R(cov.discretize(normalizedInputSample_));
  LOGINFO("Compute the Cholesky factor C of R");
  Bool continuationCondition(true);
  const NumericalScalar startingScaling(ResourceMap::GetAsNumericalScalar("KrigingAlgorithm-StartingScaling"));
  const NumericalScalar maximalScaling(ResourceMap::GetAsNumericalScalar("KrigingAlgorithm-MaximalScaling"));
  NumericalScalar cumulatedScaling(0.0);
  NumericalScalar scaling(startingScaling);
  TriangularMatrix C;
  while (continuationCondition && (cumulatedScaling < maximalScaling))
  {
    try
    {
      C = R.computeCholesky();
      continuationCondition = false;
    }
    // If it has not yet been computed, compute it and store it
    catch (InternalException & ex)
    {
      cumulatedScaling += scaling ;
      // Unroll the regularization to optimize the computation
      for (UnsignedInteger i = 0; i < size; ++i) R(i, i) += scaling;
      scaling *= 2.0;
    }
  }
  if (scaling >= maximalScaling)
    throw InvalidArgumentException(HERE) << "Error; Could not compute the Cholesky factor"
                                         << " Scaling up to "  << cumulatedScaling << " was not enough";
  if (cumulatedScaling > 0.0)
    LOGWARN(OSS() <<  "Warning! Scaling up to "  << cumulatedScaling << " was needed in order to get an admissible covariance. ");

  NumericalPoint y(size);
  for ( UnsignedInteger i = 0; i < size; ++ i ) y[i] = outputSample_[i][outputIndex_];
  LOGINFO("Solve C.psi = y");
  NumericalPoint psi(C.solveLinearSystem(y));
  NumericalPoint rho(psi);
  // If  trend to estimate
  if (basis_.getSize() > 0)
  {
    // Phi = C^{-1}F
    LOGINFO("Solve C.Phi = F");
    Matrix Phi(C.solveLinearSystem(F_));

    Matrix G;
    LOGINFO("Decompose Phi = Q.G with G triangular");
    Matrix Q(Phi.computeQR(G));
    LOGINFO("Solve Q.b = psi taking into account the orthogonality of Q");
    NumericalPoint b(Q.transpose() * psi);
    LOGINFO("Solve G.beta = b");
    beta_ = G.solveLinearSystem(b);

    LOGINFO("Compute rho = psi - Phi.beta");
    rho -= Phi * beta_;
  }
  LOGINFO("Solve C^t.gamma = rho");
  gamma_ = C.transpose().solveLinearSystem(rho);

  LOGINFO("Compute log(|det(R)|)");
  NumericalScalar sigma2(rho.normSquare() / size);
  if (sigma2 <= 0) return SpecFunc::MaxNumericalScalar;
  NumericalScalar logDetR(0.0);
  for ( UnsignedInteger i = 0; i < size; ++ i )
  {
    const NumericalScalar cii(C(i, i));
    if (cii <= 0.0) return SpecFunc::MaxNumericalScalar;
    logDetR += log(cii);
  }
  const NumericalScalar logLikelihood(-((size - F_.getNbColumns()) * log(sigma2) + 2.0 * logDetR));
  LOGINFO(OSS(false) << "Compute the unbiased log-likelihood=" << logLikelihood);

  return logLikelihood;
}


NumericalPoint KrigingAlgorithm::optimizeLogLikelihood()
{
  // initial guess
  const NumericalPoint initialTheta(covarianceModel_.getScale());
  const NumericalScalar initialLogLikelihood(computeLogLikelihood(initialTheta));
  LOGINFO(OSS() << "Initial theta=" << initialTheta << ", log-likelihood=" << initialLogLikelihood);

  BoundConstrainedAlgorithm optimizer(optimizer_);
  optimizer.setObjectiveFunction(getLogLikelihoodFunction(outputIndex_));
  optimizer.setOptimizationProblem(BoundConstrainedAlgorithmImplementationResult::MAXIMIZATION);
  optimizer.setStartingPoint(initialTheta);
  optimizer.run();

  // check result
  const NumericalScalar optimizedLogLikelihood(optimizer.getResult().getOptimalValue());
  const NumericalPoint optimizedTheta(optimizer.getResult().getOptimizer());
  LOGINFO(OSS() << "Optimized theta=" << optimizedTheta << ", log-likelihood=" << optimizedLogLikelihood);
  const NumericalPoint finalTheta(optimizedLogLikelihood > initialLogLikelihood ? optimizedTheta : initialTheta);
  // the last optimized point is not necessarily the last evaluated, so update intermediate results
  const NumericalScalar finalLogLikelihood(computeLogLikelihood(finalTheta));
  LOGINFO(OSS() << "Final theta=" << finalTheta << ", log-likelihood=" << finalLogLikelihood);

  return finalTheta;
}


void KrigingAlgorithm::setOptimizer(const BoundConstrainedAlgorithm& optimizer)
{
  optimizer_ = optimizer;
  optimizerProvided_ = true;
}


BoundConstrainedAlgorithm KrigingAlgorithm::getOptimizer() const
{
  return optimizer_;
}


void KrigingAlgorithm::setInputTransformation(const NumericalMathFunction& inputTransformation)
{
  if (inputTransformation.getInputDimension() != inputSample_.getDimension()) throw InvalidDimensionException(HERE)
        << "Input dimension of the transformation (" << inputTransformation.getInputDimension() << ") should match input sample dimension (" << inputSample_.getDimension() << ")";
  if (inputTransformation.getOutputDimension() != inputSample_.getDimension()) throw InvalidDimensionException(HERE)
        << "Output dimension of the transformation (" << inputTransformation.getOutputDimension() << ") should match output sample dimension (" << inputSample_.getDimension() << ")";
  inputTransformation_ = inputTransformation;
}

NumericalMathFunction KrigingAlgorithm::getInputTransformation() const
{
  if (!normalize_)
  {
    const UnsignedInteger dimension(inputSample_.getDimension());
    return LinearNumericalMathFunction(NumericalPoint(dimension), NumericalPoint(dimension), IdentityMatrix(dimension));
  }
  return inputTransformation_;
}


/* Virtual constructor */
KrigingAlgorithm * KrigingAlgorithm::clone() const
{
  return new KrigingAlgorithm(*this);
}


/* String converter */
String KrigingAlgorithm::__repr__() const
{
  return OSS() << "class=" << getClassName();
}


NumericalSample KrigingAlgorithm::getInputSample() const
{
  return inputSample_;
}


NumericalSample KrigingAlgorithm::getOutputSample() const
{
  return outputSample_;;
}


KrigingResult KrigingAlgorithm::getResult()
{
  return result_;
}


NumericalMathFunction KrigingAlgorithm::getLogLikelihoodFunction(const UnsignedInteger outputIndex)
{
  outputIndex_ = outputIndex;
  LOGINFO("normalize the data");
  normalizeInputSample();
  LOGINFO("Compute the design matrix");
  computeF();
  return bindMethod <KrigingAlgorithm, NumericalScalar, NumericalPoint> ( *this, &KrigingAlgorithm::computeLogLikelihood, covarianceModel_.getScale().getSize(), 1 );
}


/* Method save() stores the object through the StorageManager */
void KrigingAlgorithm::save(Advocate & adv) const
{
  MetaModelAlgorithm::save(adv);
  adv.saveAttribute( "inputSample_", inputSample_ );
  adv.saveAttribute( "inputTransformation_", inputTransformation_ );
  adv.saveAttribute( "normalize_", normalize_ );
  adv.saveAttribute( "outputSample_", outputSample_ );
  adv.saveAttribute( "basis_", basis_ );
  adv.saveAttribute( "covarianceModel_", covarianceModel_ );
  adv.saveAttribute( "optimizer_", optimizer_ );
  adv.saveAttribute( "optimizerProvided_", optimizerProvided_ );
  adv.saveAttribute( "result_", result_ );
}


/* Method load() reloads the object from the StorageManager */
void KrigingAlgorithm::load(Advocate & adv)
{
  MetaModelAlgorithm::load(adv);
  adv.loadAttribute( "inputSample_", inputSample_ );
  adv.loadAttribute( "inputTransformation_", inputTransformation_ );
  adv.loadAttribute( "normalize_", normalize_ );
  adv.loadAttribute( "outputSample_", outputSample_ );
  adv.loadAttribute( "basis_", basis_ );
  adv.loadAttribute( "covarianceModel_", covarianceModel_ );
  adv.loadAttribute( "optimizer_", optimizer_ );
  adv.loadAttribute( "optimizerProvided_", optimizerProvided_ );
  adv.loadAttribute( "result_", result_ );
}

END_NAMESPACE_OPENTURNS
