// Copyright (c) 2021 Stig Rune Sellevag
//
// This file is distributed under the MIT License. See the accompanying file
// LICENSE.txt or http://www.opensource.org/licenses/mit-license.php for terms
// and conditions.

#ifndef SCILIB_LINALG_BLAS2_MATRIX_VECTOR_PRODUCT_H
#define SCILIB_LINALG_BLAS2_MATRIX_VECTOR_PRODUCT_H

#ifdef USE_MKL
#include <mkl.h>
#else
#include <cblas.h>
#endif

#include <scilib/mdarray.h>
#include <scilib/linalg_impl/lapack_types.h>
#include <experimental/mdspan>
#include <cassert>
#include <complex>
#include <type_traits>
#include <iostream>

namespace Sci {
namespace Linalg {

namespace stdex = std::experimental;

template <class T_a,
          stdex::extents<>::size_type nrows_a,
          stdex::extents<>::size_type ncols_a,
          class Layout_a,
          class Accessor_a,
          class T_x,
          stdex::extents<>::size_type ext_x,
          class Layout_x,
          class Accessor_x,
          class T_y,
          stdex::extents<>::size_type ext_y,
          class Layout_y,
          class Accessor_y>
    requires(!std::is_const_v<T_y>)
inline void matrix_vector_product(
    stdex::mdspan<T_a, stdex::extents<nrows_a, ncols_a>, Layout_a, Accessor_a>
        a,
    stdex::mdspan<T_x, stdex::extents<ext_x>, Layout_x, Accessor_x> x,
    stdex::mdspan<T_y, stdex::extents<ext_y>, Layout_y, Accessor_y> y)
{
    static_assert(x.static_extent(0) == a.static_extent(1));
    using size_type = stdex::extents<>::size_type;

    for (size_type i = 0; i < a.extent(0); ++i) {
        y(i) = T_y{0};
        for (size_type j = 0; j < a.extent(1); ++j) {
            y(i) += a(i, j) * x(j);
        }
    }
}

template <class Layout>
inline void matrix_vector_product(Sci::Matrix_view<double, Layout> a,
                                  Sci::Vector_view<double, Layout> x,
                                  Sci::Vector_view<double, Layout> y)
{
    static_assert(x.static_extent(0) == a.static_extent(1));

    constexpr double alpha = 1.0;
    constexpr double beta = 0.0;

    const BLAS_INT m = static_cast<BLAS_INT>(a.extent(0));
    const BLAS_INT n = static_cast<BLAS_INT>(a.extent(1));
    const BLAS_INT incx = static_cast<BLAS_INT>(x.stride(0));
    const BLAS_INT incy = static_cast<BLAS_INT>(y.stride(0));

    auto matrix_layout = CblasRowMajor;
    BLAS_INT lda = n;

    if constexpr (std::is_same_v<Layout, stdex::layout_left>) {
        matrix_layout = CblasColMajor;
        lda = m;
    }
    cblas_dgemv(matrix_layout, CblasNoTrans, m, n, alpha, a.data(), lda,
                x.data(), incx, beta, y.data(), incy);
}

template <class Layout>
inline void matrix_vector_product(Sci::Matrix_view<const double, Layout> a,
                                  Sci::Vector_view<const double, Layout> x,
                                  Sci::Vector_view<double, Layout> y)
{
    static_assert(x.static_extent(0) == a.static_extent(1));

    constexpr double alpha = 1.0;
    constexpr double beta = 0.0;

    const BLAS_INT m = static_cast<BLAS_INT>(a.extent(0));
    const BLAS_INT n = static_cast<BLAS_INT>(a.extent(1));
    const BLAS_INT incx = static_cast<BLAS_INT>(x.stride(0));
    const BLAS_INT incy = static_cast<BLAS_INT>(y.stride(0));

    auto matrix_layout = CblasRowMajor;
    BLAS_INT lda = n;

    if constexpr (std::is_same_v<Layout, stdex::layout_left>) {
        matrix_layout = CblasColMajor;
        lda = m;
    }
    cblas_dgemv(matrix_layout, CblasNoTrans, m, n, alpha, a.data(), lda,
                x.data(), incx, beta, y.data(), incy);
}

#ifdef USE_MKL
// Does not work with OpenBLAS version 0.2.14.1
template <class Layout>
inline void
matrix_vector_product(Sci::Matrix_view<std::complex<double>, Layout> a,
                      Sci::Vector_view<std::complex<double>, Layout> x,
                      Sci::Vector_view<std::complex<double>, Layout> y)
{
    static_assert(x.static_extent(0) == a.static_extent(1));

    constexpr std::complex<double> alpha = {1.0, 0.0};
    constexpr std::complex<double> beta = {0.0, 0.0};

    const BLAS_INT m = static_cast<BLAS_INT>(a.extent(0));
    const BLAS_INT n = static_cast<BLAS_INT>(a.extent(1));
    const BLAS_INT incx = static_cast<BLAS_INT>(x.stride(0));
    const BLAS_INT incy = static_cast<BLAS_INT>(y.stride(0));

    auto matrix_layout = CblasRowMajor;
    BLAS_INT lda = n;

    if constexpr (std::is_same_v<Layout, stdex::layout_left>) {
        matrix_layout = CblasColMajor;
        lda = m;
    }
    cblas_zgemv(matrix_layout, CblasNoTrans, m, n, &alpha, a.data(), lda,
                x.data(), incx, &beta, y.data(), incy);
}

template <class Layout>
inline void
matrix_vector_product(Sci::Matrix_view<const std::complex<double>, Layout> a,
                      Sci::Vector_view<const std::complex<double>, Layout> x,
                      Sci::Vector_view<std::complex<double>, Layout> y)
{
    static_assert(x.static_extent(0) == a.static_extent(1));

    constexpr std::complex<double> alpha = {1.0, 0.0};
    constexpr std::complex<double> beta = {0.0, 0.0};

    const BLAS_INT m = static_cast<BLAS_INT>(a.extent(0));
    const BLAS_INT n = static_cast<BLAS_INT>(a.extent(1));
    const BLAS_INT incx = static_cast<BLAS_INT>(x.stride(0));
    const BLAS_INT incy = static_cast<BLAS_INT>(y.stride(0));

    auto matrix_layout = CblasRowMajor;
    BLAS_INT lda = n;

    if constexpr (std::is_same_v<Layout, stdex::layout_left>) {
        matrix_layout = CblasColMajor;
        lda = m;
    }
    cblas_zgemv(matrix_layout, CblasNoTrans, m, n, &alpha, a.data(), lda,
                x.data(), incx, &beta, y.data(), incy);
}
#endif

template <class T, class Layout, class Allocator>
inline Sci::Vector<T, Layout, Allocator>
matrix_vector_product(const Sci::Matrix<T, Layout, Allocator>& a,
                      const Sci::Vector<T, Layout, Allocator>& x)
{
    Sci::Vector<T, Layout, Allocator> res(a.extent(0));
    matrix_vector_product(a.view(), x.view(), res.view());
    return res;
}

template <class T, class Layout, class Allocator>
inline void matrix_vector_product(const Sci::Matrix<T, Layout, Allocator>& a,
                                  const Sci::Vector<T, Layout, Allocator>& x,
                                  Sci::Vector<T, Layout, Allocator>& res)
{
    assert(res.size() == a.extent(0));
    matrix_vector_product(a.view(), x.view(), res.view());
}

} // namespace Linalg
} // namespace Sci

#endif // SCILIB_LINALG_BLAS2_MATRIX_VECTOR_PRODUCT_H


///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

/*
 Contains the PERDIF program that implements personalized diffusions for
 recommendations.

 Code by: Dimitris Berberidis and Athanasios N. Nikolakopoulos
 University of Minnesota 2018-2019
*/

//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

#include <inttypes.h>
#include <math.h>
#include <stdbool.h>
#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
//#include <cblas.h>
#include <mkl.h>

#include "perdif_fit.h"
#include "rec_defs.h"

#if DEBUG
static double cost_func(double *, double *, double *, sz_med_t);
#endif

static void matrix_matrix_product(double *C, double *A, double *B, int, int,  int);
static double *unfold_matrix(d_mat_t);
static void constr_QP_with_PG(double *, double *, double *, sz_med_t);
static void grammian_matrix(double *, double *, int, int);
static double max_abs_dif(double *, double *, sz_long_t);
// static void project_to_simplex(double *, sz_med_t);
// static void project_to_simplex_thanos(double *, double *, sz_med_t length);
static void simplexproj_Duchi(double *y, double *x, const unsigned int length);
static void matvec(double *, double *, double *, int, int);
static double frob_norm(double *, sz_med_t);

///////////////////////////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////////////////////

void k_simplex(double *theta, d_mat_t P, int max_walk, double *dummy, bool rmse) {

  if (!rmse) {
    // adjust for BPR  metric
    for (int i = 1; i < P.num_rows; i++) {
      for (int j = 0; j < P.num_cols; j++)
        P.val[i][j] = P.val[0][j] - P.val[i][j];
    }
  }

  // Unfold P for BLAS compatibility
  double *P_unfold = unfold_matrix(P);

  // Prepare parameters for constraned QP

  // first the Hessian matrix
  double A[max_walk * max_walk];
  grammian_matrix(A, P_unfold, (int)P.num_rows, (int)P.num_cols);

  // then the linear part
  double b[max_walk];

  if (rmse) {
    for (int i = 0; i < max_walk; i++)
      b[i] = -2.0 * (P_unfold[i]);
  } else {
    for (int i = 0; i < max_walk; i++) {
      b[i] = 0.0;
      for (int j = 0; j < P.num_rows; j++) {
        b[i] += P.val[i][j];
      }
      b[i] *= -2.0;
    }
  }

  constr_QP_with_PG(theta, A, b, (sz_med_t)max_walk);

  // free
  free(P_unfold);
}

// Simple search returns the single best step ( i.e. theta = [0 0 0 1 0 0 .... 0
// 0 ] )

void single_best(double *theta, d_mat_t P, int max_walk, double *dummy, bool rmse) {

  if (!rmse) {
    // adjust for BPR  metric
    for (int i = 1; i < P.num_rows; i++) {
      for (int j = 0; j < P.num_cols; j++)
        P.val[i][j] = P.val[0][j] - P.val[i][j];
    }
  }

  for (int k = 0; k < max_walk; k++)
    theta[k] = 0.0;

  double loss[max_walk];

  // Compute loss of each k-step and find smallest
  int best_k = 0;
  double min_loss = 1.0e5;
  for (int k = 0; k < max_walk; k++) {
    loss[k] += pow(1.0 - P.val[0][k], 2);
    if (rmse) {
      for (int j = 1; j < P.num_rows; j++)
        loss[k] += P.val[j][k] * P.val[j][k];
    } else {
      for (int j = 1; j < P.num_rows; j++)
        loss[k] += pow(1.0 - P.val[j][k], 2);
    }
    if (loss[k] < min_loss) {
      min_loss = loss[k];
      best_k = k;
    }
  }

  theta[best_k] = 1.0;
}

void ppr_grid(double *theta, d_mat_t P, int max_walk, double *ppr_c, bool rmse) {

  if (!rmse) {
    // adjust for BPR  metric
    for (int i = 1; i < P.num_rows; i++) {
      for (int j = 0; j < P.num_cols; j++)
        P.val[i][j] = P.val[0][j] - P.val[i][j];
    }
  }

  // Unfold P for BLAS compatibility
  double *P_unfold = unfold_matrix(P);

  // Candidate diffusions
  double *C = (double *)malloc(P.num_rows * GRID_POINTS * sizeof(double));
  matrix_matrix_product(C, P_unfold, ppr_c, P.num_rows, max_walk, GRID_POINTS);

  // find the best one
  double loss[GRID_POINTS];
  int best_point = 0;
  double min_loss = 1.0e5;
  for (int i = 0; i < GRID_POINTS; i++) {
    // first the regularizer
    loss[i] = 0.0;
    // then the ratings
    if (rmse) {
      loss[i] += pow(1.0 - C[i], 2);
      for (int j = 1; j < P.num_rows; j++)
        loss[i] += pow(C[j * GRID_POINTS + i], 2);
    } else {
      for (int j = 0; j < P.num_rows; j++)
        loss[i] += pow(1.0 - C[j * GRID_POINTS + i], 2);
    }

    if (loss[i] < min_loss) {
      min_loss = loss[i];
      best_point = i;
    }
  }

  // return the best one
  for (int k = 0; k < max_walk; k++)
    theta[k] = ppr_c[k * GRID_POINTS + best_point];

  // free
  free(P_unfold);
  free(C);
}

void hk_grid(double *theta, d_mat_t P, int max_walk, double *hk_c, bool rmse) {

  if (!rmse) {
    // adjust for BPR  metric
    for (int i = 1; i < P.num_rows; i++) {
      for (int j = 0; j < P.num_cols; j++)
        P.val[i][j] = P.val[0][j] - P.val[i][j];
    }
  }

  // Unfold P for BLAS compatibility
  double *P_unfold = unfold_matrix(P);

  double *C = (double *)malloc(P.num_rows * GRID_POINTS * sizeof(double));
  matrix_matrix_product(C, P_unfold, hk_c, P.num_rows, max_walk, GRID_POINTS);

  // find the best one
  double loss[GRID_POINTS];
  int best_point = 0;
  double min_loss = 1.0e5;
  for (int i = 0; i < GRID_POINTS; i++) {
    // first the regularizer
    loss[i] = 0.0;

    // then the ratings
    if (rmse) {
      loss[i] += pow(1.0 - C[i], 2);
      for (int j = 1; j < P.num_rows; j++)
        loss[i] += pow(C[j * GRID_POINTS + i], 2);
    } else {
      for (int j = 0; j < P.num_rows; j++)
        loss[i] += pow(1.0 - C[j * GRID_POINTS + i], 2);
    }

    if (loss[i] < min_loss) {
      min_loss = loss[i];
      best_point = i;
    }
  }

  // return the best one
  for (int k = 0; k < max_walk; k++)
    theta[k] = hk_c[k * GRID_POINTS + best_point];

  // free
  free(P_unfold);
  free(C);
}

void dictionary_single(double *theta, d_mat_t P, int max_walk, double *dict, bool rmse) {

  if (!rmse) {
    // adjust for BPR  metric
    for (int i = 1; i < P.num_rows; i++) {
      for (int j = 0; j < P.num_cols; j++)
        P.val[i][j] = P.val[0][j] - P.val[i][j];
    }
  }

  // Unfold P for BLAS compatibility
  double *P_unfold = unfold_matrix(P);

  // Prepare matrices
  double *PD = (double *)malloc(P.num_rows * GRID_POINTS * sizeof(double));
  matrix_matrix_product(PD, P_unfold, dict, (int)P.num_rows, (int)P.num_cols,
                        GRID_POINTS);
  double DD[GRID_POINTS * GRID_POINTS];
  grammian_matrix(DD, dict, max_walk, GRID_POINTS);
  double A[GRID_POINTS * GRID_POINTS];
  grammian_matrix(A, PD, (int)P.num_rows, GRID_POINTS);

  double g_D[GRID_POINTS];
  for (int i = 0; i < GRID_POINTS; i++) {
    g_D[i] = 0.0;
  }

  if (rmse) {
    for (int i = 0; i < GRID_POINTS; i++)
      g_D[i] += -2.0 * PD[i];
  } else {
    for (int i = 0; i < GRID_POINTS; i++) {
      for (int j = 0; j < P.num_rows; j++)
        g_D[i] += -2.0 * PD[j * GRID_POINTS + i];
    }
  }

  double theta_dict[GRID_POINTS];

  // call QP solver for dictionary coefficients
  constr_QP_with_PG(theta_dict, A, g_D, (sz_med_t)GRID_POINTS);

  // construct diffusion from dictionary coefficients
  matvec(theta, dict, theta_dict, max_walk, GRID_POINTS);

  // free
  free(PD);
  free(P_unfold);
}

void dictionary(double *theta, d_mat_t P, int max_walk, double *dict, bool rmse) {

  if (!rmse) {
    // adjust for BPR  metric
    for (int i = 1; i < P.num_rows; i++) {
      for (int j = 0; j < P.num_cols; j++)
        P.val[i][j] = P.val[0][j] - P.val[i][j];
    }
  }

  // Unfold P for BLAS compatibility
  double *P_unfold = unfold_matrix(P);

  // Prepare matrices
  double *PD = (double *)malloc(P.num_rows * 2 * GRID_POINTS * sizeof(double));
  matrix_matrix_product(PD, P_unfold, dict, (int)P.num_rows, (int)P.num_cols,
                        2 * GRID_POINTS);
  double DD[4 * GRID_POINTS * GRID_POINTS];
  grammian_matrix(DD, dict, max_walk, 2 * GRID_POINTS);
  double A[4 * GRID_POINTS * GRID_POINTS];
  grammian_matrix(A, PD, (int)P.num_rows, 2 * GRID_POINTS);

  double g_D[2 * GRID_POINTS];
  for (int i = 0; i < 2 * GRID_POINTS; i++) {
    g_D[i] = 0.0;
  }

  if (rmse) {
    for (int i = 0; i < 2 * GRID_POINTS; i++)
      g_D[i] += -2.0 * PD[i];
  } else {
    for (int i = 0; i < 2 * GRID_POINTS; i++) {
      for (int j = 0; j < P.num_rows; j++)
        g_D[i] += -2.0 * PD[j * 2 * GRID_POINTS + i];
    }
  }

  double theta_dict[2 * GRID_POINTS];

  // call QP solver for dictionary coefficients
  constr_QP_with_PG(theta_dict, A, g_D, (sz_med_t)2 * GRID_POINTS);

  // construct diffusion from dictionary coefficients
  matvec(theta, dict, theta_dict, max_walk, 2 * GRID_POINTS);

  // free
  free(PD);
  free(P_unfold);
}

// Returns PPR coefficients for given points on a grid of parameter a
double *ppr_coefficients(int max_walk) {

  int grid_points = GRID_POINTS;

  double *coef = (double *)malloc(max_walk * grid_points * sizeof(double));
  double width = 1.0 / (double)(grid_points + 2);

  double col_sum[grid_points];
  for (int j = 0; j < grid_points; j++)
    col_sum[j] = 0.0;

  for (int i = 0; i < max_walk; i++) {
    for (int j = 0; j < grid_points; j++) {
      // pagerank weights per step
      coef[i * grid_points + j] = (i == 0)
                                      ? (1 - (0.01 + width * ((double)j)))
                                      : (1 - (0.01 + width * ((double)j))) *
                                            pow(0.01 + width * ((double)j), i);
      col_sum[j] += coef[i * grid_points + j];
    }
  }

  // Normalize coefficients so tha they sum to 1
  for (int i = 0; i < max_walk; i++) {
    for (int j = 0; j < grid_points; j++)
      coef[i * grid_points + j] /= col_sum[j];
  }

  return coef;
}

// Returns HK coefficients for given points on a grid of parameter \tau
double *hk_coefficients(int max_walk) {

  int grid_points = GRID_POINTS;

  double *coef = (double *)malloc(max_walk * grid_points * sizeof(double));
  double width = (double)MAX_HK_COEF / (double)(grid_points + 2);

  // compute factorials from 1 to max_walk
  sz_long_t factorial[max_walk];
  factorial[0] = 1;
  for (int j = 1; j < max_walk; j++)
    factorial[j] = factorial[j - 1] * j;

  double col_sum[grid_points];
  for (int j = 0; j < grid_points; j++)
    col_sum[j] = 0.0;

  for (int i = 0; i < max_walk; i++) {
    for (int j = 0; j < grid_points; j++) {
      coef[i * grid_points + j] =
          pow(1.0 + width * ((double)j), i) / (double)factorial[i];
      col_sum[j] += coef[i * grid_points + j];
    }
  }
  // Normalize coefficients so that they sum to 1
  for (int i = 0; i < max_walk; i++) {
    for (int j = 0; j < grid_points; j++)
      coef[i * grid_points + j] /= col_sum[j];
  }

  return coef;
}

// Returns dictionary coefficients by concatenating PPR and HK
double *dict_coefficients(int max_walk) {
  int grid_points = GRID_POINTS;

  double *coef = (double *)malloc(2 * max_walk * grid_points * sizeof(double));

  double *hk = hk_coefficients(max_walk);
  double *ppr = ppr_coefficients(max_walk);

  for (int i = 0; i < max_walk; i++) {
    for (int j = 0; j < grid_points; j++) {
      coef[2 * i * grid_points + j] = ppr[i * grid_points + j];
    }
    for (int j = grid_points; j < 2 * grid_points; j++) {
      coef[2 * i * grid_points + j] = hk[i * grid_points + j - grid_points];
    }
  }

  // free temporary coefs
  free(ppr);
  free(hk);
  return coef;
}

///////////////////////////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////////////////////

// AUXILIARY FUNCTIONS

// Interface for CBLAS mutrix matrix product
// C =A*B
// A : m x k
// B : k x n
static void matrix_matrix_product(double *C, double *A, double *B, int m, int k,
                                  int n) {

  for (int i = 0; i < m * n; i++)
    C[i] = 0.0f;

  cblas_dgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, m, n, k, 1.0f, A, k, B,
              n, 0.0f, C, n);
}

// Unfold matrix from A[i[j] to A[i*M + j]
static double *unfold_matrix(d_mat_t A) {

  double *A_u = (double *)malloc(A.num_cols * A.num_rows * sizeof(double));

  for (sz_long_t i = 0; i < A.num_rows; i++)
    memcpy(A_u + i * A.num_cols, A.val[i], A.num_cols * sizeof(double));

  return A_u;
}

// Solving constrained quadratic minimization via projected gradient descent
// The following function returns x =arg min {x^T*A*x +x^T*B} s.t. x in Prob.
// Simplex
static void constr_QP_with_PG(double *x, double *A, double *b, sz_med_t K) {
  double inf_norm, step_size;
  double x_temp[K];
  double x_prev[K];

  step_size = STEPSIZE_2 / frob_norm(A, K);

  // Initialize to uniform
  for (sz_med_t i = 0; i < K; i++)
    x[i] = 1.0f / (double)K;

  sz_med_t iter = 0;
  memcpy(x_prev, x, K * sizeof(double));
  // int count=0;
  do {
    iter++;
    // step_size = 1 / ((double)iter) * step_size;
    // step_size = 1/(sqrt(iter));//*step_size;
    // step_size = 1/(double)iter;
    // Take gradient step
    matvec(x_temp, A, x, K, K);

    for (sz_med_t j = 0; j < K; j++)
      x[j] -= step_size * (2.0f * x_temp[j] + b[j]);

    // project to feasible set
    simplexproj_Duchi(x, x, K);

#if DEBUG
    printf("\n COST: ");
    printf(" %lf ", cost_func(A, b, x, K));
#endif

    inf_norm = max_abs_dif(x_prev, x, (sz_long_t)K);

    memcpy(x_prev, x, K * sizeof(double));

  } while (iter < MAXIT_GD && inf_norm > GD_TOL_2);

  // printf("\n Counter = %d", count);
  // printf("\n Optimization finished after: %"PRIu32" iterations\n", (uint32_t)
  // iter);
}

// Grammian matrix G =A'*A using CBLAS
// A : m x n
static void grammian_matrix(double *G, double *A, int m, int n) {

  for (int i = 0; i < n * n; i++)
    G[i] = 0.0f;

  double *A_copy = (double *)malloc(m * n * sizeof(double));

  memcpy(A_copy, A, m * n * sizeof(double));

  cblas_dgemm(CblasRowMajor, CblasTrans, CblasNoTrans, n, n, m, 1.0f, A_copy, n,
              A, n, 0.0f, G, n);

  free(A_copy);
}

// Interface for CBLAS matrix vector product
// A : M x N
static void matvec(double *y, double *A, double *x, int M, int N) {

  // for (int i = 0; i < M; i++)
  // {
  // 	y[i] = 0.0f;
  // }

  cblas_dgemv(CblasRowMajor, CblasNoTrans, M, N, 1.0f, A, N, x, 1, 0.0f, y, 1);
}

//

/* Algorithm using partitioning with respect to a pivot, chosen randomly,
as given by Duchi et al. in "Efficient Projections onto the l1-Ball for
Learning in High Dimensions" */
static void simplexproj_Duchi(double *y, double *x, const unsigned int length) {
  double *auxlower = (x == y ? (double *)malloc(length * sizeof(double)) : x);
  double *auxupper = (double *)malloc(length * sizeof(double));
  double *aux = auxlower;
  double pivot;
  int auxlowerlength = 0;
  int auxupperlength = 1;
  int upperlength;
  int auxlength;
  int i = 0;
  int pospivot = (int)(rand() / (((double)RAND_MAX + 1.0) / length));
  double tauupper;
  double tau = (pivot = y[pospivot]) - 1.0;
  while (i < pospivot)
    if (y[i] < pivot)
      auxlower[auxlowerlength++] = y[i++];
    else {
      auxupper[auxupperlength++] = y[i];
      tau += (y[i++] - tau) / auxupperlength;
    }
  i++;
  while (i < length)
    if (y[i] < pivot)
      auxlower[auxlowerlength++] = y[i++];
    else {
      auxupper[auxupperlength++] = y[i];
      tau += (y[i++] - tau) / auxupperlength;
    }
  if (tau < pivot) {
    upperlength = auxupperlength;
    tauupper = tau;
    auxlength = auxlowerlength;
  } else {
    tauupper = 0.0;
    upperlength = 0;
    aux = auxupper + 1;
    auxlength = auxupperlength - 1;
  }
  while (auxlength > 0) {
    pospivot = (int)(rand() / (((double)RAND_MAX + 1.0) / auxlength));
    if (upperlength == 0)
      tau = (pivot = aux[pospivot]) - 1.0;
    else
      tau = tauupper + ((pivot = aux[pospivot]) - tauupper) / (upperlength + 1);
    i = 0;
    auxlowerlength = 0;
    auxupperlength = 1;
    while (i < pospivot)
      if (aux[i] < pivot)
        auxlower[auxlowerlength++] = aux[i++];
      else {
        auxupper[auxupperlength++] = aux[i];
        tau += (aux[i++] - tau) / (upperlength + auxupperlength);
      }
    i++;
    while (i < auxlength)
      if (aux[i] < pivot)
        auxlower[auxlowerlength++] = aux[i++];
      else {
        auxupper[auxupperlength++] = aux[i];
        tau += (aux[i++] - tau) / (upperlength + auxupperlength);
      }
    if (tau < pivot) {
      upperlength += auxupperlength;
      tauupper = tau;
      auxlength = auxlowerlength;
      aux = auxlower;
    } else {
      aux = auxupper + 1;
      auxlength = auxupperlength - 1;
    }
  }
  for (i = 0; i < length; i++)
    x[i] = (y[i] > tau ? y[i] - tauupper : 0.0000000001);
  double sum = 0.0;
  for (i = 0; i < length; i++)
    sum += x[i];
  for (i = 0; i < length; i++)
    x[i] /= sum;
  if (x == y)
    free(auxlower);
  free(auxupper);
}

// Infinity norm
static double max_abs_dif(double *a, double *b, sz_long_t len) {
  double dif, max = 0.0;

  for (sz_long_t i = 0; i < len; i++) {
    dif = fabs(a[i] - b[i]);
    max = (dif > max) ? dif : max;
  }

  return max;
}

#if DEBUG
// Evaluates quadratic with Hessian A and linear part b at x
static double cost_func(double *A, double *b, double *x, sz_med_t len) {

  double quad = 0.0f, lin = 0.0f;

  for (sz_med_t i = 0; i < len; i++) {
    for (sz_med_t j = 0; j < len; j++) {
      quad += A[i * len + j] * x[i] * x[j];
    }
    lin += b[i] * x[i];
  }
  return quad + lin;
}
#endif

// frobenious norm of double-valued square matrix
static double frob_norm(double *A, sz_med_t dim) {
  double norm = 0.0f;

  for (sz_med_t i = 0; i < dim * dim; i++) {
    norm += pow(A[i], 2.0f);
  }

  return sqrt(norm);
}


// Copyright (c) 2021 Stig Rune Sellevag
//
// This file is distributed under the MIT License. See the accompanying file
// LICENSE.txt or http://www.opensource.org/licenses/mit-license.php for terms
// and conditions.

#ifndef SCILIB_LINALG_INV_H
#define SCILIB_LINALG_INV_H

#ifdef USE_MKL
#include <mkl.h>
#else
#include <lapacke.h>
#endif

#include <scilib/mdarray.h>
#include <scilib/linalg_impl/lapack_types.h>
#include <experimental/mdspan>
#include <exception>
#include <cassert>
#include <type_traits>

namespace Sci {
namespace Linalg {

// Matrix inversion.
template <class T_a, class T_res, class Layout>
    requires(std::is_same_v<std::remove_cv_t<T_a>, double>)
inline void inv(Sci::Matrix_view<T_a, Layout> a,
                Sci::Matrix_view<T_res, Layout> res)
{
    namespace stdex = std::experimental;

    static_assert(a.is_contiguous());
    static_assert(res.is_contiguous());

    assert(a.extent(0) == a.extent(1));

    if (det(a) == 0.0) {
        throw std::runtime_error("inv: matrix not invertible");
    }
    const BLAS_INT n = static_cast<BLAS_INT>(a.extent(0));

    auto matrix_layout = LAPACK_ROW_MAJOR;
    BLAS_INT lda = n;

    if constexpr (std::is_same_v<Layout, stdex::layout_left>) {
        matrix_layout = LAPACK_COL_MAJOR;
    }

    Sci::copy(a, res);

    Sci::Vector<BLAS_INT, Layout> ipiv(n);
    Sci::Linalg::lu(res, ipiv.view()); // perform LU factorization

    BLAS_INT info =
        LAPACKE_dgetri(matrix_layout, n, res.data(), lda, ipiv.data());
    if (info != 0) {
        throw std::runtime_error("dgetri: matrix inversion failed");
    }
}

template <class Layout, class Allocator>
inline Sci::Matrix<double, Layout, Allocator>
inv(const Sci::Matrix<double, Layout, Allocator>& a)
{
    Sci::Matrix<double, Layout, Allocator> res(a.extent(0), a.extent(1));
    inv(a.view(), res.view());
    return res;
}

} // namespace Linalg
} // namespace Sci

#endif // SCILIB_LINALG_INV_H


/* specfunc/hyperg.c
 *
 * Copyright (C) 1996, 1997, 1998, 1999, 2000 Gerard Jungman
 *
 * 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 2 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., 675 Mass Ave, Cambridge, MA 02139, USA.
 */

/* Author:  G. Jungman */

/* Miscellaneous implementations of use
 * for evaluation of hypergeometric functions.
 */
#include <config.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_errno.h>
#include "gsl_sf_exp.h"
#include "gsl_sf_gamma.h"

#include "error.h"
#include "hyperg.h"

#define SUM_LARGE  (1.0e-5*GSL_DBL_MAX)

int
gsl_sf_hyperg_1F1_series_e(const double a, const double b, const double x,
                              gsl_sf_result * result
                              )
{
  double an  = a;
  double bn  = b;
  double n   = 1.0;
  double del = 1.0;
  double abs_del = 1.0;
  double max_abs_del = 1.0;
  double sum_val = 1.0;
  double sum_err = 0.0;

  while(abs_del/fabs(sum_val) > GSL_DBL_EPSILON) {
    double u, abs_u;

    if(bn == 0.0) {
      DOMAIN_ERROR(result);
    }
    if(an == 0.0 || n > 1000.0) {
      result->val  = sum_val;
      result->err  = sum_err;
      result->err += 2.0 * GSL_DBL_EPSILON * n * fabs(sum_val);
      return GSL_SUCCESS;
    }

    u = x * (an/(bn*n));
    abs_u = fabs(u);
    if(abs_u > 1.0 && max_abs_del > GSL_DBL_MAX/abs_u) {
      result->val = sum_val;
      result->err = fabs(sum_val);
      GSL_ERROR ("overflow", GSL_EOVRFLW);
    }
    del *= u;
    sum_val += del;
    if(fabs(sum_val) > SUM_LARGE) {
      result->val = sum_val;
      result->err = fabs(sum_val);
      GSL_ERROR ("overflow", GSL_EOVRFLW);
    }

    abs_del = fabs(del);
    max_abs_del = GSL_MAX_DBL(abs_del, max_abs_del);
    sum_err += 2.0*GSL_DBL_EPSILON*abs_del;

    an += 1.0;
    bn += 1.0;
    n  += 1.0;
  }

  result->val  = sum_val;
  result->err  = sum_err;
  result->err += abs_del;
  result->err += 2.0 * GSL_DBL_EPSILON * n * fabs(sum_val);

  return GSL_SUCCESS;
}

int
gsl_sf_hyperg_1F1_large_b_e(const double a, const double b, const double x, gsl_sf_result * result)
{
  if(fabs(x/b) < 1.0) {
    const double u = x/b;
    const double v = 1.0/(1.0-u);
    const double pre = pow(v,a);
    const double uv  = u*v;
    const double uv2 = uv*uv;
    const double t1  = a*(a+1.0)/(2.0*b)*uv2;
    const double t2a = a*(a+1.0)/(24.0*b*b)*uv2;
    const double t2b = 12.0 + 16.0*(a+2.0)*uv + 3.0*(a+2.0)*(a+3.0)*uv2;
    const double t2  = t2a*t2b;
    result->val  = pre * (1.0 - t1 + t2);
    result->err  = pre * GSL_DBL_EPSILON * (1.0 + fabs(t1) + fabs(t2));
    result->err += 2.0 * GSL_DBL_EPSILON * fabs(result->val);
    return GSL_SUCCESS;
  }
  else {
    DOMAIN_ERROR(result);
  }
}

int
gsl_sf_hyperg_U_large_b_e(const double a, const double b, const double x,
                             gsl_sf_result * result,
			     double * ln_multiplier
			     )
{
  double N   = floor(b);  /* b = N + eps */
  double eps = b - N;

  if(fabs(eps) < GSL_SQRT_DBL_EPSILON) {
    double lnpre_val;
    double lnpre_err;
    gsl_sf_result M;
    if(b > 1.0) {
      double tmp = (1.0-b)*log(x);
      gsl_sf_result lg_bm1;
      gsl_sf_result lg_a;
      gsl_sf_lngamma_e(b-1.0, &lg_bm1);
      gsl_sf_lngamma_e(a, &lg_a);
      lnpre_val = tmp + x + lg_bm1.val - lg_a.val;
      lnpre_err = lg_bm1.err + lg_a.err + GSL_DBL_EPSILON * (fabs(x) + fabs(tmp));
      gsl_sf_hyperg_1F1_large_b_e(1.0-a, 2.0-b, -x, &M);
    }
    else {
      gsl_sf_result lg_1mb;
      gsl_sf_result lg_1pamb;
      gsl_sf_lngamma_e(1.0-b, &lg_1mb);
      gsl_sf_lngamma_e(1.0+a-b, &lg_1pamb);
      lnpre_val = lg_1mb.val - lg_1pamb.val;
      lnpre_err = lg_1mb.err + lg_1pamb.err;
      gsl_sf_hyperg_1F1_large_b_e(a, b, x, &M);
    }

    if(lnpre_val > GSL_LOG_DBL_MAX-10.0) {
      result->val  = M.val;
      result->err  = M.err;
      *ln_multiplier = lnpre_val;
      GSL_ERROR ("overflow", GSL_EOVRFLW);
    }
    else {
      gsl_sf_result epre;
      int stat_e = gsl_sf_exp_err_e(lnpre_val, lnpre_err, &epre);
      result->val  = epre.val * M.val;
      result->err  = epre.val * M.err + epre.err * fabs(M.val);
      result->err += 2.0 * GSL_DBL_EPSILON * fabs(result->val);
      *ln_multiplier = 0.0;
      return stat_e;
    }
  }
  else {
    double omb_lnx = (1.0-b)*log(x);
    gsl_sf_result lg_1mb;    double sgn_1mb;
    gsl_sf_result lg_1pamb;  double sgn_1pamb;
    gsl_sf_result lg_bm1;    double sgn_bm1;
    gsl_sf_result lg_a;      double sgn_a;
    gsl_sf_result M1, M2;
    double lnpre1_val, lnpre2_val;
    double lnpre1_err, lnpre2_err;
    double sgpre1, sgpre2;
    gsl_sf_hyperg_1F1_large_b_e(    a,     b, x, &M1);
    gsl_sf_hyperg_1F1_large_b_e(1.0-a, 2.0-b, x, &M2);

    gsl_sf_lngamma_sgn_e(1.0-b,   &lg_1mb,   &sgn_1mb);
    gsl_sf_lngamma_sgn_e(1.0+a-b, &lg_1pamb, &sgn_1pamb);

    gsl_sf_lngamma_sgn_e(b-1.0, &lg_bm1, &sgn_bm1);
    gsl_sf_lngamma_sgn_e(a,     &lg_a,   &sgn_a);

    lnpre1_val = lg_1mb.val - lg_1pamb.val;
    lnpre1_err = lg_1mb.err + lg_1pamb.err;
    lnpre2_val = lg_bm1.val - lg_a.val - omb_lnx - x;
    lnpre2_err = lg_bm1.err + lg_a.err + GSL_DBL_EPSILON * (fabs(omb_lnx)+fabs(x));
    sgpre1 = sgn_1mb * sgn_1pamb;
    sgpre2 = sgn_bm1 * sgn_a;

    if(lnpre1_val > GSL_LOG_DBL_MAX-10.0 || lnpre2_val > GSL_LOG_DBL_MAX-10.0) {
      double max_lnpre_val = GSL_MAX(lnpre1_val,lnpre2_val);
      double max_lnpre_err = GSL_MAX(lnpre1_err,lnpre2_err);
      double lp1 = lnpre1_val - max_lnpre_val;
      double lp2 = lnpre2_val - max_lnpre_val;
      double t1  = sgpre1*exp(lp1);
      double t2  = sgpre2*exp(lp2);
      result->val  = t1*M1.val + t2*M2.val;
      result->err  = fabs(t1)*M1.err + fabs(t2)*M2.err;
      result->err += GSL_DBL_EPSILON * exp(max_lnpre_err) * (fabs(t1*M1.val) + fabs(t2*M2.val));
      result->err += 2.0 * GSL_DBL_EPSILON * fabs(result->val);
      *ln_multiplier = max_lnpre_val;
      GSL_ERROR ("overflow", GSL_EOVRFLW);
    }
    else {
      double t1 = sgpre1*exp(lnpre1_val);
      double t2 = sgpre2*exp(lnpre2_val);
      result->val  = t1*M1.val + t2*M2.val;
      result->err  = fabs(t1) * M1.err + fabs(t2)*M2.err;
      result->err += GSL_DBL_EPSILON * (exp(lnpre1_err)*fabs(t1*M1.val) + exp(lnpre2_err)*fabs(t2*M2.val));
      result->err += 2.0 * GSL_DBL_EPSILON * fabs(result->val);
      *ln_multiplier = 0.0;
      return GSL_SUCCESS;
    }
  }
}

/* [Carlson, p.109] says the error in truncating this asymptotic series
 * is less than the absolute value of the first neglected term.
 *
 * A termination argument is provided, so that the series will
 * be summed at most up to n=n_trunc. If n_trunc is set negative,
 * then the series is summed until it appears to start diverging.
 */
int
gsl_sf_hyperg_2F0_series_e(const double a, const double b, const double x,
                              int n_trunc,
                              gsl_sf_result * result
                              )
{
  const int maxiter = 2000;
  double an = a;
  double bn = b;
  double n   = 1.0;
  double sum = 1.0;
  double del = 1.0;
  double abs_del = 1.0;
  double max_abs_del = 1.0;
  double last_abs_del = 1.0;

  while(abs_del/fabs(sum) > GSL_DBL_EPSILON && n < maxiter) {

    double u = an * (bn/n * x);
    double abs_u = fabs(u);

    if(abs_u > 1.0 && (max_abs_del > GSL_DBL_MAX/abs_u)) {
      result->val = sum;
      result->err = fabs(sum);
      GSL_ERROR ("overflow", GSL_EOVRFLW);
    }

    del *= u;
    sum += del;

    abs_del = fabs(del);

    if(abs_del > last_abs_del) break; /* series is probably starting to grow */

    last_abs_del = abs_del;
    max_abs_del  = GSL_MAX(abs_del, max_abs_del);

    an += 1.0;
    bn += 1.0;
    n  += 1.0;

    if(an == 0.0 || bn == 0.0) break;        /* series terminated */

    if(n_trunc >= 0 && n >= n_trunc) break;  /* reached requested timeout */
  }

  result->val = sum;
  result->err = GSL_DBL_EPSILON * n + abs_del;
  if(n >= maxiter)
    GSL_ERROR ("error", GSL_EMAXITER);
  else
    return GSL_SUCCESS;
}


#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION

#include <stdio.h>
#include <stdlib.h>
#include <gsl/gsl_multifit.h>
#include <gsl/gsl_vector.h>
#include <gsl/gsl_matrix.h>
#include <Python.h>
#include <numpy/arrayobject.h>

void mark_debug (void)
{
    static int number;
    printf("---> Mark %d <---\n", number);
    fflush(stdout);
    number++;
}

void label_debug (const char * label)
{
    printf("---> Mark %s <---\n", label);
    fflush(stdout);
}

/* =========================== Pure-C source code =========================== */

/* Regularization workspace */
struct rglz_wksp {
    double lamb; /* Lambda value */
    gsl_vector *Ldiag; /* Pointer to vector with diagonal values of the
			      * regularization matrix L. */
    const gsl_vector *Y; /* Pointer to original Y */
    gsl_vector *Ys; /* Ponter to Y vector of the standard form. */
    const gsl_matrix *X; /* Pointer to original X matrix  */
    gsl_matrix *Xs; /* Pointer to X matrix of the standard form. */
    gsl_multifit_linear_workspace *lwksp; /* Pointer to GSL workspace where SVDs are stored. */
    /* L-curve data */
    struct lcurve {
	gsl_vector *reg_param;
	gsl_vector *rho;
	gsl_vector *eta;
    } lcurve;
};
typedef struct rglz_wksp rglz_wksp;

/* Allocate and initialize inner variables of a regularization workspace. */
rglz_wksp * rglz_init(const double *Ldiag, const int p,
		      const gsl_vector *Y, const gsl_matrix *X, const gsl_vector *W,
		      gsl_multifit_linear_workspace *lwksp)
{
    rglz_wksp *rw;

    /* Allocate space for the workspace */
    rw = (rglz_wksp*) malloc(sizeof(rglz_wksp));

    rw->lamb = 0.0; /* Initial value */
    rw->lcurve.reg_param = NULL;
    rw->lcurve.rho = NULL;
    rw->lcurve.eta = NULL;

    /* Allocate space for the standard-form matrixes */
    rw->Xs = gsl_matrix_alloc(X->size1, X->size2);
    rw->Ys = gsl_vector_alloc(Y->size);

    /* Initialize the regularization weights Ldiag */
    rw->Ldiag = gsl_vector_alloc(p);
    for (int i = 0; i < p; i++)
        gsl_vector_set(rw->Ldiag, i, Ldiag[i]);

    /* Convert the system to standard form */
    gsl_multifit_linear_wstdform1(rw->Ldiag, X, W, Y, rw->Xs, rw->Ys, lwksp);

    /* Calculate SVD of standard-form matrix Xs */
    gsl_multifit_linear_svd(rw->Xs, lwksp);

    rw->lwksp = lwksp;
    rw->X = X;
    rw->Y = Y;
    return rw;
}

void rglz_estimate_lambda(rglz_wksp *rw, const int lcsize)
{
    size_t idx; /* Temporary index to store corner of the L-curve */
    /* Allocate vectors */
    rw->lcurve.reg_param = gsl_vector_alloc(lcsize);
    rw->lcurve.rho = gsl_vector_alloc(lcsize);
    rw->lcurve.eta = gsl_vector_alloc(lcsize);
    /* Create Lcurve */
    gsl_multifit_linear_lcurve(rw->Y, rw->lcurve.reg_param, rw->lcurve.rho, rw->lcurve.eta, rw->lwksp);
#ifdef PRINT_LCURVE
    for (int i = 0; i < lcsize; i++)
    {
        printf("%12.4e%12.4e\n",
                gsl_vector_get(rw->lcurve.rho, i),
                gsl_vector_get(rw->lcurve.eta, i));
    }
#endif
    /* Estimate corner */
    gsl_multifit_linear_lcorner(rw->lcurve.rho, rw->lcurve.eta, &idx);
    /* Set lambda estimate */
    rw->lamb = gsl_vector_get(rw->lcurve.reg_param, idx);
    fprintf(stdout, "From gslpy_ridge: lambda = %.4e\n", rw->lamb);
}

void rglz_solve(gsl_vector *c, double *rnorm, double *snorm, rglz_wksp *rw)
{
    /* Solve in standard form */
    gsl_multifit_linear_solve(rw->lamb, rw->Xs, rw->Ys, c, rnorm, snorm, rw->lwksp);
    /* Convert back to original parameters */
    gsl_multifit_linear_genform1(rw->Ldiag, c, c, rw->lwksp);
    return;
}

inline void rglz_set_lambda(rglz_wksp *rw, const int lamb)
{
    rw->lamb = lamb;
}

void rglz_close(rglz_wksp *rw)
{
    gsl_vector_free(rw->Ldiag);
    gsl_matrix_free(rw->Xs);
    gsl_vector_free(rw->Ys);
    if (rw->lcurve.reg_param != NULL)
	gsl_vector_free(rw->lcurve.reg_param);
    if (rw->lcurve.rho != NULL)
	gsl_vector_free(rw->lcurve.rho);
    if (rw->lcurve.eta != NULL)
	gsl_vector_free(rw->lcurve.eta);
    free(rw);
}

int
gslpy_ridge_regression(const gsl_vector * y,
		       const gsl_matrix * X,
		       const gsl_vector * w,
		       const double lambda,
		       const double * L,
		       const unsigned long Lcurve_size,
		       gsl_vector * coefs)
{

    unsigned long n_samples = X->size1;
    unsigned long n_features = X->size2;

    double snorm;
    double chisq;

    gsl_multifit_linear_workspace
	* gmlw = gsl_multifit_linear_alloc(n_samples, n_features);

    rglz_wksp * rw = rglz_init(L, n_features, y, X, w, gmlw);

    /* If the user specified gamma, do not attempt to optimize it. */
    if (lambda >= 0)
	rglz_set_lambda(rw, lambda);
    else
	rglz_estimate_lambda(rw, Lcurve_size);

    rglz_solve(coefs, &chisq, &snorm, rw);

    rglz_close(rw);

    gsl_multifit_linear_free(gmlw);
    return 0;
}

/* ============================= Python wrapper ============================= */

static PyObject*
gslpy_ridge_fit(PyObject* self, PyObject* args, PyObject* kwargs)
{
    static char * keywords[] = {
	"X",
	"y",
	"sample_weight",
	"lamb",
	"ldiag",
	"lcurve_size",
	NULL,
    };

    PyArrayObject
	*X=NULL,
	*y=NULL,
	*w=NULL,
	*Ldiag=NULL,
	*coefs=NULL;

    /* Negative lambda means that lambda is optimized */
    double lambda = -1;
    unsigned long Lcurve_size = 10000;

    /* Flag to check if things are set */
    int wei_set = 0;
    int L_set = 0;

    /* GSL stuff */
    gsl_vector * g_y;
    gsl_matrix * g_X;
    gsl_vector * g_w;
    gsl_vector * g_coefs;

    if (!PyArg_ParseTupleAndKeywords(args,
                                     kwargs,
                                     "O!O!|O!dO!k",
                                     keywords,
                                     &PyArray_Type, &X,
				     &PyArray_Type, &y,
				     &PyArray_Type, &w,
				     &lambda,
				     &PyArray_Type, &Ldiag,
				     &Lcurve_size)) {
        return NULL;
    }

    /* If weights are not given, set them as 1 */
    if (w == NULL)
    {
	wei_set = 1;
	w = (PyArrayObject*) PyArray_ZEROS(PyArray_NDIM(y),
					   PyArray_DIMS(y),
					   NPY_DOUBLE,
					   0);

	npy_double * w_data = PyArray_DATA(w);

	for (npy_intp i = 0; i < PyArray_SIZE(w); ++i)
	    w_data[i] = 1.0;
    }

    /* If Ldiag is not given, set it to identity. */
    if (Ldiag == NULL)
    {
	npy_intp Lshape[] = {PyArray_DIM(X,1)};
	L_set = 1;
	Ldiag = (PyArrayObject*) PyArray_ZEROS(1,
					       Lshape,
					       NPY_DOUBLE,
					       0);

	for (npy_intp i = 0; i < PyArray_DIM(X, 1); i++)
	    ((double *) PyArray_DATA(Ldiag))[i] = 1.0;
    }

    /* Convert PyArray objects to gsl types */
    g_y = gsl_vector_calloc(PyArray_DIM(y, 0));
    g_w = gsl_vector_calloc(PyArray_DIM(w, 0));
    g_X = gsl_matrix_calloc(PyArray_DIMS(X)[0], PyArray_DIMS(X)[1]);
    g_coefs = gsl_vector_calloc(PyArray_DIM(X, 1));

    for (npy_intp i = 0; i < (npy_intp) g_y->size; i++)
    {
	gsl_vector_set(g_y, i, ((double *) PyArray_DATA(y))[i]);
	gsl_vector_set(g_w, i, ((double *) PyArray_DATA(w))[i]);
    }

    for (npy_intp i = 0; i < PyArray_SIZE(X); i++)
	g_X->data[i] = ((double *) PyArray_DATA(X))[i];

    /* Run regression */
    gslpy_ridge_regression(g_y, g_X,
			   g_w, lambda,
			   (double *) PyArray_DATA(Ldiag),
			   Lcurve_size, g_coefs);

    /* Convert g_coefs back to a Python object */
    /* Local scope just to delete ``coefs_shape``after. */
    {
	npy_intp coefs_shape[] = {g_coefs->size, 1};
	coefs = (PyArrayObject *) PyArray_ZEROS(2,
						coefs_shape,
						NPY_DOUBLE,
						0);
	for (npy_intp i = 0; i < (npy_intp) g_coefs->size; i++)
	    ((double *) PyArray_DATA(coefs))[i] = gsl_vector_get(g_coefs, i);

    }

    if (wei_set)
	Py_XDECREF(w);

    if (L_set)
	Py_XDECREF(Ldiag);

    gsl_matrix_free(g_X);
    gsl_vector_free(g_y);
    gsl_vector_free(g_w);
    gsl_vector_free(g_coefs);

    return (PyObject *) coefs;
}

/* =========================== Python boilerplate =========================== */

PyDoc_STRVAR(module_doc,
	     "Ridge regression model methods using the GNU Scientific Library.");

PyDoc_STRVAR(fit_doc,
	     "Fits a ridge regression model to data.\n"
	     "fit(X, y, sample_weights=None, lambda=None, ldiag=None, Lcurve_size=10000).\n");

PyMethodDef methods[] = {
    {
	"fit",
	(PyCFunction) gslpy_ridge_fit,
	METH_VARARGS | METH_KEYWORDS,
	fit_doc,
    },
    {NULL},
};

PyModuleDef gslpyridge_module = {
    PyModuleDef_HEAD_INIT,
    "gslpyridge",
    module_doc,
    -1,
    methods,
    NULL,
    NULL,
    NULL,
    NULL
};

PyMODINIT_FUNC
PyInit_gslpyridge(void)
{
    import_array();
    return PyModule_Create(&gslpyridge_module);
}


/*
Matrix multiplication.

Based on the amazing:
https://github.com/HandsOnOpenCL/Exercises-Solutions/tree/a908ac3f0fadede29f2735eb1264b0db7f4311a0/Solutions/Exercise08

The most basic / useful application where OpenCL might be faster than CPU.

TODO: make a SERIOUS matrix implementation. Also compare with existing SERIOUS CPU and GPU implementations:

- http://stackoverflow.com/questions/1907557/optimized-matrix-multiplication-in-c
- http://stackoverflow.com/questions/12289235/simple-and-fast-matrix-vector-multiplication-in-c-c
- https://www.quora.com/What-is-the-best-way-to-multiply-two-matrices-in-C++
- http://stackoverflow.com/questions/25900312/optimizing-batched-matrix-multiplication-opencl-code

Serious CPU implementation means it considers:

- caching
- SIMD

Articles:

- http://www.netlib.org/utk/papers/autoblock/node2.html
- http://codereview.stackexchange.com/questions/101144/simd-matrix-multiplication
*/

#include "common.h"

#include <cblas.h>
#include <clBLAS.h>

#define NELEMS(a) sizeof(a)/sizeof(a[0])
#define UNUSED(x) (void)(x)

/* TODO how to auto tune this?
 * GCC 6 was not smart enough to use widest type automatically:
 * anything larger than the widest type just dropped perf drastically. */
#define VECTOR_NELEMS 4
#define VECTOR_SIZEOF (VECTOR_NELEMS * sizeof(F))

typedef cl_float F;
typedef F Vec __attribute__ ((vector_size(VECTOR_SIZEOF)));
/* Stuff reused across OpenCL calls.
 * It would not be fair to consider their alocation time for each call,
 * as the functions will be called millions of times.
 *
 * TODO: also factor out the kernel compile time.
 * Currently recompiled every time.
 * Implementations already cache compiled kernels, but we would need
 * to call each function multiple times.
 * */
typedef struct {
    Common common;
    /* We only factor out the allocation of those buffers.
     * Data transfer time with clEnqueueWriteBuffer must still be considered. */
    cl_mem buf_a, buf_b, buf_c;
} Cache;
typedef void (*MatMul)(const F *A, const F *B, F *C, size_t n, Cache *cache);

void cl_buf_init(Cache *cache, size_t mat_sizeof) {
    cache->buf_a = clCreateBuffer(cache->common.context, CL_MEM_READ_ONLY, mat_sizeof, NULL, NULL);
    cache->buf_b = clCreateBuffer(cache->common.context, CL_MEM_READ_ONLY, mat_sizeof, NULL, NULL);
    cache->buf_c = clCreateBuffer(cache->common.context, CL_MEM_WRITE_ONLY, mat_sizeof, NULL, NULL);
}

void cl_buf_deinit(Cache *cache) {
    clReleaseMemObject(cache->buf_a);
    clReleaseMemObject(cache->buf_b);
    clReleaseMemObject(cache->buf_c);
}

/* No, this was not created for debugging, my code is flawless from the first try. */
void mat_print(const F *A, size_t n) {
    size_t i, j;
    for (i = 0; i < n; ++i) {
        for (j = 0; j < n; ++j) {
            printf("%f ", A[i*n+j]);
        }
        puts("");
    }
}

/* Zero a matrix. */
void mat_zero(F *A, size_t n) {
    size_t i, n2;
    n2 = n*n;
    for (i = 0; i < n2; ++i) {
        A[i] = 0.0;
    }
}

/* Initialize a random matrix. */
void mat_rand(F *A, size_t n) {
    size_t i, n2;
    n2 = n*n;
    for (i = 0; i < n2; ++i) {
        A[i] = ((float)rand()) / ((float)RAND_MAX);
    }
}

/* Initialize a random matrix. */
void mat_trans(F *A, size_t n) {
    size_t i, j, i1, i2;
    F tmp;
    for (i = 0; i < n; ++i) {
        for (j = 0; j < i; ++j) {
            i1 = i*n+j;
            i2 = j*n+i;
            tmp = A[i1];
            A[i1] = A[i2];
            A[i2] = tmp;
        }
    }
}

/* Check if two matrices are equal with given mean squared err_maxor. */
int mat_eq(const F *A, const F *B, size_t n) {
    const F err_max = 10e-3;
    F err, diff, a, b;
    size_t i, i_max;

    err = 0.0;
    i_max = n*n;
    for (i = 0; i < i_max; ++i) {
        a = A[i];
        b = B[i];
        diff = a - b;
        err += diff * diff;
    }
    return (sqrt(err) / i_max) < err_max;
}

void mat_assert_eq(const F *A, const F *B, size_t n) {
    if (!mat_eq(A, B, n)) {
        puts("");
        mat_print(A, n);
        puts("");
        mat_print(B, n);
        assert(0);
    }
}

size_t zmin(size_t x, size_t y) {
    return (x < y) ? x : y;
}

/* C = A*B, width n, naive. */
void mat_mul_cpu(const F *A, const F *B, F *C, size_t n, Cache *cache) {
    F tmp;
    size_t i, j, k;
    UNUSED(cache);

    for (i = 0; i < n; ++i) {
        for (j = 0; j < n; ++j) {
            tmp = 0.0;
            for (k = 0; k < n; ++k) {
                tmp += A[i*n+k] * B[k*n+j];
            }
            C[i*n+j] = tmp;
        }
    }
}

/* Transpose matrix B to increase cache hits. */
void mat_mul_cpu_trans(const F *A, const F *B, F *C, size_t n, Cache *cache) {
    F tmp;
    size_t i, j, k;
    UNUSED(cache);

    mat_trans((F*)B, n);
    for (i = 0; i < n; ++i) {
        for (j = 0; j < n; ++j) {
            tmp = 0.0;
            for (k = 0; k < n; ++k) {
                tmp += A[i*n+k] * B[j*n+k];
            }
            C[i*n+j] = tmp;
        }
    }
    mat_trans((F*)B, n);
}

/* Zero vector. */
void vec_zero(Vec *vec, size_t vec_n) {
    size_t i;
    for (i = 0; i < vec_n; ++i) {
        (*vec)[i] = 0.0;
    }
}

/* Load vector from array. */
void vec_load(Vec *vec, size_t vec_n, const F* A, size_t A_base) {
    size_t i;
    for (i = 0; i < vec_n; ++i) {
        (*vec)[i] = A[A_base + i];
    }
}

/* Sum elements of the vector. */
F vec_sum(Vec vec, size_t vec_n) {
    size_t i;
    F sum;
    sum = 0.0;
    for (i = 0; i < vec_n; ++i) {
        sum += vec[i];
    }
    return sum;
}

/* Transpose matrix B to both:
 *
 * - increase cache hits
 * - simd GCC vector extensions which is made possible.
 *   by the transposition, to increase likelyhood of SIMDs.
 *
 * Note that GCC 6 O=3 is smart enough to use SIMD
 * even for the naive CPU method. However this was still way faster.
 * */
void mat_mul_cpu_trans_vec(const F *A, const F *B, F *C, size_t n, Cache *cache) {
    F tmpf;
    size_t i, j, k, k_max, ai, bi;
    Vec tmp, a, b;
    UNUSED(cache);

    mat_trans((F*)B, n);
    k_max = (n / VECTOR_NELEMS) * VECTOR_NELEMS;
    for (i = 0; i < n; ++i) {
        for (j = 0; j < n; ++j) {
            vec_zero(&tmp, VECTOR_NELEMS);
            for (k = 0; k < k_max; k += VECTOR_NELEMS) {
                ai = i * n + k;
                bi = j * n + k;
                vec_load(&a, VECTOR_NELEMS, A, ai);
                vec_load(&b, VECTOR_NELEMS, B, bi);
                tmp += a * b;
            }
            tmpf = 0.0;
            for (; k < n; ++k) {
                tmpf += A[i*n+k] * B[j*n+k];
            }
            C[i*n+j] = vec_sum(tmp, VECTOR_NELEMS) + tmpf;
        }
    }
    mat_trans((F*)B, n);
}

/* Blocked matrix multiplication.
 * TODO slower than transpose, sometimes similar timing to naive.
 * Why do they say that this is better?
 * */
void mat_mul_cpu_block(const F *A, const F *B, F *C, size_t n, Cache *cache) {
    size_t ib, jb, kb, i, j, k, i_max, j_max, k_max, nb ;
    F tmp;
    UNUSED(cache);

    nb = lround(sqrt(n));
    for (ib = 0; ib < n; ib += nb) {
        i_max = zmin(ib + nb, n);
        for (jb = 0; jb < n; jb += nb) {
            k_max = zmin(jb + nb, n);
            for (kb = 0; kb < n; kb += nb) {
                j_max = zmin(kb + nb, n);

                /* TODO would be cool to recurse here, but would require more offset parameters,
                 * likely similar to tose of CBLAS. */
                for (i = ib; i < i_max; ++i) {
                    for (j = kb; j < j_max; ++j) {
                        tmp = 0.0;
                        for (k = jb; k < k_max; ++k) {
                            tmp += A[i*n+k] * B[k*n+j];
                        }
                        C[i*n+j] += tmp;
                    }
                }
            }
        }
    }
}

/* The golden single thread CPU standard. */
void mat_mul_cpu_cblas(const F *A, const F *B, F *C, size_t n, Cache *cache) {
    UNUSED(cache);
    cblas_sgemm(
        CblasRowMajor,
        CblasNoTrans,
        CblasNoTrans,
        n,
        n,
        n,
        1.0,
        A,
        n,
        B,
        n,
        0.0,
        C,
        n
    );
}

/* Simplest possible CL implementation. No speedup. */
void mat_mul_cl(const F *A, const F *B, F *C, size_t n, Cache *cache) {
    cl_uint ncl;
    size_t global_work_size[2], mat_sizeof;

    /* Setup variables. */
    global_work_size[0] = n;
    global_work_size[1] = n;
    mat_sizeof = n * n * sizeof(F);
    ncl = n;

    /* Run kernel. */
    common_create_kernel_file(&cache->common, "matmul.cl", NULL);
    clEnqueueWriteBuffer(cache->common.command_queue, cache->buf_a, CL_TRUE, 0, mat_sizeof, (F*)A, 0, NULL, NULL);
    clEnqueueWriteBuffer(cache->common.command_queue, cache->buf_b, CL_TRUE, 0, mat_sizeof, (F*)B, 0, NULL, NULL);
    clSetKernelArg(cache->common.kernel, 0, sizeof(cache->buf_a), &cache->buf_a);
    clSetKernelArg(cache->common.kernel, 1, sizeof(cache->buf_b), &cache->buf_b);
    clSetKernelArg(cache->common.kernel, 2, sizeof(cache->buf_c), &cache->buf_c);
    clSetKernelArg(cache->common.kernel, 3, sizeof(ncl), &ncl);
    clEnqueueNDRangeKernel(cache->common.command_queue, cache->common.kernel, 2, NULL, global_work_size, NULL, 0, NULL, NULL);
    clFlush(cache->common.command_queue);
    clFinish(cache->common.command_queue);
    clEnqueueReadBuffer(cache->common.command_queue, cache->buf_c, CL_TRUE, 0, mat_sizeof, C, 0, NULL, NULL);
}

/* Cache rows in private memory. Drastic speedups expected over naive CPU. */
void mat_mul_cl_row_priv(const F *A, const F *B, F *C, size_t n, Cache *cache) {
    char options[256];
    cl_uint ncl;
    size_t global_work_size, mat_sizeof;

    /* Setup variables. */
    global_work_size = n;
    mat_sizeof = n * n * sizeof(F);
    ncl = n;

    /* Run kernel. */
    snprintf(options, sizeof(options), "-DPRIV_ROW_SIZE=%ju", n);
    common_create_kernel_file(&cache->common, "matmul_row_priv.cl", options);
    clEnqueueWriteBuffer(cache->common.command_queue, cache->buf_a, CL_TRUE, 0, mat_sizeof, (F*)A, 0, NULL, NULL);
    clEnqueueWriteBuffer(cache->common.command_queue, cache->buf_b, CL_TRUE, 0, mat_sizeof, (F*)B, 0, NULL, NULL);
    clSetKernelArg(cache->common.kernel, 0, sizeof(cache->buf_a), &cache->buf_a);
    clSetKernelArg(cache->common.kernel, 1, sizeof(cache->buf_b), &cache->buf_b);
    clSetKernelArg(cache->common.kernel, 2, sizeof(cache->buf_c), &cache->buf_c);
    clSetKernelArg(cache->common.kernel, 3, sizeof(ncl), &ncl);
    clEnqueueNDRangeKernel(cache->common.command_queue, cache->common.kernel, 1, NULL, &global_work_size, NULL, 0, NULL, NULL);
    clFlush(cache->common.command_queue);
    clFinish(cache->common.command_queue);
    clEnqueueReadBuffer(cache->common.command_queue, cache->buf_c, CL_TRUE, 0, mat_sizeof, C, 0, NULL, NULL);
}

/* Let's see if this is any different from local memory.
 * Outcome: much slower than private memory, slower than naive method. */
void mat_mul_cl_row_local(const F *A, const F *B, F *C, size_t n, Cache *cache) {
    cl_uint ncl;
    size_t global_work_size, local_work_size, mat_sizeof;

    /* Setup variables. */
    /* Cannot be larger than 1 on this example, otherwise memory conflicts
     * will happen between work items. */
    local_work_size = 1;
    global_work_size = n;
    mat_sizeof = n * n * sizeof(F);
    ncl = n;

    /* Run kernel. */
    common_create_kernel_file(&cache->common, "matmul_row_local.cl", NULL);
    clEnqueueWriteBuffer(cache->common.command_queue, cache->buf_a, CL_TRUE, 0, mat_sizeof, (F*)A, 0, NULL, NULL);
    clEnqueueWriteBuffer(cache->common.command_queue, cache->buf_b, CL_TRUE, 0, mat_sizeof, (F*)B, 0, NULL, NULL);
    clSetKernelArg(cache->common.kernel, 0, sizeof(cache->buf_a), &cache->buf_a);
    clSetKernelArg(cache->common.kernel, 1, sizeof(cache->buf_b), &cache->buf_b);
    clSetKernelArg(cache->common.kernel, 2, sizeof(cache->buf_c), &cache->buf_c);
    clSetKernelArg(cache->common.kernel, 3, n * sizeof(F), NULL);
    clSetKernelArg(cache->common.kernel, 4, sizeof(ncl), &ncl);
    clEnqueueNDRangeKernel(cache->common.command_queue, cache->common.kernel, 1, NULL, &global_work_size, &local_work_size, 0, NULL, NULL);
    clFlush(cache->common.command_queue);
    clFinish(cache->common.command_queue);
    clEnqueueReadBuffer(cache->common.command_queue, cache->buf_c, CL_TRUE, 0, mat_sizeof, C, 0, NULL, NULL);
}

/* Like row private, but to reduce global memory accesses,
 * we copy only once per work group to local memory.
 *
 * Each work group contains a few rows of A.
 *
 * We load the first column, multiply all rows by that column, synrhconize,
 * load another column, and so on.
 *
 * This leads to a thread blockage / memory access tradeoff.
 *
 * We make work groups as large as possible to reload memory less times. */
void mat_mul_cl_row_priv_col_local(const F *A, const F *B, F *C, size_t n, Cache *cache) {
    char options[256];
    cl_uint ncl;
    size_t global_work_size, local_work_size, mat_sizeof;

    /* Setup variables. */
    global_work_size = n;
    mat_sizeof = n * n * sizeof(F);
    ncl = n;

    /* Run kernel. */
    snprintf(options, sizeof(options), "-DPRIV_ROW_SIZE=%ju", n);
    common_create_kernel_file(&cache->common, "matmul_row_priv_col_local.cl", options);
    local_work_size = 0;
    clGetDeviceInfo(cache->common.device, CL_DEVICE_MAX_WORK_GROUP_SIZE, sizeof(local_work_size), &local_work_size, NULL);
    local_work_size = zmin(local_work_size, n);
    clEnqueueWriteBuffer(cache->common.command_queue, cache->buf_a, CL_TRUE, 0, mat_sizeof, (F*)A, 0, NULL, NULL);
    clEnqueueWriteBuffer(cache->common.command_queue, cache->buf_b, CL_TRUE, 0, mat_sizeof, (F*)B, 0, NULL, NULL);
    clSetKernelArg(cache->common.kernel, 0, sizeof(cache->buf_a), &cache->buf_a);
    clSetKernelArg(cache->common.kernel, 1, sizeof(cache->buf_b), &cache->buf_b);
    clSetKernelArg(cache->common.kernel, 2, sizeof(cache->buf_c), &cache->buf_c);
    clSetKernelArg(cache->common.kernel, 3, n * sizeof(F), NULL);
    clSetKernelArg(cache->common.kernel, 4, sizeof(ncl), &ncl);
    clEnqueueNDRangeKernel(cache->common.command_queue, cache->common.kernel, 1, NULL, &global_work_size, &local_work_size, 0, NULL, NULL);
    clFlush(cache->common.command_queue);
    clFinish(cache->common.command_queue);
    clEnqueueReadBuffer(cache->common.command_queue, cache->buf_c, CL_TRUE, 0, mat_sizeof, C, 0, NULL, NULL);
}

/* Copy as many cols from B as possibl to the local memory, only then start multiplying.
 * This leads to less memory barrier hits.
 * How many rows we copy is limited by the local memory size, ideally the entire matrix will fit. */
void mat_mul_cl_row_priv_cols_local(const F *A, const F *B, F *C, size_t n, Cache *cache) {
    char options[256];
    cl_uint ncl, n_local_cols;
    cl_ulong local_mem_size;
    size_t col_size, global_work_size, local_work_size, mat_sizeof;

    /* Setup variables. */
    col_size = n * sizeof(F);
    global_work_size = n;
    mat_sizeof = n * n * sizeof(F);
    ncl = n;

    /* Run kernel. */
    snprintf(options, sizeof(options), "-DPRIV_ROW_SIZE=%ju", n);
    common_create_kernel_file(&cache->common, "matmul_row_priv_cols_local.cl", options);
    local_work_size = 0;
    clGetDeviceInfo(cache->common.device, CL_DEVICE_MAX_WORK_GROUP_SIZE, sizeof(local_work_size), &local_work_size, NULL);
    local_work_size = zmin(local_work_size, n);
    local_mem_size = 0;
    clGetDeviceInfo(cache->common.device, CL_DEVICE_LOCAL_MEM_SIZE, sizeof(local_mem_size), &local_mem_size, NULL);
    /* TODO: can blow up without that - 1. Why?
     * It only reaches the max without it, not crosses, right?
     * So bug in the kernel? */
    n_local_cols = zmin(local_mem_size / col_size, n) - 1;
    /*puts("");*/
    /*printf("max memory %llu\n", (unsigned long long)local_mem_size);*/
    /*printf("n_local_cols %llu\n", (unsigned long long)n_local_cols);*/
    /*printf("memory %llu\n", (unsigned long long)n_local_cols * n * sizeof(F));*/
    clEnqueueWriteBuffer(cache->common.command_queue, cache->buf_a, CL_TRUE, 0, mat_sizeof, (F*)A, 0, NULL, NULL);
    clEnqueueWriteBuffer(cache->common.command_queue, cache->buf_b, CL_TRUE, 0, mat_sizeof, (F*)B, 0, NULL, NULL);
    clSetKernelArg(cache->common.kernel, 0, sizeof(cache->buf_a), &cache->buf_a);
    clSetKernelArg(cache->common.kernel, 1, sizeof(cache->buf_b), &cache->buf_b);
    clSetKernelArg(cache->common.kernel, 2, sizeof(cache->buf_c), &cache->buf_c);
    clSetKernelArg(cache->common.kernel, 3, n_local_cols * col_size, NULL);
    clSetKernelArg(cache->common.kernel, 4, sizeof(ncl), &ncl);
    clSetKernelArg(cache->common.kernel, 5, sizeof(n_local_cols), &n_local_cols);
    clEnqueueNDRangeKernel(cache->common.command_queue, cache->common.kernel, 1, NULL, &global_work_size, &local_work_size, 0, NULL, NULL);
    clFlush(cache->common.command_queue);
    clFinish(cache->common.command_queue);
    clEnqueueReadBuffer(cache->common.command_queue, cache->buf_c, CL_TRUE, 0, mat_sizeof, C, 0, NULL, NULL);
}

void mat_mul_cl_block(const F *A, const F *B, F *C, size_t n, Cache *cache) {
    cl_uint ncl, nblkcl;
    size_t global_work_size[2], local_work_size[2], mat_sizeof, nblk;

    /* Setup variables. */
    global_work_size[0] = n;
    global_work_size[1] = n;
    mat_sizeof = n * n * sizeof(F);
    ncl = n;

    /* Run kernel. */
    common_create_kernel_file(&cache->common, "matmul_block.cl", NULL);
    clGetDeviceInfo(cache->common.device, CL_DEVICE_MAX_WORK_GROUP_SIZE, sizeof(nblk), &nblk, NULL);
    nblk = sqrt(zmin(nblk, n));
    nblk = zmin(nblk, 3);
    nblkcl = nblk;
    local_work_size[0] = nblk;
    local_work_size[1] = nblk;
    clEnqueueWriteBuffer(cache->common.command_queue, cache->buf_a, CL_TRUE, 0, mat_sizeof, (F*)A, 0, NULL, NULL);
    clEnqueueWriteBuffer(cache->common.command_queue, cache->buf_b, CL_TRUE, 0, mat_sizeof, (F*)B, 0, NULL, NULL);
    clSetKernelArg(cache->common.kernel, 0, sizeof(cache->buf_a), &cache->buf_a);
    clSetKernelArg(cache->common.kernel, 1, sizeof(cache->buf_b), &cache->buf_b);
    clSetKernelArg(cache->common.kernel, 2, sizeof(cache->buf_c), &cache->buf_c);
    clSetKernelArg(cache->common.kernel, 3, nblk * nblk * sizeof(F), NULL);
    printf("nblk = %llu\n", (unsigned long long)nblk);
    printf("local memory = %llu\n", (unsigned long long)2 * nblk * nblk * sizeof(F));
    clSetKernelArg(cache->common.kernel, 4, nblk * nblk * sizeof(F), NULL);
    clSetKernelArg(cache->common.kernel, 5, sizeof(ncl), &ncl);
    clSetKernelArg(cache->common.kernel, 6, sizeof(nblkcl), &nblkcl);
    clEnqueueNDRangeKernel(
        cache->common.command_queue, cache->common.kernel, 2, NULL,
        global_work_size, local_work_size, 0, NULL, NULL
    );
    clFlush(cache->common.command_queue);
    clFinish(cache->common.command_queue);
    clEnqueueReadBuffer(cache->common.command_queue, cache->buf_c, CL_TRUE, 0, mat_sizeof, C, 0, NULL, NULL);
}

void mat_mul_cl_clblas(const F *A, const F *B, F *C, size_t n, Cache *cache) {
    cl_event event;
    size_t mat_sizeof;

    mat_sizeof = n * n * sizeof(F);
    clEnqueueWriteBuffer(cache->common.command_queue, cache->buf_a, CL_TRUE, 0, mat_sizeof, (F*)A, 0, NULL, NULL);
    clEnqueueWriteBuffer(cache->common.command_queue, cache->buf_b, CL_TRUE, 0, mat_sizeof, (F*)B, 0, NULL, NULL);
    clblasSgemm(
        clblasRowMajor,
        clblasNoTrans,
        clblasNoTrans,

        n,
        n,
        n,
        1.0,

        cache->buf_a,
        0,
        n,

        cache->buf_b,
        0,
        n,

        0.0,
        cache->buf_c,
        0,
        n,

        1,
        &(cache->common.command_queue),
        0,
        NULL,
        &event
    );
    clWaitForEvents(1, &event);
    clEnqueueReadBuffer(cache->common.command_queue, cache->buf_c, CL_TRUE, 0, mat_sizeof, C, 0, NULL, NULL);
}

double bench(MatMul f, const F *A, const F *B, F *C, F *C_ref, size_t n, Cache *cache) {
    double dt, time;
    mat_zero(C, n);
    time = common_get_nanos();
    f(A, B, C, n, cache);
    dt = common_get_nanos() - time;
    if (C_ref != NULL)
        mat_assert_eq(C, C_ref, n);
    printf("%.3e ", dt);
    fflush(stdout);
    return dt;
}

int main(int argc, char **argv) {
    srand(time(NULL));
    Cache cache;
    double max_runtime;
    /* Overly slow ones commented out by default. */
    MatMul mat_mul_funcs[] = {
        /*mat_mul_cpu,*/
        mat_mul_cpu_trans,
        mat_mul_cpu_trans_vec,
        mat_mul_cpu_block,
        mat_mul_cpu_cblas,
        /*mat_mul_cl,*/
        mat_mul_cl_row_priv,
        mat_mul_cl_row_local,
        mat_mul_cl_row_priv_col_local,
        mat_mul_cl_row_priv_cols_local,
        /* TODO broken for larger matrics, some cells contain trash.
         * Likey some memory overflow problem. */
        /*mat_mul_cl_block,*/
        mat_mul_cl_clblas,
    };
    int first, func_done[NELEMS(mat_mul_funcs)] = {0};
    size_t f, i;
    size_t mat_sizeof;

    /* CLI args. */
    if (argc > 1) {
        max_runtime = strtod(argv[1], NULL);
    } else {
        max_runtime = 1.0;
    }

    common_init(&(cache.common), NULL);

    /* Unit test 2x2. */
    {
        const F A[] = {
            1.0, 2.0,
            3.0, 4.0
        };
        const F B[] = {
            5.0, 6.0,
            7.0, 8.0
        };
        enum N { n = 2 };
        F C[n*n];
        const F C_ref[] = {
            19.0, 22.0,
            43.0, 50.0
        };
        cl_buf_init(&cache, n * n * sizeof(F));
        for (f = 0; f < sizeof(mat_mul_funcs)/sizeof(mat_mul_funcs[0]); ++f) {
            mat_zero(C, n);
            mat_mul_funcs[f](A, B, C, n, &cache);
            mat_assert_eq(C, C_ref, n);
        }
        cl_buf_deinit(&cache);
    }

    /* Unit test 4x4. */
    {
        const F A[] = {
             1.0,  2.0,  3.0,  4.0,
             5.0,  6.0,  7.0,  8.0,
             9.0, 10.0, 11.0, 12.0,
            13.0, 14.0, 15.0, 16.0,
        };
        const F B[] = {
            17.0, 18.0, 19.0, 20.0,
            21.0, 22.0, 23.0, 24.0,
            25.0, 26.0, 27.0, 28.0,
            29.0, 30.0, 31.0, 32.0,
        };
        const F C_ref[] = {
             250.0,  260.0,  270.0,  280.0,
             618.0,  644.0,  670.0,  696.0,
             986.0, 1028.0, 1070.0, 1112.0,
            1354.0, 1412.0, 1470.0, 1528.0,
        };
        enum N { n = 4 };
        F C[n*n];
        cl_buf_init(&cache, n * n * sizeof(F));
        for (f = 0; f < NELEMS(mat_mul_funcs); ++f) {
            mat_zero(C, n);
            mat_mul_funcs[f](A, B, C, n, &cache);
            mat_assert_eq(C, C_ref, n);
        }
        cl_buf_deinit(&cache);
    }

    /* Benchmarks. */
    {
        double dt;
        F *A = NULL, *B = NULL, *C = NULL, *C_ref = NULL, *dst = NULL, *ref = NULL;
        int done;
        size_t n = 2;

        puts("#matmul");
        done = 0;
        while(1) {
            printf("%zu ", (size_t)log2(n));
            mat_sizeof = n * n * sizeof(F);

            /* CPU setup. */
            A = aligned_alloc(VECTOR_SIZEOF, mat_sizeof);
            B = aligned_alloc(VECTOR_SIZEOF, mat_sizeof);
            C = aligned_alloc(VECTOR_SIZEOF, mat_sizeof);
            C_ref = aligned_alloc(VECTOR_SIZEOF, mat_sizeof);
            if (NULL == A || NULL == B || NULL == C) {
                printf("error: could not allocate memory for n = %zu", n);
                break;
            }
            mat_rand(A, n);
            mat_rand(B, n);

            cl_buf_init(&cache, mat_sizeof);
            first = 1;
            for (f = 0; f < NELEMS(mat_mul_funcs); ++f) {
                if (func_done[f]) {
                    printf("%*s", 10, "");
                } else {
                    if (first) {
                        dst = C_ref;
                        ref = NULL;
                        first = 0;
                    } else {
                        dst = C;
                        ref = C_ref;
                    }
                    dt = bench(mat_mul_funcs[f], A, B, dst, ref, n, &cache);
                    if (dt > max_runtime)
                        func_done[f] = 1;
                }
            }
            puts("");
            done = 1;
            for (i = 0; i < NELEMS(mat_mul_funcs); ++i) {
                if (!func_done[i]) {
                    done = 0;
                    break;
                }
            }
            if (done)
                break;
            n *= 2;

            /* CPU deinit. */
            free(A);
            free(B);
            free(C);
            free(C_ref);

            cl_buf_deinit(&cache);
        }
        common_deinit(&cache.common);
    }

    return EXIT_SUCCESS;
}


#if !defined (RAND_H)
#define RAND_H

#include <cstdlib>
#include <iostream>
#include <cmath>
#include <vector>
#include <list>
#include <set>
#include <limits>
#include <algorithm>

#include <gsl/gsl_sf_zeta.h>

using namespace std;

const double TOLERANCE = numeric_limits<double>::epsilon();
const int    INTMAX    = numeric_limits<int>::max();

const int BIGPRIME = 899999963;
const double PI    = 3.1415926535897932384626433832795028841971693993751;

typedef double (*FUNC)(const double, const double, const double);
inline double fx(const double htilde, const double y, const double u)
{
  return (pow(htilde,y+1) - (y+1)*htilde + fabs(2*u-1)*y);
}

inline int Getnode(list<int>& X, int index)
{
  int pos=0;
  for(list<int>::iterator p = X.begin(); p!= X.end(); p++, pos++) {
    if(pos==index)
      return (*p);
  }

  return -1;
}

class Rand {
  double c,cd,cm,u[97];
  int i97,j97 ;
  bool outputReady;
  double output;

  int idum;

 public:
  Rand(int seed) {
    ranmarin(fabs(seed%BIGPRIME));
    outputReady=false;
    idum=-1*seed;

    flag_setFvector = false;
    flag_setPvector = false;
  }

  Rand() {
    ranmarin(1);
    outputReady=false;
    idum=-1;

    flag_setFvector = false;
    flag_setPvector = false;
  }

  void seed(int seed) {
    ranmarin(fabs(seed%BIGPRIME));
    outputReady=false;
    idum=-1*seed;

    flag_setFvector = false;
    flag_setPvector = false;
  }

  double ran() { return rand()/(double)RAND_MAX;}

  double uniform()
  {
    double uni;

    uni=u[i97]-u[j97];
    if (uni<0.0) uni+=1.0;
    u[i97]=uni;
    if (--i97<0) i97=96;
    if (--j97<0) j97=96;
    c-=cd;
    if (c<0.0) c+=cm;
    uni-=c;
    if (uni<0.0) uni+=1.0;
    return(uni);
  }

  double ran1()
  {
    const int IA=16807,IM=2147483647,IQ=127773,IR=2836,NTAB=32;
    const int NDIV=(1+(IM-1)/NTAB);
    const double EPS=3.0e-16,AM=1.0/IM,RNMX=(1.0-EPS);
    static int iy=0;
    static vector<int> iv(NTAB);
    int j,k;
    double temp;

    if (idum <= 0 || !iy) {
      if (-idum < 1) idum=1;
      else idum = -idum;
      for (j=NTAB+7;j>=0;j--) {
	k=idum/IQ;
	idum=IA*(idum-k*IQ)-IR*k;
	if (idum < 0) idum += IM;
	if (j < NTAB) iv[j] = idum;
      }
      iy=iv[0];
    }
    k=idum/IQ;
    idum=IA*(idum-k*IQ)-IR*k;
    if (idum < 0) idum += IM;
    j=iy/NDIV;
    iy=iv[j];
    iv[j] = idum;
    if ((temp=AM*iy) > RNMX) return RNMX;
    else return temp;
  }

  int discrete(int min, int max)
  {
    int ret= (int)(floor(ran1()*(max+1-min)+min));
    return ret;
  }

  void discrete(int min, int max, int& a, int& b)
  {
    a = (int)(floor(ran1()*(max+1-min)+min));

    do {
      b = (int)(floor(ran1()*(max+1-min)+min));
    } while (a==b);

  }

  void discrete(int min, int max, int& a, int& b, int& c)
  {
    a = (int)(floor(ran1()*(max+1-min)+min));

    do {
      b = (int)(floor(ran1()*(max+1-min)+min));
    } while (b==a);

    do {
      c = (int)(floor(ran1()*(max+1-min)+min));
    } while (c==b || c==a);

  }

  void discrete(int min, int max, int& a, int& b, int& c, int& d)
  {
    a = (int)(floor(ran1()*(max+1-min)+min));

    do {
      b = (int)(floor(ran1()*(max+1-min)+min));
    } while (b==a);

    do {
      c = (int)(floor(ran1()*(max+1-min)+min));
    } while (c==b || c==a);

    do {
      d = (int)(floor(ran1()*(max+1-min)+min));
    } while (d==c || d==b || d==a);

  }

  void discrete(int min, int max, int n, set<int>& Set)
  {
    Set.clear();
    for(int i=0; i<n; ) {
      int a = (int)(floor(ran1()*(max+1-min)+min));
      if(Set.find(a)==Set.end()) {
	Set.insert(a);
	i++;
      }
    }
  }

  void discrete(int min, int max, int n, vector<int>& Vec)
  {
    set<int> Set;

    Vec.clear();
    for(int i=0; i<n; ) {
      int a = (int)(floor(ran1()*(max+1-min)+min));
      if(Set.find(a)==Set.end()) {
	Set.insert(a);
	Vec.push_back(a);
	i++;
      }
    }
  }

  int discrete_logbin(int xmin, int xmax)
  {
    double ratio = 1.15;
    if (xmin==0)
      xmin=1;

    int Nbins = (int)(floor(log(xmax/xmin)/log(ratio))) + 1;
    int i = discrete(0, Nbins-1);
    int a = xmin * pow(ratio, (double)i);
    int b = a*ratio;
    if(b>xmax)
      b = xmax;

    int x = discrete(a,b);
    return x;
  }

  void randomsubset(int N, int Nc, vector<int>& S)
  {
    list<int> Lbig;
    for(int i=0; i<N; i++) {
      Lbig.push_back(i);
    }

    int N0= N;
    int M = (2*Nc>N0) ? (N0-Nc) : Nc;

    list<int> Lsmall;
    for(int i=0; i<M; i++) {
      int index = discrete(0, N-1);

      int j = Getnode(Lbig, index);
      Lsmall.push_back(j);
      Lbig.remove(j);
      N--;
    }

    S.clear();
    S.resize(Nc);

    if(2*Nc>N0)
      copy(Lbig.begin(), Lbig.end(), S.begin());
    else
      copy(Lsmall.begin(), Lsmall.end(), S.begin());

  }

  void randomsubset_knuth(int n, int m, vector<int>& S)
  {
    for (int i = 0; i < n; i++)
      if ((discrete(1,n) % (n-i)) < m) {
	S.push_back(i);
	m--;
      }
  }

  void randomsubset_shuf(int n, int m, vector<int>& S)
  {
    int i,j,t;
    vector<int> x(n);
    for (i = 0; i < n; i++)
      x[i] = i;
    for (i = 0; i < m; i++) {
      j = discrete(i, n-1);
      t = x[i];
      x[i] = x[j];
      x[j] = t;
    }
    sort(x.begin(), x.begin()+m);
    for (i = 0; i< m; i++)
      S.push_back(x[i]);
  }

  void randomsubset_shuffle(int n, int m, vector<int>& S)
  {
    vector<int> x(n);
    for(int i=0; i<n; i++)
      x[i] = i;

    random_shuffle(x.begin(), x.end());

    for (int i=0; i<m; i++)
      S.push_back(x[i]);
  }

  double gaussian(double sd)
  {
    double v1,v2,Rsq,z;
    if(outputReady)
      {
	outputReady = false;
	return output;
      }

    do
      {

	v1=(uniform()*2.0)-1.0;
	v2=(uniform()*2.0)-1.0;

	Rsq=v1*v1+v2*v2;
      } while (Rsq>=1.0);

    z=sqrt(-2.0*log(Rsq)/Rsq);
    output = v1*z*sd;
    outputReady=true;
    return v2*z*sd;
  }

  double lorentzian(double sd)
  {
    double v1,v2,z;
    do {
      v2 = uniform();
      v1 = 2.0 * uniform() - 1.0;
    } while ((v1*v1+v2*v2)>1.0);
    z = v2/v1;
    return sd*z/2.0;
  }

  double powerlaw_continuous(double alpha, double xmin)
  {
    return ( xmin*pow(1-ran1(), 1./(1-alpha) ) );
  }

  double powerlaw_continuous_interval(double a, double beta)
  {
    return pow( (1- (1-ran1())*(1-pow(a, 1-beta))), 1/(1-beta) );
  }

  int powerlaw_discrete_approximate(double alpha, double xmin)
  {
    double r = ran1();
    int k = (int)( (xmin-0.5)*pow(1-r, 1./(1-alpha) ) + 0.5 );
    if(k<0) {
      cout << 1-r << ',' << k << endl;
      k =  INTMAX;
    }
    return k;

  }

  int powerlaw_discrete(double alpha, int xmin)
  {
    if(!flag_setPvector)
      SetPvector(alpha,xmin);

    double r = ran1();
    int j = Findj(P, 1-r);
    return j;

  }

  int powerlaw_exponentialcutoff_discrete(double gamma, double kappa)
  {
    if(!flag_setPvector)
      SetPvector_SFEC(gamma, kappa);

    double r = ran1();

    int j = Findj(P, 1-r);
    return j;

  }

  int arbitrary_discrete(vector<double>& CDF)
  {
    double r = ran1();
    int j = Findj(CDF, r);
    return j+1;
  }

  int arbitrary_discrete(vector<long double>& CDF)
  {
    double r = ran1();
    int j = Findj(CDF, r);
    return j+1;
  }

  double lognormal(double sd)
  {
    double v1,v2,Rsq,z;
    if(outputReady)
      {
	outputReady = false;
	return output;
      }

    do
      {
	v1=(ran1()*2.0)-1.0;
	v2=(ran1()*2.0)-1.0;
	Rsq=v1*v1+v2*v2;
      } while (Rsq>=1.0);

    z=sqrt(-2.0*log(Rsq)/Rsq);
    output = exp(v1*z*sd);
    outputReady=true;
    return exp(v2*z*sd);
  }

  double bimodal(double R)
  {

    return (ran1()>0.5) ? (R) : (-R);
  }

  double rectangular(double R)
  {

    return ((ran1()*2-1)*R);
  }

  double parabolic(double R)
  {

    double x = ran1();

    int sign =  (x>0.5) ? (-1) : (+1);

    double R2 = R*R;
    double R3 = R2*R;

    double q     = 4*R3*(x-0.5);
    double p     = -3*R2;
    double Disc  = q*q/4. + p*p*p/27.;
    double r     = R3;
    double theta = atan(sqrt(-Disc)/(-q/2.));
    double ronethird = pow(r,1/3.);

    double h1 = sign * 2 * ronethird * cos(theta/3.);
    double h2 = sign * 2 * ronethird * cos((theta+2*PI)/3.);
    double h3 = sign * 2 * ronethird * cos((theta-2*PI)/3.);

    double y = 0;
    if(h1<R && h1>(-R))
      y = h1;
    else if(h2<R && h2>(-R))
      y = h2;
    else if(h3<R && h3>(-R))
      y = h3;
    else
      cout << "range is wrong!";

    return y;
  }

  double antiparabolic(double R)
  {

    double x = ran1();

    int sign =  (x>0.5) ? (+1) : (-1);

    double h = sign * pow(sign*(2*x-1),1./3.) * R;
    return h;
  }

  double DMI_lookuptable(double y, double R)

  {
    if(!flag_setFvector)
      SetFvector(y,R);

    double u = ran1();
    int j = Findj(F, u);
    double h = -R + j*dh;

    return h;
  }

  double DMI(double y, double R)
  {
    double u = ran1();
    int sign =  (u>0.5) ? (+1) : (-1);

    const double hacc = 1e-12;
    double htilde = rtbis(&fx, y, u, 0, 1, hacc);
    double h = sign*htilde*R;
    return h;
  }

  int exponential(double kappa)
  {
    double u = ran1();

    int k = -kappa * log(1-u);
    return k;
  }

 private:

  bool flag_setFvector;
  vector<double> F;
  double dh;

  double P_DMI(double h, double y, double R)
  {
    if(h<-R || h>R) return 0;
    else            return (y+1)/(2*y*R)*(1-pow(fabs(h)/R,y));
  }

  void SetFvector(double y, double R)
  {
    int N = 10000001;
    dh = R/(double)(N-1);
    F.resize(2*N-1,0);

    F[0]     = 0;
    F[2*N-2] = 1 - F[0];
    for(int i=1;i<N;i++)
      {
	double h1 = -R + i*dh;
	double h2 = h1 + dh;

	F[i]       = F[i-1] + 0.5* (P_DMI(h1,y,R)+P_DMI(h2,y,R))*dh;
	F[2*N-2-i] = 1- F[i];
      }

    flag_setFvector = true;
  }

  int Findj(vector<double>& xx, const double x)
  {
    int j,jl,jm,ju;
    int n=xx.size();

    jl=-1; ju=n;

    bool ascnd = (xx[n-1] >= xx[0]);
    while (ju-jl > 1)
      {
	jm=(ju+jl) >> 1;
	if ((x >= xx[jm]) == ascnd) jl=jm;
	else    	          ju=jm;
      }

    if (x == xx[0]) j=0;
    else if (x == xx[n-1]) j=n-2;
    else j=jl;

    return j;
  }

  int Findj(vector<long double>& xx, const double x)
  {
    int j,jl,jm,ju;
    int n=xx.size();

    jl=-1; ju=n;

    bool ascnd = (xx[n-1] >= xx[0]);
    while (ju-jl > 1)
      {
	jm=(ju+jl) >> 1;
	if ((x >= xx[jm]) == ascnd) jl=jm;
	else    	          ju=jm;
      }

    if (x == xx[0]) j=0;
    else if (x == xx[n-1]) j=n-2;
    else j=jl;

    return j;
  }

  double rtbis(FUNC func, const double y, const double u, const double x1, const double x2, const double xacc)
  {
    const int JMAX=40;
    int j;
    double dx,f,fmid,xmid,rtb;

    f=func(x1,y,u);
    fmid=func(x2,y,u);
    if (f*fmid >= 0.0)
      cout << "Root must be bracketed for bisection in rtbis\n";

    rtb = f < 0.0 ? (dx=x2-x1,x1) : (dx=x1-x2,x2);
    for (j=0;j<JMAX;j++) {
      fmid=func(xmid=rtb+(dx *= 0.5),y,u);
      if (fmid <= 0.0) rtb=xmid;
      if (fabs(dx) < xacc || fmid == 0.0) return rtb;
    }
    cout << "Too many bisections in rtbis\n";
    return 0.0;
  }

  void ranmarin(int ijkl)
  {
    int i,ii,j,jj,k,l,m ;
    double s,t ;

    int ij=ijkl/30082;
    int kl=ijkl-30082*ij;

    i=((ij/177)%177)+2 ;
    j=(ij%177)+2 ;
    k=((kl/169)%178)+1 ;
    l=kl%169 ;
    for (ii=0;ii<97;ii++) {
      s=0.0 ;
      t=0.5 ;
      for (jj=0;jj<24;jj++) {
	m=(((i*j)%179)*k)%179 ;
	i=j;
	j=k;
	k=m;
	l=(53*l+1)%169;
	if (((l*m)%64)>=32) s+=t;
	t*=0.5;
      }
      u[ii]=s;
    }
    c=362436.0/16777216.0;
    cd=7654321.0/16777216.0;
    cm=16777213.0/16777216.0;
    i97=96;
    j97=32;
  }

  bool flag_setPvector;
  vector<double> P;

  double Hzeta(double s, double q)
  {
    double sum = 0;
    double nterm = 0;
    int n;
    for(n=0; ;n++)
      {
	nterm = pow((n+q),-s);
	sum += nterm;

	if(fabs(nterm/sum) < TOLERANCE)
	  break;
      }
    cout << "nmax= " << n << endl;
    return sum;
  }

  void SetPvector(double alpha, int xmin)
  {

    double a = gsl_sf_hzeta(alpha, xmin);

    int N = 10000000+xmin;

    P.resize(N,2);

    P[xmin] = 1;
    for(int i=xmin;i<N-1;i++)
      {

	P[i+1] = P[i] - 1/(pow((double)i,alpha)*a);
      }

    flag_setPvector = true;
  }

double polylog(double s, double z)
{
  double sum = 0;
  double zn = z;
  int nmax = 1;
  for(int n=1; ; n++) {
    double term = zn/pow((double)n,s);
    sum += term;
    zn *= z;
    if(term < TOLERANCE) {
      nmax = n;
      break;
    }
  }

  return sum;
}

  void SetPvector_SFEC(double gamma, double kappa)
  {
    double enkappa= exp(-1/kappa);
    double C = polylog(gamma, enkappa);

    int N = 10000000;

    P.resize(N,2);

    P[1] = 1;
    double a = enkappa;
    for(int i=1;i<=N;i++) {
      P[i+1] = P[i] - pow((double)i,-gamma) * a/ C ;
      a *= enkappa;
    }

    flag_setPvector = true;
  }

};

#endif


/* fft/gsl_fft_complex.h
 *
 * Copyright (C) 1996, 1997, 1998, 1999, 2000 Brian Gough
 *
 * 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 2 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., 675 Mass Ave, Cambridge, MA 02139, USA.
 */

#ifndef __GSL_FFT_COMPLEX_H__
#define __GSL_FFT_COMPLEX_H__

#include <stddef.h>

#include <gsl/gsl_math.h>
#include <gsl/gsl_complex.h>
#include <gsl/gsl_fft.h>

#undef __BEGIN_DECLS
#undef __END_DECLS
#ifdef __cplusplus
# define __BEGIN_DECLS extern "C" {
# define __END_DECLS }
#else
# define __BEGIN_DECLS /* empty */
# define __END_DECLS /* empty */
#endif

__BEGIN_DECLS

/*  Power of 2 routines  */

int gsl_fft_complex_radix2_forward (gsl_complex_packed_array data,
                                    const size_t stride,
                                    const size_t n);

int gsl_fft_complex_radix2_backward (gsl_complex_packed_array data,
                                     const size_t stride,
                                     const size_t n);

int gsl_fft_complex_radix2_inverse (gsl_complex_packed_array data,
                                    const size_t stride,
                                    const size_t n);

int gsl_fft_complex_radix2_transform (gsl_complex_packed_array data,
                                      const size_t stride,
                                      const size_t n,
                                      const gsl_fft_direction sign);

int gsl_fft_complex_radix2_dif_forward (gsl_complex_packed_array data,
                                        const size_t stride,
                                        const size_t n);

int gsl_fft_complex_radix2_dif_backward (gsl_complex_packed_array data,
                                         const size_t stride,
                                         const size_t n);

int gsl_fft_complex_radix2_dif_inverse (gsl_complex_packed_array data,
                                        const size_t stride,
                                        const size_t n);

int gsl_fft_complex_radix2_dif_transform (gsl_complex_packed_array data,
                                          const size_t stride,
                                          const size_t n,
                                          const gsl_fft_direction sign);

int gsl_fft_complex_bitreverse_order (gsl_complex_packed_array data,
                                      size_t stride,
                                      size_t n,
                                      size_t n_bits);

/*  Mixed Radix general-N routines  */

typedef struct
  {
    size_t n;
    size_t nf;
    size_t factor[64];
    gsl_complex *twiddle[64];
    gsl_complex *trig;
  }
gsl_fft_complex_wavetable;

typedef struct
{
  size_t n;
  double *scratch;
}
gsl_fft_complex_workspace;

gsl_fft_complex_wavetable *gsl_fft_complex_wavetable_alloc (size_t n);

void gsl_fft_complex_wavetable_free (gsl_fft_complex_wavetable * wavetable);

gsl_fft_complex_workspace *gsl_fft_complex_workspace_alloc (size_t n);

void gsl_fft_complex_workspace_free (gsl_fft_complex_workspace * workspace);

int gsl_fft_complex_memcpy (gsl_fft_complex_wavetable * dest,
                            gsl_fft_complex_wavetable * src);

int gsl_fft_complex_forward (gsl_complex_packed_array data,
                             const size_t stride,
                             const size_t n,
                             const gsl_fft_complex_wavetable * wavetable,
                             gsl_fft_complex_workspace * work);

int gsl_fft_complex_backward (gsl_complex_packed_array data,
                              const size_t stride,
                              const size_t n,
                              const gsl_fft_complex_wavetable * wavetable,
                              gsl_fft_complex_workspace * work);

int gsl_fft_complex_inverse (gsl_complex_packed_array data,
                             const size_t stride,
                             const size_t n,
                             const gsl_fft_complex_wavetable * wavetable,
                             gsl_fft_complex_workspace * work);

int gsl_fft_complex_transform (gsl_complex_packed_array data,
                               const size_t stride, const size_t n,
                               const gsl_fft_complex_wavetable * wavetable,
                               gsl_fft_complex_workspace * work,
                               const gsl_fft_direction sign);

__END_DECLS

#endif /* __GSL_FFT_COMPLEX_H__ */