/* randist/wishart.c
 *
 * Copyright (C) 2017 Timothée Flutre
 *
 * This program 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.
 *
 * This program 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 this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
 */

#include <config.h>
#include <math.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>
#include <gsl/gsl_vector.h>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_blas.h>
#include <gsl/gsl_linalg.h>
#include <gsl/gsl_sf_gamma.h>


/* Generate a random matrix from a Wishart distribution using the Bartlett
 * decomposition, following Smith and Hocking, Journal of the Royal Statistical
 * Society. Series C (Applied Statistics), Vol. 21, No. 3 (1972), pp. 341-345.
 *
 * df      degrees of freedom
 * L       matrix resulting from the Cholesky decomposition of
 *         the scale matrix V = L L^T (dimension d x d)
 * result  output matrix (dimension d x d)
 * work    matrix used for intermediate computations (dimension d x d)
 */
int
gsl_ran_wishart (const gsl_rng * r,
                 const double df,
                 const gsl_matrix * L,
                 gsl_matrix * result,
                 gsl_matrix * work)
{
  if (L->size1 != L->size2)
    {
      GSL_ERROR("L should be a square matrix", GSL_ENOTSQR);
    }
  else if (result->size1 != result->size2)
    {
      GSL_ERROR("result should be a square matrix", GSL_ENOTSQR);
    }
  else if (work->size1 != work->size2)
    {
      GSL_ERROR("work should be a square matrix", GSL_ENOTSQR);
    }
  else if (result->size1 != L->size1)
    {
      GSL_ERROR("incompatible dimensions of result matrix", GSL_EBADLEN);
    }
  else if (work->size1 != L->size1)
    {
      GSL_ERROR("incompatible dimensions of work matrix", GSL_EBADLEN);
    }
  else if (df <= L->size1 - 1)
    {
      GSL_ERROR("incompatible degrees of freedom", GSL_EDOM);
    }
  else
    {
      /* result: X = L A A^T L^T */

      size_t d = L->size1, i, j;

      /* insure the upper part of A is zero before filling its lower part */
      gsl_matrix_set_zero(work);
      for (i = 0; i < d; ++i)
        {
          gsl_matrix_set(work, i, i, sqrt(gsl_ran_chisq(r, df - i)));

          for (j = 0; j < i; ++j)
            {
              gsl_matrix_set(work, i, j, gsl_ran_ugaussian(r));
            }
        }

      /* compute L * A */
      gsl_blas_dtrmm(CblasLeft, CblasLower, CblasNoTrans, CblasNonUnit, 1.0,
                     L, work);

      /* compute (L * A) * (L * A)^T */
      gsl_blas_dsyrk(CblasUpper, CblasNoTrans, 1.0, work, 0.0, result);
      for (i = 0; i < d; ++i)
        {
          for (j = 0; j < i; ++j)
            {
              gsl_matrix_set(result, i, j, gsl_matrix_get(result, j, i));
            }
        }

      return GSL_SUCCESS;
    }
}

/* Compute the log of the probability density function at a given quantile
 * matrix for a Wishart distribution using the Cholesky decomposition of the
 * scale matrix.
 *
 * X       quantile matrix at which to evaluate the PDF (dimension d x d)
 * L_X     matrix resulting from the Cholesky decomposition of
 *         of the quantile matrix at which to evaluate the PDF
 *         X = L_X L_X^T (dimension d x d)
 * df      degrees of freedom
 * L       matrix resulting from the Cholesky decomposition of
 *         the scale matrix V = L L^T (dimension d x d)
 * result  output of the density (dimension 1)
 * work    matrix used for intermediate computations (dimension d x d)
 */
int
gsl_ran_wishart_log_pdf (const gsl_matrix * X,
                         const gsl_matrix * L_X,
                         const double df,
                         const gsl_matrix * L,
                         double * result,
                         gsl_matrix * work)
{
  if (L->size1 != L->size2)
    {
      GSL_ERROR("L should be a square matrix", GSL_ENOTSQR);
    }
  else if (X->size1 != X->size2)
    {
      GSL_ERROR("X should be a square matrix", GSL_ENOTSQR);
    }
  else if (L_X->size1 != L_X->size2)
    {
      GSL_ERROR("L_X should be a square matrix", GSL_ENOTSQR);
    }
  else if (X->size1 != L->size1)
    {
      GSL_ERROR("incompatible dimensions of X matrix", GSL_EBADLEN);
    }
  else if (L_X->size1 != L->size1)
    {
      GSL_ERROR("incompatible dimensions of L_X matrix", GSL_EBADLEN);
    }
  else if (df <= L->size1 - 1)
    {
      GSL_ERROR("incompatible degrees of freedom", GSL_EDOM);
    }
  else
    {
      size_t d = L->size1, i;
      int status;
      double log_mv_Ga, log_det_V, log_det_X, tr_Vinv_X;

      /* compute the log of the multivariate Gamma */
      log_mv_Ga = d * (d-1) * 0.25 * log(M_PI);
      for (i = 0; i < d; ++i)
        {
          log_mv_Ga += gsl_sf_lngamma((df - i + 1) * 0.5);
        }

      /* compute the log of the determinant of the scale matrix */
      log_det_V = log(gsl_matrix_get(L, 0, 0));
      for (i = 1; i < d; ++i)
        {
          log_det_V += log(gsl_matrix_get(L, i, i));
        }
      log_det_V = 2 * log_det_V;

      /* compute the log of the determinant of the quantile matrix */
      log_det_X = log(gsl_matrix_get(L_X, 0, 0));
      for (i = 1; i < d; ++i)
        {
          log_det_X += log(gsl_matrix_get(L_X, i, i));
        }
      log_det_X = 2 * log_det_X;

      /* compute the trace of V^(-1) X */
      status = gsl_linalg_cholesky_solve_mat(L, X, work);
      if (status)
        return status;
      tr_Vinv_X = gsl_matrix_get(work, 0, 0);
      for (i = 1; i < d; ++i)
        {
          tr_Vinv_X += gsl_matrix_get(work, i, i);
        }

      /* add all to get the log of the PDF */
      *result = - (0.5 * df * d) * log(2.0)
                - (0.5 * df) * log_det_V
                - log_mv_Ga
                + 0.5 * (df - d - 1) * log_det_X
                - 0.5 * tr_Vinv_X;

      return GSL_SUCCESS;
    }
}

int
gsl_ran_wishart_pdf (const gsl_matrix * X,
                     const gsl_matrix * L_X,
                     const double df,
                     const gsl_matrix * L,
                     double * result,
                     gsl_matrix * work)
{
  double logpdf;
  int status = gsl_ran_wishart_log_pdf(X, L_X, df, L, &logpdf, work);

  if (status == GSL_SUCCESS)
    *result = exp(logpdf);

  return status;
}