/* tsqr.c
 * 
 * Copyright (C) 2015 Patrick Alken
 * 
 * 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.
 */

/*
 * This module implements the sequential TSQR algorithm
 * described in
 *
 * [1] Demmel, J., Grigori, L., Hoemmen, M. F., and Langou, J.
 *     "Communication-optimal parallel and sequential QR and LU factorizations",
 *     UCB Technical Report No. UCB/EECS-2008-89, 2008.
 *
 * The algorithm operates on a tall least squares system:
 *
 * [ A_1 ] x = [ b_1 ]
 * [ A_2 ]     [ b_2 ]
 * [ ... ]     [ ... ]
 * [ A_k ]     [ b_k ]
 *
 * as follows:
 *
 * 1. Initialize
 *    a. [Q_1,R_1] = qr(A_1)
 *    b. z_1 = Q_1^T b_1
 * 2. Loop i = 2:k
 *    a. [Q_i,R_i] = qr( [ R_{i-1} ; A_i ] )
 *    b. z_i = Q_i^T [ z_{i-1} ; b_i ]
 * 3. Output:
 *    a. R = R_k
 *    b. Q^T b = z_k
 *
 * Step 2(a) is optimized to take advantage
 * of the sparse structure of the matrix
 */

#include <config.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_vector.h>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_linalg.h>
#include <gsl/gsl_errno.h>
#include <gsl/gsl_blas.h>
#include <gsl/gsl_multilarge.h>
#include <gsl/gsl_multifit.h>

typedef struct
{
  size_t p;             /* number of columns of LS matrix */
  int init;             /* QR system has been initialized */
  int svd;              /* SVD of R has been computed */
  double normb;         /* || b || for computing residual norm */

  gsl_vector *tau;      /* Householder scalars, p-by-1 */
  gsl_matrix *R;        /* [ R ; A_i ], size p-by-p */
  gsl_vector *QTb;      /* [ Q^T b ; b_i ], size p-by-1 */

  gsl_multifit_linear_workspace *multifit_workspace_p;
} tsqr_state_t;

static void *tsqr_alloc(const size_t p);
static void tsqr_free(void *vstate);
static int tsqr_reset(void *vstate);
static int tsqr_accumulate(gsl_matrix * A, gsl_vector * b,
                           void * vstate);
static int tsqr_solve(const double lambda, gsl_vector * x,
                      double * rnorm, double * snorm,
                      void * vstate);
static int tsqr_rcond(double * rcond, void * vstate);
static int tsqr_lcurve(gsl_vector * reg_param, gsl_vector * rho,
                       gsl_vector * eta, void * vstate);
static int tsqr_svd(tsqr_state_t * state);
static double tsqr_householder_transform (double *v0, gsl_vector * v);
static int tsqr_householder_hv (const double tau, const gsl_vector * v, double *w0,
                                gsl_vector * w);
static int tsqr_householder_hm (const double tau, const gsl_vector * v, gsl_matrix * R,
                                gsl_matrix * A);
static int tsqr_QR_decomp (gsl_matrix * R, gsl_matrix * A, gsl_vector * tau);

/*
tsqr_alloc()
  Allocate workspace for solving large linear least squares
problems using the TSQR approach

Inputs: p    - number of columns of LS matrix

Return: pointer to workspace
*/

static void *
tsqr_alloc(const size_t p)
{
  tsqr_state_t *state;

  if (p == 0)
    {
      GSL_ERROR_NULL("p must be a positive integer",
                     GSL_EINVAL);
    }

  state = calloc(1, sizeof(tsqr_state_t));
  if (!state)
    {
      GSL_ERROR_NULL("failed to allocate tsqr state", GSL_ENOMEM);
    }

  state->p = p;
  state->init = 0;
  state->svd = 0;
  state->normb = 0.0;

  state->R = gsl_matrix_alloc(p, p);
  if (state->R == NULL)
    {
      tsqr_free(state);
      GSL_ERROR_NULL("failed to allocate R matrix", GSL_ENOMEM);
    }

  state->QTb = gsl_vector_alloc(p);
  if (state->QTb == NULL)
    {
      tsqr_free(state);
      GSL_ERROR_NULL("failed to allocate QTb vector", GSL_ENOMEM);
    }

  state->tau = gsl_vector_alloc(p);
  if (state->tau == NULL)
    {
      tsqr_free(state);
      GSL_ERROR_NULL("failed to allocate tau vector", GSL_ENOMEM);
    }

  state->multifit_workspace_p = gsl_multifit_linear_alloc(p, p);
  if (state->multifit_workspace_p == NULL)
    {
      tsqr_free(state);
      GSL_ERROR_NULL("failed to allocate multifit workspace", GSL_ENOMEM);
    }

  return state;
}

static void
tsqr_free(void *vstate)
{
  tsqr_state_t *state = (tsqr_state_t *) vstate;

  if (state->R)
    gsl_matrix_free(state->R);

  if (state->QTb)
    gsl_vector_free(state->QTb);

  if (state->tau)
    gsl_vector_free(state->tau);

  if (state->multifit_workspace_p)
    gsl_multifit_linear_free(state->multifit_workspace_p);

  free(state);
}

static int
tsqr_reset(void *vstate)
{
  tsqr_state_t *state = (tsqr_state_t *) vstate;

  gsl_matrix_set_zero(state->R);
  gsl_vector_set_zero(state->QTb);
  state->init = 0;
  state->svd = 0;
  state->normb = 0.0;

  return GSL_SUCCESS;
}

/*
tsqr_accumulate()
  Add a new block of rows to the QR system

Inputs: A      - new block of rows, n-by-p
        b      - new rhs vector n-by-1
        vstate - workspace

Return: success/error

Notes:
1) On output, the upper triangular portion of state->R(1:p,1:p)
contains current R matrix

2) state->QTb(1:p) contains current Q^T b vector

3) A and b are destroyed
*/

static int
tsqr_accumulate(gsl_matrix * A, gsl_vector * b, void * vstate)
{
  tsqr_state_t *state = (tsqr_state_t *) vstate;
  const size_t n = A->size1;
  const size_t p = A->size2;

  if (p != state->p)
    {
      GSL_ERROR("columns of A do not match workspace", GSL_EBADLEN);
    }
  else if (n != b->size)
    {
      GSL_ERROR("A and b have different numbers of rows", GSL_EBADLEN);
    }
  else if (state->init == 0)
    {
      int status;
      const size_t npmin = GSL_MIN(n, p);
      gsl_vector_view tau = gsl_vector_subvector(state->tau, 0, npmin);
      gsl_matrix_view R = gsl_matrix_submatrix(state->R, 0, 0, npmin, p);
      gsl_matrix_view Av = gsl_matrix_submatrix(A, 0, 0, npmin, p);
      gsl_vector_view QTb = gsl_vector_subvector(state->QTb, 0, npmin);
      gsl_vector_view bv = gsl_vector_subvector(b, 0, npmin);

      /* this is the first matrix block A_1, compute its (dense) QR decomposition */

      /* compute QR decomposition of A */
      status = gsl_linalg_QR_decomp(A, &tau.vector);
      if (status)
        return status;

      /* store upper triangular R factor in state->R */
      gsl_matrix_tricpy('U', 1, &R.matrix, &Av.matrix);

      /* compute ||b|| */
      state->normb = gsl_blas_dnrm2(b);

      /* compute Q^T b and keep the first p elements */
      gsl_linalg_QR_QTvec(A, &tau.vector, b);
      gsl_vector_memcpy(&QTb.vector, &bv.vector);

      state->init = 1;

      return GSL_SUCCESS;
    }
  else
    {
      int status;

      /* compute QR decomposition of [ R_{i-1} ; A_i ], accounting for
       * sparse structure */
      status = tsqr_QR_decomp(state->R, A, state->tau);
      if (status)
        return status;

      /* update ||b|| */
      state->normb = gsl_hypot(state->normb, gsl_blas_dnrm2(b));

      /*
       * compute Q^T [ QTb_{i - 1}; b_i ], accounting for the sparse
       * structure of the Householder reflectors
       */
      {
        size_t i;

        for (i = 0; i < p; i++)
          {
            const double ti = gsl_vector_get (state->tau, i);
            gsl_vector_const_view h = gsl_matrix_const_column (A, i);
            double *wi = gsl_vector_ptr(state->QTb, i);
            tsqr_householder_hv (ti, &(h.vector), wi, b);
          }
      }

      return GSL_SUCCESS;
    }
}

/*
tsqr_solve()
  Solve the least squares system:

chi^2 = || QTb - R x ||^2 + lambda^2 || x ||^2

using the SVD of R

Inputs: lambda - regularization parameter
        x      - (output) solution vector p-by-1
        rnorm  - (output) residual norm ||b - A x||
        snorm  - (output) solution norm ||x||
        vstate - workspace

Return: success/error
*/

static int
tsqr_solve(const double lambda, gsl_vector * x,
           double * rnorm, double * snorm,
           void * vstate)
{
  tsqr_state_t *state = (tsqr_state_t *) vstate;
  const size_t p = x->size;

  if (p != state->p)
    {
      GSL_ERROR("solution vector does not match workspace", GSL_EBADLEN);
    }
  else
    {
      int status;

      /* compute SVD of R if not already computed */
      if (state->svd == 0)
        {
          status = tsqr_svd(state);
          if (status)
            return status;
        }

      status = gsl_multifit_linear_solve(lambda, state->R, state->QTb, x, rnorm, snorm,
                                         state->multifit_workspace_p);
      if (status)
        return status;

      /*
       * Since we're solving a reduced square system above, we need
       * to account for the full residual vector:
       *
       * rnorm = || [ Q1^T b - R x ; Q2^T b ] ||
       *
       * where Q1 is the thin Q factor of X, and Q2
       * are the remaining columns of Q. But:
       *
       * || Q2^T b ||^2 = ||b||^2 - ||Q1^T b||^2
       * 
       * so add this into the rnorm calculation
       */
      {
        double norm_Q1Tb = gsl_blas_dnrm2(state->QTb);
        double ratio = norm_Q1Tb / state->normb;
        double diff = 1.0 - ratio*ratio;

        if (diff > GSL_DBL_EPSILON)
          {
            double norm_Q2Tb = state->normb * sqrt(diff);
            *rnorm = gsl_hypot(*rnorm, norm_Q2Tb);
          }
      }

      return GSL_SUCCESS;
    }
}

/*
tsqr_lcurve()
  Compute L-curve of least squares system

Inputs: reg_param - (output) vector of regularization parameters
        rho       - (output) vector of residual norms
        eta       - (output) vector of solution norms
        vstate    - workspace

Return: success/error
*/

static int
tsqr_lcurve(gsl_vector * reg_param, gsl_vector * rho,
            gsl_vector * eta, void * vstate)
{
  tsqr_state_t *state = (tsqr_state_t *) vstate;
  int status;

  /* compute SVD of R if not already computed */
  if (state->svd == 0)
    {
      status = tsqr_svd(state);
      if (status)
        return status;
    }

  status = gsl_multifit_linear_lcurve(state->QTb, reg_param, rho, eta,
                                      state->multifit_workspace_p);

  /* now add contribution to rnorm from Q2 factor */
  {
    double norm_Q1Tb = gsl_blas_dnrm2(state->QTb);
    double ratio = norm_Q1Tb / state->normb;
    double diff = 1.0 - ratio*ratio;
    size_t i;

    if (diff > GSL_DBL_EPSILON)
      {
        double norm_Q2Tb = state->normb * sqrt(diff);

        for (i = 0; i < rho->size; ++i)
          {
            double *rhoi = gsl_vector_ptr(rho, i);
            *rhoi = gsl_hypot(*rhoi, norm_Q2Tb);
          }
      }
  }

  return status;
}

static int
tsqr_rcond(double * rcond, void * vstate)
{
  tsqr_state_t *state = (tsqr_state_t *) vstate;

  /* compute SVD of R if not already computed */
  if (state->svd == 0)
    {
      int status = tsqr_svd(state);
      if (status)
        return status;
    }

  *rcond = gsl_multifit_linear_rcond(state->multifit_workspace_p);

  return GSL_SUCCESS;
}

/*
tsqr_svd()
  Compute the SVD of the upper triangular
R factor. This allows us to compute the upper/lower
bounds on the regularization parameter and compute
the matrix reciprocal condition number.

Inputs: state - workspace

Return: success/error
*/

static int
tsqr_svd(tsqr_state_t * state)
{
  int status;

  status = gsl_multifit_linear_svd(state->R, state->multifit_workspace_p);
  if (status)
    {
      GSL_ERROR("error computing SVD of R", status);
    }

  state->svd = 1;

  return GSL_SUCCESS;
}

/*
tsqr_householder_transform()
  This routine is an optimized version of
gsl_linalg_householder_transform(), designed for the QR
decomposition of M-by-N matrices of the form:

T = [ R ]
    [ A ]

where R is N-by-N upper triangular, and A is (M-N)-by-N dense.
This routine computes a householder transformation (tau,v) of a 
x so that P x = [ I - tau*v*v' ] x annihilates x(1:n-1). x will
be a subcolumn of the matrix T, and so its structure will be:

x = [ x0 ] <- 1 nonzero value for the diagonal element of R
    [ 0  ] <- N - j - 1 zeros, where j is column of matrix in [0,N-1]
    [ x  ] <- M-N nonzero values for the dense part A

Inputs: v0 - pointer to diagonal element of R
             on input, v0 = x0;
        v  - on input, x vector
             on output, householder vector v
*/

static double
tsqr_householder_transform (double *v0, gsl_vector * v)
{
  /* replace v[0:M-1] with a householder vector (v[0:M-1]) and
     coefficient tau that annihilate v[1:M-1] */

  double alpha, beta, tau ;
  
  /* compute xnorm = || [ 0 ; v ] ||, ignoring zero part of vector */
  double xnorm = gsl_blas_dnrm2(v);

  if (xnorm == 0) 
    {
      return 0.0; /* tau = 0 */
    }

  alpha = *v0;
  beta = - (alpha >= 0.0 ? +1.0 : -1.0) * hypot(alpha, xnorm) ;
  tau = (beta - alpha) / beta ;
  
  {
    double s = (alpha - beta);
    
    if (fabs(s) > GSL_DBL_MIN) 
      {
        gsl_blas_dscal (1.0 / s, v);
        *v0 = beta;
      }
    else
      {
        gsl_blas_dscal (GSL_DBL_EPSILON / s, v);
        gsl_blas_dscal (1.0 / GSL_DBL_EPSILON, v);
        *v0 = beta;
      }
  }
  
  return tau;
}

/*
tsqr_householder_hv()
  Apply Householder reflector to a vector. The Householder
reflectors are for the QR decomposition of the matrix

  [ R ]
  [ A ]

where R is p-by-p upper triangular and A is n-by-p dense.
Therefore all relevant components of the Householder
vector are stored in the columns of A, while the components
in R are 0, except for diag(R) which are 1.

The vector w to be transformed is partitioned as

  [ w1 ]
  [ w2 ]

where w1 is p-by-1 and w2 is n-by-1. The w2 portion
of w is transformed by v, but most of w1 remains unchanged
except for the first element, w0

Inputs: tau - Householder scalar
        v   - Householder vector, n-by-1
        w0  - (input/output)
              on input, w1(0);
              on output, transformed w1(0)
        w   - (input/output) n-by-1
              on input, vector w2;
              on output, P*w2
*/

static int
tsqr_householder_hv (const double tau, const gsl_vector * v, double *w0, gsl_vector * w)
{
  /* applies a householder transformation v to vector w */
 
  if (tau == 0)
    return GSL_SUCCESS ;

  {
    double d1, d;

    /* compute d1 = v(2:n)' w(2:n) */
    gsl_blas_ddot (v, w, &d1);

    /* compute d = v'w = w(1) + d1 since v(1) = 1 */
    d = *w0 + d1;

    /* compute w = w - tau (v) (v'w) */
    *w0 -= tau * d;
    gsl_blas_daxpy (-tau * d, v, w);
  }
  
  return GSL_SUCCESS;
}

/*
tsqr_householder_hm()
  Apply Householder reflector to a submatrix of

  [ R ]
  [ A ]

where R is p-by-p upper triangular and A is n-by-p dense.
The diagonal terms of R are already transformed by
tsqr_householder_transform(), so we just need to operate
on the submatrix A(:,i:p) as well as the superdiagonal
elements of R

Inputs: tau - Householder scalar
        v   - Householder vector
        R   - upper triangular submatrix of R, (p-i)-by-(p-i-1)
        A   - dense submatrix of A, n-by-(p-i)
*/

static int
tsqr_householder_hm (const double tau, const gsl_vector * v, gsl_matrix * R,
                     gsl_matrix * A)
{
  /* applies a householder transformation v,tau to matrix [ R ; A ] */

  if (tau == 0.0)
    {
      return GSL_SUCCESS;
    }
  else
    {
      size_t j;

      for (j = 0; j < A->size2; j++)
        {
          double R0j = gsl_matrix_get (R, 0, j);
          double wj;
          gsl_vector_view A1j = gsl_matrix_column(A, j);

          gsl_blas_ddot (&A1j.vector, v, &wj);
          wj += R0j;

          gsl_matrix_set (R, 0, j, R0j - tau *  wj);

          gsl_blas_daxpy (-tau * wj, v, &A1j.vector);
        }

      return GSL_SUCCESS;
    }
}

/*
tsqr_QR_decomp()
  Compute the QR decomposition of the matrix

  [ R ]
  [ A ]

where R is p-by-p upper triangular and A is n-by-p dense.

Inputs: R   - upper triangular p-by-p matrix
        A   - dense n-by-p matrix
        tau - Householder scalars
*/

static int
tsqr_QR_decomp (gsl_matrix * R, gsl_matrix * A, gsl_vector * tau)
{
  const size_t n = A->size1;
  const size_t p = R->size2;

  if (R->size2 != A->size2)
    {
      GSL_ERROR ("R and A have different number of columns", GSL_EBADLEN);
    }
  else if (tau->size != p)
    {
      GSL_ERROR ("size of tau must be p", GSL_EBADLEN);
    }
  else
    {
      size_t i;

      for (i = 0; i < p; i++)
        {
          /* Compute the Householder transformation to reduce the j-th
             column of the matrix [ R ; A ] to a multiple of the j-th unit vector,
             taking into account the sparse structure of R */

          gsl_vector_view c = gsl_matrix_column(A, i);
          double *Rii = gsl_matrix_ptr(R, i, i);
          double tau_i = tsqr_householder_transform(Rii, &c.vector);

          gsl_vector_set (tau, i, tau_i);

          /* Apply the transformation to the remaining columns and
             update the norms */

          if (i + 1 < p)
            {
              gsl_matrix_view Rv = gsl_matrix_submatrix(R, i, i + 1, p - i, p - (i + 1));
              gsl_matrix_view Av = gsl_matrix_submatrix(A, 0, i + 1, n, p - (i + 1));
              tsqr_householder_hm (tau_i, &(c.vector), &(Rv.matrix), &(Av.matrix));
            }
        }

      return GSL_SUCCESS;
    }
}

static const gsl_multilarge_linear_type tsqr_type =
{
  "tsqr",
  tsqr_alloc,
  tsqr_reset,
  tsqr_accumulate,
  tsqr_solve,
  tsqr_rcond,
  tsqr_lcurve,
  tsqr_free
};

const gsl_multilarge_linear_type * gsl_multilarge_linear_tsqr =
  &tsqr_type;