/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#ifndef SINGA_CORE_TENSOR_TENSOR_MATH_CPP_H_
#define SINGA_CORE_TENSOR_TENSOR_MATH_CPP_H_

#include "./tensor_math.h"
#include <cfloat>
#include "singa/core/common.h"
#include <math.h>

#ifdef USE_CBLAS
#include <cblas.h>
#endif

namespace singa {

template <>
void Abs<float, lang::Cpp>(const size_t num, const Block *in, Block *out,
                           Context *ctx) {
  float *outPtr = static_cast<float *>(out->mutable_data());
  const float *inPtr = static_cast<const float *>(in->data());
  for (size_t i = 0; i < num; i++) {
    outPtr[i] = fabs(inPtr[i]);
  }
}

template <>
void Add<float, lang::Cpp>(const size_t num, const Block *in, const float x,
                           Block *out, Context *ctx) {
  float *outPtr = static_cast<float *>(out->mutable_data());
  const float *inPtr = static_cast<const float *>(in->data());
  for (size_t i = 0; i < num; i++) {
    outPtr[i] = inPtr[i] + x;
  }
}

template <>
void Add<float, lang::Cpp>(const size_t num, const Block *in1, const Block *in2,
                           Block *out, Context *ctx) {
  // CHECK_EQ(ctx->stream, nullptr);
  float *outPtr = static_cast<float *>(out->mutable_data());
  const float *in1Ptr = static_cast<const float *>(in1->data());
  const float *in2Ptr = static_cast<const float *>(in2->data());
  for (size_t i = 0; i < num; i++) {
    outPtr[i] = in1Ptr[i] + in2Ptr[i];
  }
}

template <>
void Clamp<float, lang::Cpp>(const size_t num, const float low,
                             const float high, const Block *in, Block *out,
                             Context *ctx) {
  float *outPtr = static_cast<float *>(out->mutable_data());
  const float *inPtr = static_cast<const float *>(in->data());
  for (size_t i = 0; i < num; i++) {
    if (inPtr[i] > high) {
      outPtr[i] = high;
    } else if (inPtr[i] < low) {
      outPtr[i] = low;
    } else {
      outPtr[i] = inPtr[i];
    }
  }
}

template <>
void Div<float, lang::Cpp>(const size_t num, const Block *in1, const Block *in2,
                           Block *out, Context *ctx) {
  float *outPtr = static_cast<float *>(out->mutable_data());
  const float *in1Ptr = static_cast<const float *>(in1->data());
  const float *in2Ptr = static_cast<const float *>(in2->data());
  for (size_t i = 0; i < num; i++) {
    CHECK_NE(in2Ptr[i], 0.f);
    outPtr[i] = in1Ptr[i] / in2Ptr[i];
  }
}

template <>
void Div<float, lang::Cpp>(const size_t num, const float x, const Block *in,
                           Block *out, Context *ctx) {
  const float *inPtr = static_cast<const float *>(in->data());
  float *outPtr = static_cast<float *>(out->mutable_data());
  for (size_t i = 0; i < num; i++) {
    CHECK_NE(inPtr[i], 0.f);
    outPtr[i] = x / inPtr[i];
  }
}

template <>
void EltwiseMult<float, lang::Cpp>(const size_t num, const Block *in,
                                   const float x, Block *out, Context *ctx) {
  float *outPtr = static_cast<float *>(out->mutable_data());
  const float *inPtr = static_cast<const float *>(in->data());
  for (size_t i = 0; i < num; i++) {
    outPtr[i] = inPtr[i] * x;
  }
}

template <>
void EltwiseMult<float, lang::Cpp>(const size_t num, const Block *in1,
                                   const Block *in2, Block *out, Context *ctx) {
  float *outPtr = static_cast<float *>(out->mutable_data());
  const float *in1Ptr = static_cast<const float *>(in1->data());
  const float *in2Ptr = static_cast<const float *>(in2->data());
  for (size_t i = 0; i < num; i++) {
    outPtr[i] = in1Ptr[i] * in2Ptr[i];
  }
}
template <>
void Exp<float, lang::Cpp>(const size_t num, const Block *in, Block *out,
                           Context *ctx) {
  float *outPtr = static_cast<float *>(out->mutable_data());
  const float *inPtr = static_cast<const float *>(in->data());
  for (size_t i = 0; i < num; i++) {
    outPtr[i] = exp(inPtr[i]);
  }
}

template <>
void GE<float, lang::Cpp>(const size_t num, const Block *in, const float x,
                          Block *out, Context *ctx) {
  float *outPtr = static_cast<float *>(out->mutable_data());
  const float *inPtr = static_cast<const float *>(in->data());
  for (size_t i = 0; i < num; i++) {
    outPtr[i] = (inPtr[i] >= x) ? 1.f : 0.f;
  }
}

template <>
void GE<float, lang::Cpp>(const size_t num, const Block *in1, const Block *in2,
                          Block *out, Context *ctx) {
  float *outPtr = static_cast<float *>(out->mutable_data());
  const float *inPtr1 = static_cast<const float *>(in1->data());
  const float *inPtr2 = static_cast<const float *>(in2->data());
  for (size_t i = 0; i < num; i++) {
    outPtr[i] = (inPtr1[i] >= inPtr2[i]) ? 1.f : 0.f;
  }
}
template <>
void GT<float, lang::Cpp>(const size_t num, const Block *in, const float x,
                          Block *out, Context *ctx) {
  float *outPtr = static_cast<float *>(out->mutable_data());
  const float *inPtr = static_cast<const float *>(in->data());
  for (size_t i = 0; i < num; i++) {
    outPtr[i] = (inPtr[i] > x) ? 1.f : 0.f;
  }
}
template <>
void GT<float, lang::Cpp>(const size_t num, const Block *in1, const Block *in2,
                          Block *out, Context *ctx) {
  float *outPtr = static_cast<float *>(out->mutable_data());
  const float *inPtr1 = static_cast<const float *>(in1->data());
  const float *inPtr2 = static_cast<const float *>(in2->data());
  for (size_t i = 0; i < num; i++) {
    outPtr[i] = (inPtr1[i] > inPtr2[i]) ? 1.f : 0.f;
  }
}

template <>
void LE<float, lang::Cpp>(const size_t num, const Block *in, const float x,
                          Block *out, Context *ctx) {
  float *outPtr = static_cast<float *>(out->mutable_data());
  const float *inPtr = static_cast<const float *>(in->data());
  for (size_t i = 0; i < num; i++) {
    outPtr[i] = (inPtr[i] <= x) ? 1.f : 0.f;
  }
}
template <>
void LE<float, lang::Cpp>(const size_t num, const Block *in1, const Block *in2,
                          Block *out, Context *ctx) {
  float *outPtr = static_cast<float *>(out->mutable_data());
  const float *inPtr1 = static_cast<const float *>(in1->data());
  const float *inPtr2 = static_cast<const float *>(in2->data());
  for (size_t i = 0; i < num; i++) {
    outPtr[i] = (inPtr1[i] <= inPtr2[i]) ? 1.f : 0.f;
  }
}
template <>
void Log<float, lang::Cpp>(const size_t num, const Block *in, Block *out,
                           Context *ctx) {
  float *outPtr = static_cast<float *>(out->mutable_data());
  const float *inPtr = static_cast<const float *>(in->data());
  for (size_t i = 0; i < num; i++) {
    CHECK_GT(inPtr[i], 0.f);
    outPtr[i] = log(inPtr[i]);
  }
}
template <>
void LT<float, lang::Cpp>(const size_t num, const Block *in, const float x,
                          Block *out, Context *ctx) {
  float *outPtr = static_cast<float *>(out->mutable_data());
  const float *inPtr = static_cast<const float *>(in->data());
  for (size_t i = 0; i < num; i++) {
    outPtr[i] = (inPtr[i] < x) ? 1.f : 0.f;
  }
}
template <>
void LT<float, lang::Cpp>(const size_t num, const Block *in1, const Block *in2,
                          Block *out, Context *ctx) {
  float *outPtr = static_cast<float *>(out->mutable_data());
  const float *inPtr1 = static_cast<const float *>(in1->data());
  const float *inPtr2 = static_cast<const float *>(in2->data());
  for (size_t i = 0; i < num; i++) {
    outPtr[i] = (inPtr1[i] < inPtr2[i]) ? 1.f : 0.f;
  }
}

template <>
void Pow<float, lang::Cpp>(const size_t num, const Block *in, const float x,
                           Block *out, Context *ctx) {
  float *outPtr = static_cast<float *>(out->mutable_data());
  const float *inPtr = static_cast<const float *>(in->data());
  for (size_t i = 0; i < num; i++) {
    outPtr[i] = pow(inPtr[i], x);
  }
}

template <>
void Pow<float, lang::Cpp>(const size_t num, const Block *in1, const Block *in2,
                           Block *out, Context *ctx) {
  float *outPtr = static_cast<float *>(out->mutable_data());
  const float *in1Ptr = static_cast<const float *>(in1->data());
  const float *in2Ptr = static_cast<const float *>(in2->data());
  for (size_t i = 0; i < num; i++) {
    outPtr[i] = pow(in1Ptr[i], in2Ptr[i]);
  }
}
template <>
void ReLU<float, lang::Cpp>(const size_t num, const Block *in, Block *out,
                            Context *ctx) {
  float *outPtr = static_cast<float *>(out->mutable_data());
  const float *inPtr = static_cast<const float *>(in->data());
  for (size_t i = 0; i < num; i++) {
    outPtr[i] = (inPtr[i] >= 0.f) ? inPtr[i] : 0.f;
  }
}
template <>
void Set<float, lang::Cpp>(const size_t num, const float x, Block *out,
                           Context *ctx) {
  float *outPtr = static_cast<float *>(out->mutable_data());
  for (size_t i = 0; i < num; i++) outPtr[i] = x;
}
template <>
void Set<int, lang::Cpp>(const size_t num, const int x, Block *out,
                           Context *ctx) {
  int *outPtr = static_cast<int *>(out->mutable_data());
  for (size_t i = 0; i < num; i++) outPtr[i] = x;
}

template <>
void Sigmoid<float, lang::Cpp>(const size_t num, const Block *in, Block *out,
                               Context *ctx) {
  float *outPtr = static_cast<float *>(out->mutable_data());
  const float *inPtr = static_cast<const float *>(in->data());
  for (size_t i = 0; i < num; i++) {
    outPtr[i] = 1.f / (1.f + exp(-inPtr[i]));
  }
}

template <>
void Sign<float, lang::Cpp>(const size_t num, const Block *in, Block *out,
                            Context *ctx) {
  float *outPtr = static_cast<float *>(out->mutable_data());
  const float *inPtr = static_cast<const float *>(in->data());
  for (size_t i = 0; i < num; i++) {
    outPtr[i] = (inPtr[i] > 0) - (inPtr[i] < 0);
  }
}

template <>
void Sqrt<float, lang::Cpp>(const size_t num, const Block *in, Block *out,
                            Context *ctx) {
  float *outPtr = static_cast<float *>(out->mutable_data());
  const float *inPtr = static_cast<const float *>(in->data());
  for (size_t i = 0; i < num; i++) {
    CHECK_GE(inPtr[i], 0.f);
    outPtr[i] = sqrt(inPtr[i]);
  }
}
/*
template <>
void Square<float, lang::Cpp>(const size_t num, const Block *in, Block *out,
                              Context *ctx) {
  float *outPtr = static_cast<float *>(out->mutable_data());
  const float *inPtr = static_cast<const float *>(in->data());
  for (size_t i = 0; i < num; i++) {
    outPtr[i] = inPtr[i] * inPtr[i];
  }
}
*/

template <>
void Sub<float, lang::Cpp>(const size_t num, const Block *in1, const Block *in2,
                           Block *out, Context *ctx) {
  // CHECK_EQ(ctx->stream, nullptr);
  float *outPtr = static_cast<float *>(out->mutable_data());
  const float *in1Ptr = static_cast<const float *>(in1->data());
  const float *in2Ptr = static_cast<const float *>(in2->data());
  for (size_t i = 0; i < num; i++) {
    outPtr[i] = in1Ptr[i] - in2Ptr[i];
  }
}

// sum all elements of input into out
// TODO(wangwei) optimize using omp
template <>
void Sum<float, lang::Cpp>(const size_t num, const Block *in, float *out,
                           Context *ctx) {
  float s = 0.f;
  const float *inPtr = static_cast<const float *>(in->data());
  for (size_t i = 0; i < num; i++) {
    s += inPtr[i];
  }
  *out = s;
}

template <>
void Tanh<float, lang::Cpp>(const size_t num, const Block *in, Block *out,
                            Context *ctx) {
  float *outPtr = static_cast<float *>(out->mutable_data());
  const float *inPtr = static_cast<const float *>(in->data());
  for (size_t i = 0; i < num; i++) {
    outPtr[i] = tanh(inPtr[i]);
  }
}

// ===============Random operations==========================================
template <>
void Bernoulli<float, lang::Cpp>(const size_t num, const float p, Block *out,
                                 Context *ctx) {
  std::bernoulli_distribution distribution(p);
  float *outPtr = static_cast<float *>(out->mutable_data());
  for (size_t i = 0; i < num; i++) {
    outPtr[i] = distribution(ctx->random_generator) ? 1.0f : 0.0f;
  }
}

template <>
void Gaussian<float, lang::Cpp>(const size_t num, const float mean,
                                const float std, Block *out, Context *ctx) {
  std::normal_distribution<float> distribution(mean, std);
  float *outPtr = static_cast<float *>(out->mutable_data());
  for (size_t i = 0; i < num; i++) {
    outPtr[i] = static_cast<float>(distribution(ctx->random_generator));
  }
}
template <>
void Uniform<float, lang::Cpp>(const size_t num, const float low,
                               const float high, Block *out, Context *ctx) {
  std::uniform_real_distribution<float> distribution(low, high);
  float *outPtr = static_cast<float *>(out->mutable_data());
  for (size_t i = 0; i < num; i++) {
    outPtr[i] = static_cast<float>(distribution(ctx->random_generator));
  }
}

// ====================Blas operations======================================

template <>
void DGMM<float, lang::Cpp>(const bool side_right, const size_t nrow,
                            const size_t ncol, const Block *M, const Block *v,
                            Block *out, Context *ctx) {
  const float *MPtr = static_cast<const float *>(M->data());
  const float *vPtr = static_cast<const float *>(v->data());
  float *outPtr = static_cast<float *>(out->mutable_data());
  if (side_right) {
    for (size_t r = 0; r < nrow; r++) {
      size_t offset = r * ncol;
      for (size_t c = 0; c < ncol; c++) {
        outPtr[offset + c] = MPtr[offset + c] * vPtr[c];
      }
    }
  } else {
    for (size_t r = 0; r < nrow; r++) {
      size_t offset = r * ncol;
      for (size_t c = 0; c < ncol; c++) {
        outPtr[offset + c] = MPtr[offset + c] * vPtr[r];
      }
    }
  }
}

#ifdef USE_CBLAS
template <>
void Amax<float, lang::Cpp>(const size_t num, const Block *in, size_t *out,
                            Context *ctx) {
  const float *inPtr = static_cast<const float *>(in->data());
  *out = cblas_isamax(num, inPtr, 1);
}

template <>
void Asum<float, lang::Cpp>(const size_t num, const Block *in, float *out,
                            Context *ctx) {
  const float *inPtr = static_cast<const float *>(in->data());
  *out = cblas_sasum(num, inPtr, 1);
}

template <>
void Axpy<float, lang::Cpp>(const size_t num, const float alpha,
                            const Block *in, Block *out, Context *ctx) {
  const float *inPtr = static_cast<const float *>(in->data());
  float *outPtr = static_cast<float *>(out->mutable_data());
  cblas_saxpy(num, alpha, inPtr, 1, outPtr, 1);
}

template <>
void Dot<float, lang::Cpp>(const size_t num, const Block *in1, const Block *in2,
                           float *out, Context *ctx) {
  const float *in1Ptr = static_cast<const float *>(in1->data());
  const float *in2Ptr = static_cast<const float *>(in2->data());
  *out = cblas_sdot(num, in1Ptr, 1, in2Ptr, 1);
}
template <>
void Scale<float, lang::Cpp>(const size_t num, const float x, Block *out,
                             Context *ctx) {
  float *outPtr = static_cast<float *>(out->mutable_data());
  cblas_sscal(num, x, outPtr, 1);
}
template <>
void Nrm2<float, lang::Cpp>(const size_t num, const Block *in, float *out,
                            Context *ctx) {
  const float *inPtr = static_cast<const float *>(in->data());
  *out = cblas_snrm2(num, inPtr, 1);
}

template <>
void GEMV<float, lang::Cpp>(bool trans, const size_t m, const size_t n,
                            const float alpha, const Block *A, const Block *v,
                            const float beta, Block *out, Context *ctx) {
  const float *APtr = static_cast<const float *>(A->data());
  const float *vPtr = static_cast<const float *>(v->data());
  float *outPtr = static_cast<float *>(out->mutable_data());
  if (!trans) {
    cblas_sgemv(CblasRowMajor, CblasNoTrans, m, n, alpha, APtr, n, vPtr, 1,
                beta, outPtr, 1);
  } else {
    cblas_sgemv(CblasRowMajor, CblasTrans, n, m, alpha, APtr, m, vPtr, 1, beta,
                outPtr, 1);
  }
}

template <>
void GEMM<float, lang::Cpp>(const bool transA, const bool transB,
                            const size_t nrowA, const size_t ncolB,
                            const size_t ncolA, const float alpha,
                            const Block *A, const Block *B, const float beta,
                            Block *C, Context *ctx) {
  auto transa = transA ? CblasTrans : CblasNoTrans;
  auto transb = transB ? CblasTrans : CblasNoTrans;
  auto lda = transA ? nrowA : ncolA;
  auto ldb = transB ? ncolA : ncolB;
  auto ldc = ncolB;
  const float *APtr = static_cast<const float *>(A->data());
  const float *BPtr = static_cast<const float *>(B->data());
  float *CPtr = static_cast<float *>(C->mutable_data());
  cblas_sgemm(CblasRowMajor, transa, transb, nrowA, ncolB, ncolA, alpha, APtr,
	  lda, BPtr, ldb, beta, CPtr, ldc);
}

#else

template <>
void Amax<float, lang::Cpp>(const size_t num, const Block *in, size_t *out,
                            Context *ctx) {
  size_t maxPos = 0;
  float maxVal = 0;
  const float *inPtr = static_cast<const float *>(in->data());
  for (size_t i = 0; i < num; i++) {
    if (i == 0) {
      maxVal = inPtr[i];
    } else if (inPtr[i] > maxVal) {
      maxVal = inPtr[i];
      maxPos = i;
    }
  }
  *out = maxPos;
}
template <>
void Amin<float, lang::Cpp>(const size_t num, const Block *in, size_t *out,
                            Context *ctx) {
  size_t minPos = 0;
  float minVal = 0;
  const float *inPtr = static_cast<const float *>(in->data());
  for (size_t i = 0; i < num; i++) {
    if (i == 0) {
      minVal = inPtr[i];
    } else if (inPtr[i] > minVal) {
      minVal = inPtr[i];
      minPos = i;
    }
  }
  *out = minPos;
}

template <>
void Asum<float, lang::Cpp>(const size_t num, const Block *in, float *out,
                            Context *ctx) {
  float sum = 0;
  const float *inPtr = static_cast<const float *>(in->data());
  for (size_t i = 0; i < num; i++) {
    sum += fabs(inPtr[i]);
  }
}

template <>
void Axpy<float, lang::Cpp>(const size_t num, const float alpha,
                            const Block *in, Block *out, Context *ctx) {
  float *outPtr = static_cast<float *>(out->mutable_data());
  const float *inPtr = static_cast<const float *>(in->data());
  for (size_t i = 0; i < num; i++) {
    outPtr[i] += alpha * inPtr[i];
  }
}

template <>
void Scale<float, lang::Cpp>(const size_t num, const float x, Block *out,
                             Context *ctx) {
  float *outPtr = static_cast<float *>(out->mutable_data());
  for (size_t i = 0; i < num; i++) {
    outPtr[i] *= x;
  }
}

template <>
void Dot<float, lang::Cpp>(const size_t num, const Block *in1, const Block *in2,
                           float *out, Context *ctx) {
  float sum = 0;
  const float *in1Ptr = static_cast<const float *>(in1->data());
  const float *in2Ptr = static_cast<const float *>(in2->data());
  for (size_t i = 0; i < num; i++) {
    sum += in1Ptr[i] * in2Ptr[i];
  }
}

template <>
void GEMV<float, lang::Cpp>(bool trans, const size_t m, const size_t n,
                            const float alpha, const Block *A, const Block *v,
                            const float beta, Block *out, Context *ctx) {
  float *outPtr = static_cast<float *>(out->mutable_data());
  const float *APtr = static_cast<const float *>(A->data());
  const float *vPtr = static_cast<const float *>(v->data());
  for (size_t r = 0; r < m; r++) {
    float sum = 0;
    for (size_t c = 0; c < n; c++) {
      size_t idx = trans ? c * m + r : r * n + c;
      sum += APtr[idx] * vPtr[c];
    }
    outPtr[r] = alpha * sum + beta * outPtr[r];
  }
}

#endif  // USE_CBLAS
template <>
void ComputeCrossEntropy<float, lang::Cpp>(bool int_target,
                                           const size_t batchsize,
                                           const size_t dim, const Block *p,
                                           const Block *t, Block *loss,
                                           Context *ctx) {
  const float *pPtr = static_cast<const float *>(p->data());
  const int *tPtr = static_cast<const int *>(t->data());
  float *lossPtr = static_cast<float *>(loss->mutable_data());
  if (int_target) {
    for (size_t i = 0; i < batchsize; i++) {
      int truth_idx = tPtr[i];
      CHECK_GE(truth_idx, 0);
      float prob_of_truth = pPtr[i * dim + truth_idx];
      lossPtr[i] = -std::log((std::max)(prob_of_truth, FLT_MIN));
    }
  } else {
    for (size_t i = 0;i < batchsize; i++) {
      float sum = 0.f;
      for (size_t j = 0; j < dim; j++) {
        sum += tPtr[i * dim + j];
      }
      float loss = 0.f;
      for (size_t j = 0, offset = i * dim; j < dim; j++, offset++) {
        loss -= tPtr[offset] / sum * std::log((std::max)(pPtr[offset], FLT_MIN));
      }
      lossPtr[i] = loss;
    }
  }
}

template <>
void SoftmaxCrossEntropyBwd<float, lang::Cpp>(bool int_target,
                                              const size_t batchsize,
                                              const size_t dim, const Block *p,
                                              const Block *t, Block *grad,
                                              Context *ctx) {
  CHECK_EQ(p, grad) << "Use the same pointer to optimize performance";
  // const float* pPtr = static_cast<const float*>(p->data());
  const int *tPtr = static_cast<const int *>(t->data());
  float *gradPtr = static_cast<float *>(grad->mutable_data());

  if (int_target) {
    for (size_t i = 0; i < batchsize; i++) {
      int truth_idx = static_cast<int>(tPtr[i]);
      CHECK_GE(truth_idx, 0);
      gradPtr[i * dim + truth_idx] -= 1.0;
    }
  } else {
    for (size_t i = 0; i < batchsize; i++) {
      float sum = 0.f;
      for (size_t j = 0; j < dim; j++) {
        sum += tPtr[i * dim + j];
      }
      for (size_t j = 0, offset = i * dim; j < dim; j++, offset++) {
        gradPtr[offset] -= tPtr[offset] / sum;
      }
    }
  }
}

template <>
void RowMax<float, lang::Cpp>(const size_t nrow, const size_t ncol,
                              const Block *in, Block *out, Context *ctx) {
  const float *inPtr = static_cast<const float *>(in->data());
  float *outPtr = static_cast<float *>(out->mutable_data());
  for (size_t r = 0; r < nrow; r++) {
    int offset = (int)(r * ncol);
    float maxval = inPtr[offset];
    for (size_t c = 1; c < ncol; c++)
      maxval = (std::max)(maxval, inPtr[offset + c]);
    outPtr[r] = maxval;
  }
}

// =========Matrix operations ================================================
/*
template <>
void AddCol<float, lang::Cpp>(const size_t nrow, const size_t ncol,
                              const Block *A, const Block *v, Block *out,
                              Context *ctx) {
  float *outPtr = static_cast<float *>(out->mutable_data());
  const float *APtr = static_cast<const float *>(A->data());
  const float *vPtr = static_cast<const float *>(v->data());
  for (size_t r = 0; r < nrow; r++) {
    size_t offset = r * ncol;
    for (size_t c = 0; c < ncol; c++) {
      outPtr[offset + c] = APtr[offset + c] + vPtr[r];
    }
  }
}

template <>
void AddRow<float, lang::Cpp>(const size_t nrow, const size_t ncol,
                              const Block *A, const Block *v, Block *out,
                              Context *ctx) {
  float *outPtr = static_cast<float *>(out->mutable_data());
  const float *APtr = static_cast<const float *>(A->data());
  const float *vPtr = static_cast<const float *>(v->data());
  for (size_t r = 0; r < nrow; r++) {
    size_t offset = r * ncol;
    for (size_t c = 0; c < ncol; c++) {
      outPtr[offset + c] = APtr[offset + c] + vPtr[c];
    }
  }
}
template <>
void Outer<float, lang::Cpp>(const size_t m, const size_t n, const Block *in1,
                             const Block *in2, Block *out, Context *ctx) {
  float *outPtr = static_cast<float *>(out->mutable_data());
  const float *in1Ptr = static_cast<const float *>(in1->data());
  const float *in2Ptr = static_cast<const float *>(in2->data());
  for (size_t r = 0; r < m; r++) {
    size_t offset = r * n;
    for (size_t c = 0; c < n; c++) {
      outPtr[offset + c] = in1Ptr[r] * in2Ptr[c];
    }
  }
}
template <>
void Softmax<float, lang::Cpp>(const size_t nrow, const size_t ncol,
                               const Block *in, Block *out, Context *ctx) {
  float *outPtr = static_cast<float *>(out->mutable_data());
  const float *inPtr = static_cast<const float *>(in->data());
  float *bPtr = new float[ncol];
  for (size_t r = 0; r < nrow; r++) {
    size_t offset = r * ncol;
    float denom = 0.f;
    for (size_t c = 0; c < ncol; c++) {
      bPtr[c] = exp(inPtr[offset + c]);
      denom += bPtr[c];
    }
    for (size_t c = 0; c < ncol; c++) {
      size_t idx = offset + c;
      outPtr[idx] = bPtr[c] / denom;
    }
  }
  delete bPtr;
}

template <>
void SumColumns<float, lang::Cpp>(const size_t nrow, const size_t ncol,
                                  const Block *in, Block *out, Context *ctx) {
  float *outPtr = static_cast<float *>(out->mutable_data());
  const float *inPtr = static_cast<const float *>(in->data());
  for (size_t c = 0; c < ncol; c++) {
    outPtr[c] = 0.f;
  }
  for (size_t r = 0; r < nrow; r++) {
    size_t offset = r * ncol;
    for (size_t c = 0; c < ncol; c++) {
      outPtr[c] += inPtr[offset + c];
    }
  }
}

template <>
void SumRows<float, lang::Cpp>(const size_t nrow, const size_t ncol,
                               const Block *in, Block *out, Context *ctx) {
  float *outPtr = static_cast<float *>(out->mutable_data());
  const float *inPtr = static_cast<const float *>(in->data());
  for (size_t r = 0; r < nrow; r++) {
    size_t offset = r * ncol;
    outPtr[r] = 0.f;
    for (size_t c = 0; c < ncol; c++) {
      outPtr[r] += inPtr[offset + c];
    }
  }
}
*/
}  // namespace singa

#endif  // SINGA_CORE_TENSOR_TENSOR_MATH_CPP_H_


/*
 * author: Achim Gaedke
 * created: January 2001
 * file: pygsl/src/histogrammodule.c
 * $Id: matrixmodule.c,v 1.2 2003/07/25 17:15:16 schnizer Exp $
 */

#include <Python.h>
#include <gsl/gsl_errno.h>
#include <gsl/gsl_matrix.h>
#include <pygsl/error_helpers.h>
/*
 * matrix
 */

/* my typedef */

staticforward PyTypeObject matrix_matrixType;

typedef struct {
    PyObject_HEAD
    gsl_matrix* m;
} matrix_matrixObject;

static PyObject* matrix_matrix_get(PyObject *self,
				   PyObject *args
				   )
{
  gsl_matrix* matrix;
  long int i;
  long int j;
  matrix=(gsl_matrix*)((matrix_matrixObject*)self)->m;
  if (matrix==NULL) {
    PyErr_SetString(PyExc_RuntimeError, "matrix.get got a NULL pointer");
    return NULL;
  }

  /* get index */
  if (0==PyArg_ParseTuple(args,"ll",&i,&j))
      return NULL;
  if (i<0 || i>=matrix->size1) {
    gsl_error ("1st index lies outside valid range of 0 .. size1 - 1",
	       __FILE__,
	       __LINE__,
	       GSL_EDOM );
    return NULL;
  }
  if (j<0 || j>=matrix->size2) {
    gsl_error ("2nd index lies outside valid range of 0 .. size2 - 1",
	       __FILE__,
	       __LINE__,
	       GSL_EDOM );
    return NULL;
  }
  return PyFloat_FromDouble(gsl_matrix_get(matrix,i,j));
}

static PyObject* matrix_matrix_set(PyObject *self,
				   PyObject *args
				   )
{
  gsl_matrix* matrix;
  long int i;
  long int j;
  double x;
  matrix=(gsl_matrix*)((matrix_matrixObject*)self)->m;
  if (matrix==NULL) {
    PyErr_SetString(PyExc_RuntimeError, "matrix.get got a NULL pointer");
    return NULL;
  }

  /* get index */
  if (0==PyArg_ParseTuple(args,"lld",&i,&j,&x))
      return NULL;
  if (i<0 || i>=matrix->size1) {
    gsl_error ("1st index lies outside valid range of 0 .. size1 - 1",
	       __FILE__,
	       __LINE__,
	       GSL_EDOM );
    return NULL;
  }
  if (j<0 || j>=matrix->size2) {
    gsl_error ("2nd index lies outside valid range of 0 .. size2 - 1",
	       __FILE__,
	       __LINE__,
	       GSL_EDOM );
    return NULL;
  }
  gsl_matrix_set(matrix,i,j,x);
  Py_INCREF(Py_None);
  return Py_None;
}

static PyObject* matrix_matrix_set_all(PyObject *self,
				       PyObject* args)
{
  gsl_matrix* matrix;
  double x;
  matrix=(gsl_matrix*)((matrix_matrixObject*)self)->m;
  if (0==PyArg_ParseTuple(args,"d",&x))
      return NULL;
  if (matrix==NULL) {
    PyErr_SetString(PyExc_RuntimeError, "matrix.size got a NULL pointer");
    return NULL;
  }
  gsl_matrix_set_all(matrix,x);
  Py_INCREF(Py_None);
  return Py_None;
}

static PyObject* matrix_matrix_set_identity(PyObject *self)
{
  gsl_matrix* matrix;
  matrix=(gsl_matrix*)((matrix_matrixObject*)self)->m;
  if (matrix==NULL) {
    PyErr_SetString(PyExc_RuntimeError, "matrix.size got a NULL pointer");
    return NULL;
  }
  gsl_matrix_set_identity(matrix);
  Py_INCREF(Py_None);
  return Py_None;
}

static PyObject* matrix_matrix_set_zero(PyObject *self)
{
  gsl_matrix* matrix;
  matrix=(gsl_matrix*)((matrix_matrixObject*)self)->m;
  if (matrix==NULL) {
    PyErr_SetString(PyExc_RuntimeError, "matrix.size got a NULL pointer");
    return NULL;
  }
  gsl_matrix_set_zero(matrix);
  Py_INCREF(Py_None);
  return Py_None;
}

static PyObject* matrix_matrix_transpose(PyObject *self)
{
  gsl_matrix* matrix;
  matrix=(gsl_matrix*)((matrix_matrixObject*)self)->m;
  if (matrix==NULL) {
    PyErr_SetString(PyExc_RuntimeError, "matrix.size got a NULL pointer");
    return NULL;
  }
  gsl_matrix_transpose(matrix);
  Py_INCREF(Py_None);
  return Py_None;
}

static PyObject* matrix_matrix_size(PyObject *self)
{
  gsl_matrix* matrix;

  matrix=(gsl_matrix*)((matrix_matrixObject*)self)->m;
  if (matrix==NULL) {
    PyErr_SetString(PyExc_RuntimeError, "matrix.size got a NULL pointer");
    return NULL;
  }

  return Py_BuildValue("(ll)",matrix->size1,matrix->size2);
}

static PyObject* matrix_matrix_add(PyObject *self,
				   PyObject *arg) {
  gsl_matrix* matrix=(gsl_matrix*)((matrix_matrixObject*)self)->m;
  gsl_matrix* matrix2;
  int result;

  if (matrix==NULL) {
    PyErr_SetString(PyExc_RuntimeError, "matrix.add got a NULL pointer");
    return NULL;
  }
  Py_INCREF(arg);
  /* get histogram from argument */
  if (arg==NULL || !PyObject_TypeCheck(arg,&matrix_matrixType)) {
    PyErr_SetString(PyExc_TypeError, "matrix.add requires pygsl.matrix.matrix");
    Py_DECREF(arg);
    return NULL;
  }
  matrix2=(gsl_matrix*)((matrix_matrixObject*)arg)->m;
  if (matrix2==NULL) {
    PyErr_SetString(PyExc_RuntimeError, "matrix in argument has a NULL pointer");
    return NULL;
  }
  result=gsl_matrix_add(matrix,matrix2);
  Py_DECREF(arg);
  if (result!=GSL_SUCCESS) {
    return NULL;
  }
  Py_INCREF(Py_None);
  return Py_None;
}

static PyMethodDef matrix_matrix_methods[] = {
  /*  {"alloc",(PyCFunction)histogram_histogram_alloc,METH_VARARGS,"allocate necessary space"}, */
  {"get",(PyCFunction)matrix_matrix_get,METH_VARARGS,"gets the (i,j) element"},
  {"size",(PyCFunction)matrix_matrix_size,METH_NOARGS,"returns the size as tuple"},
  {"add",(PyCFunction)matrix_matrix_add,METH_O,"adds the given matrix"},
  {"set",(PyCFunction)matrix_matrix_set,METH_VARARGS,"sets the (i,j) element"},
  {"set_zero",(PyCFunction)matrix_matrix_set_zero,METH_NOARGS,"sets all to 0"},
  {"set_identity",(PyCFunction)matrix_matrix_set_identity,METH_NOARGS,"sets diagonal to 1 and others to 0"},
  {"set_all",(PyCFunction)matrix_matrix_set_identity,METH_O,"sets all entries to x"},
  {"transpose",(PyCFunction)matrix_matrix_transpose,METH_NOARGS,"transposes the matrix in place"},
  {NULL, NULL, 0, NULL}           /* sentinel */
};

static int matrix_matrix_init(PyObject *self,
			      PyObject *args,
			      PyObject *kwds
			      )
{

  PyObject* orig_matrix;
  long int n1;
  long int n2;

  static char *kwlist1[] = {"matrix",NULL};
  static char *kwlist2[] = {"size1","size2",NULL};

  /* initialise histogram pointer */
  ((matrix_matrixObject*)self)->m=NULL;

  /* test two call alternatives: */

  if (PyArg_ParseTupleAndKeywords(args, kwds,"O!",kwlist1,&matrix_matrixType,&orig_matrix)) {
    gsl_matrix* matrix;
    gsl_matrix* matrix2=(gsl_matrix*)((matrix_matrixObject*)orig_matrix)->m;
    if (matrix2==NULL) {
      PyErr_SetString(PyExc_RuntimeError, "matrix in argument has a NULL pointer");
      return -1;
    }
    matrix=gsl_matrix_calloc(matrix2->size1,matrix2->size2);
    if (matrix==NULL)
      return -1;
    gsl_matrix_memcpy(matrix,matrix2);
    ((matrix_matrixObject*)self)->m=matrix;
    return 0;
  }
  else {
    /* reset exception */
    PyErr_Clear();
  }
  if (PyArg_ParseTupleAndKeywords(args, kwds,"ll",kwlist2,&n1,&n2)) {
    gsl_matrix* matrix;
    if (n1<=0 || n2<=0) {
      gsl_error ("matrix length must be positive",
		 __FILE__,
		 __LINE__,
		 GSL_EDOM );
      return -1;
    }
    matrix=gsl_matrix_calloc(n1,n2);
    if (matrix==NULL)
      return -1;
    ((matrix_matrixObject*)self)->m=matrix;
    return 0;

  }
  else {
    /* set own error message */
    PyErr_Clear();
    PyErr_SetString(PyExc_TypeError, "matrix.init requires pygsl.matrix.matirx object or two long int argument");
    return -1;
  }
}

static PyObject*
matrix_matrix_getattr(PyObject* obj, char *name)
{
  return Py_FindMethod(matrix_matrix_methods, obj, name);
}

static void
matrix_matrix_dealloc(PyObject* self)
{
  gsl_matrix* matrix;
  matrix=(gsl_matrix*)((matrix_matrixObject*)self)->m;
  if (matrix!=NULL) gsl_matrix_free(matrix);
  self->ob_type->tp_free(self);
}

static
PyTypeObject matrix_matrixType = {
	PyObject_HEAD_INIT(NULL)
	0,
	"pygsl.matrix.matrix",
	sizeof(matrix_matrixObject),
	0,
	(destructor)matrix_matrix_dealloc,      /* tp_dealloc */
	0,                                      /* tp_print */
	matrix_matrix_getattr,                  /* tp_getattr */
	0,					/* tp_setattr */
	0,			                /* tp_compare */
	0,                  			/* tp_repr */
	0,					/* tp_as_number */
	0,	                		/* tp_as_sequence */
	0,    /* &matrix_marix_as_mapping,*/	/* tp_as_mapping */
	0,				        /* tp_hash */
	0,					/* tp_call */
	0,					/* tp_str */
	0,		                        /* tp_getattro */
	0,					/* tp_setattro */
	0,					/* tp_as_buffer */
	Py_TPFLAGS_DEFAULT,		        /* tp_flags */
        0,				        /* tp_doc */
	0, 		                        /* tp_traverse */
	0,			                /* tp_clear */
	0,              			/* tp_richcompare */
	0,					/* tp_weaklistoffset */
	0,			                /* tp_iter */
	0,					/* tp_iternext */
	0,				        /* tp_methods */
	0,					/* tp_members */
	0,					/* tp_getset */
	0,					/* tp_base */
	0,					/* tp_dict */
	0,					/* tp_descr_get */
	0,					/* tp_descr_set */
	0,					/* tp_dictoffset */
	(initproc)matrix_matrix_init,	        /* tp_init */
	NULL,              			/* tp_alloc */
	NULL,                			/* tp_new */
	NULL         			        /* tp_free */
};

static PyMethodDef matrixMethods[] = {
  {NULL, NULL, 0, NULL}        /* Sentinel */
};

void
initmatrix(void)
{
  PyObject* m;
  m=Py_InitModule("matrix", matrixMethods);
  init_pygsl();
  /* init histogram type */
  matrix_matrixType.ob_type  = &PyType_Type;
  matrix_matrixType.tp_alloc = PyType_GenericAlloc;
  matrix_matrixType.tp_new   = PyType_GenericNew;
  matrix_matrixType.tp_free  = _PyObject_Del;

  /* install histogram type */
  /* important! must increment histogram type reference counter */
  Py_INCREF((PyObject*)&matrix_matrixType);
  PyModule_AddObject(m,"matrix", (PyObject*)&matrix_matrixType);
}


#include <gsl/gsl_rng.h>
#include <easy_rng.h>
#include <time.h>

#ifndef NRUNS
#define NRUNS 10000000
#endif

void run_test(const gsl_rng_type *gtype, const easy_rng_type *etype) {
	gsl_rng *grng = gsl_rng_alloc(gtype);
	easy_rng *erng = easy_rng_alloc(etype);

	clock_t gstart, gend;
	clock_t estart, eend;
	int i;

	gstart = clock();

	for (i = 0 ; i < NRUNS ; i++)
		gsl_rng_get(grng);

	gend = clock();

	estart = clock();

	for (i = 0 ; i < NRUNS ; i++)
		easy_rng_get(erng);

	eend = clock();

	gsl_rng_free(grng);
	easy_rng_free(erng);

	fprintf(stdout, "Comparing easyRNG's %s with GSL's %s\n", etype->name, gtype->name);
	fprintf(stdout, "easyRNG %g\n", ((double) (eend - estart))/CLOCKS_PER_SEC);
	fprintf(stdout, "GSL %g\n", ((double) (gend - gstart))/CLOCKS_PER_SEC);
}

int main(int argc, char *argv[]) {

	run_test(gsl_rng_mt19937, easy_rng_mt19937);
	run_test(gsl_rng_ranlux, easy_rng_ranlux24);
	run_test(gsl_rng_ranlux389, easy_rng_ranlux48);
	run_test(gsl_rng_minstd, easy_rng_minstd_rand0);
	run_test(gsl_rng_fishman20, easy_rng_minstd_rand);

	return 0;
}


#ifndef __QUANTUM_H__
#define __QUANTUM_H__

#include <iostream>
#include <math.h>
#include <stdio.h>
#include <string.h>
#include <gsl/gsl_sf.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_linalg.h>
#include <gsl/gsl_permutation.h>
#include <gsl/gsl_complex.h>
#include <gsl/gsl_complex_math.h>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_vector.h>
#include <gsl/gsl_eigen.h>
#include <gsl/gsl_poly.h>
#include <gsl/gsl_integration.h>
#include <gsl/gsl_sort.h>
#include <gsl/gsl_sort_vector.h>

#endif

void WignerTransform_au(gsl_matrix_complex * dest, gsl_vector_complex * source, int grades, double dx, double dp);
void WignerTransform_SI(gsl_matrix_complex * dest,gsl_vector_complex * source, int grades, double dx, double dp);


#pragma once
#include "global.h"

#ifdef USE_MKL
#include <mkl.h>
inline void* _malloc(size_t size) {return mkl_malloc(size,64);}
inline void _free(void* ptr) {return mkl_free(ptr);}
#else
#include <cblas.h>
inline void* _malloc(size_t size) {return malloc(size);}
inline void _free(void* ptr) {return free(ptr);}
#endif

void init_glorot(size_t dim_x, size_t dim_y, vec_t &weight, unsigned seed);

// symmetric sparse matrix transpose
void symmetric_csr_transpose(int N, int nnz, int* A_idx_ptr, int* A_nnz_idx,
                             float* A_nonzeros, float* &B_nonzeros);

// get accuracy for single-class problems
float masked_accuracy_single(int begin, int end, int count, int num_classes,
                             mask_t* masks, float* preds, label_t* ground_truth);

// get accuracy for multi-class problems
float masked_accuracy_multi(int begin, int end, int count, int num_classes,
                            mask_t* masks, float* preds, label_t* ground_truth);

void matmul_naive(const size_t x, const size_t y, const size_t z,
                  const float* A, const float* B, float* C);
void matmul(const size_t x, const size_t y, const size_t z,
            const float_t* A, const float_t* B, float* C,
            bool transA = false, bool transB = false, bool accum = false);

void l2norm(int n, int dim, const float* in, float* out);
void d_l2norm(int n, int dim, const float* feat_in, const float* grad_in, float* grad_out);
void bias_mv(int n, int len, float* x, float* b);
void reduce_sum(int n, int len, float* x, float* a);
void reduce_sum(int n, int len, float* x, vec_t &a);

// single-precision dense matrix multiply
void sgemm_cpu(const CBLAS_TRANSPOSE TransA, const CBLAS_TRANSPOSE TransB,
               const int M, const int N, const int K, const float alpha,
               const float* A, const float* B, const float beta, float* C);

void csr2csc(int nrows, int ncols, int nnz,
             const float* values, const int* rowptr, const int* colidx,
             float* valuesT, int* rowptrT, int* colidxT);

// single-precision sparse matrix dense matrix multiply, C = A * B, A is sparse
void csrmm_cpu(const int M, const int N, const int K, const int nnz,
               const float alpha, float* A_nonzeros, int* A_idx_ptr,
               int* A_nonzero_idx, const float* B, const float beta, float* C, bool transA = false);

void spmm(size_t x, size_t y, size_t z, size_t nnz,
          float* A_nonzeros, int* A_idx_ptr, int* A_nnz_idx,
          const float* B, float* C, float* temp = NULL,
          bool transA = false, bool transB = false, bool accum = false);

// matrix-vector multiply
void mvmul(const CBLAS_TRANSPOSE TransA, const int M, const int N,
           const float alpha, const float* A, const float* x, const float beta, float* y);

//! add 2 arrays for n elements
void vadd_cpu(int n, const float* a, const float* b, float* out);
void scaled_vadd_cpu(int n, const float a, const float* x, const float* y, float* z);

// atomic y[i] += x[i]
void atomic_vreduce_cpu(size_t n, const float* x, float* y);

//! multiply n elements of vector by scalar
void scal(size_t n, const float alpha, float* x);
void scale(int n, const float alpha, const float* x, float* y);
void mul_scalar(size_t n, const float alpha, const float* x, float* y);

//! do dot product of 2 vectors
float dot(int n, const float* x, const float* y);

// concatenation of two vectors into one
void concat(size_t n, const float* x, const float* y, float* z);

// SAXPY stands for Single-precision A*X Plus Y"
void axpy(size_t n, const float a, float* x, float* y);

// Returns the index of the maximum value
int argmax(int num, const float* arr);

//! clear n elements of a vector
void clear_cpu(int n, float* in);

//! copy vector from in -> out; first len elements
void copy_cpu(size_t len, const float* in, float* out);

// dropout functions randomly remove weights
void dropout_cpu(size_t n, size_t m, float scale, float dropout_rate,
                 const float* in, mask_t* mask, float* out);

// dropout derivative: use existing dropouts in masks instead of generating them;
void d_dropout_cpu(size_t n, size_t m, float scale, const float* in,
                   const mask_t* mask, float* out);

void leaky_relu(float epsilon, float in, float &out);
void d_leaky_relu(float epsilon, float in, float data, float &out);
void relu_cpu(size_t n, const float* in, float* out);
void d_relu_cpu(size_t n, const float* in, const float* data, float* out);
void leaky_relu_cpu(size_t n, float epsilon, const float* in, float* out);
void d_leaky_relu_cpu(size_t n, float epsilon, const float* in, const float* data, float* out);
void softmax(size_t n, const float* in, float* out);
void d_softmax(int n, const float* p, const float* dp, float* dy);
void sigmoid(size_t n, const float* in, float* out);
void d_sigmoid(size_t n, const float*, const float* p, float* dy, const float* dp);
float cross_entropy(size_t n, const float* y, const float* p);
void d_cross_entropy(size_t n, const float* y, const float* p, float* d);
float sigmoid_cross_entropy(size_t n, const label_t* y, const float* p);

// use sigmoid instead of softmax for multi-class datasets, e.g. ppi, yelp and amazon
//inline float sigmoid_func(float x) { return 0.5 * tanh(0.5 * x) + 0.5; }
inline float sigmoid_func(float x) { return 1. / (1. + expf(-x)); }

// GPU operators
bool isnan_gpu(int n, const float_t* array); // does array contain any 'nan' element
void init_const_gpu(int n, float_t value, float_t* array);
void copy_gpu(int len, const float_t* in, float_t* out);
void vadd_gpu(const int n, const float_t* a, const float_t* b, float_t* out); // vector add
void axpy_gpu(const int n, const float_t a, const float_t* x, float_t* y);    // axpy
void relu_gpu(const int n, const float_t* in, float_t* out); // ReLU
void d_relu_gpu(const int n, const float_t* in_diff, const float_t* data, float_t* out_diff); // ReLU derivative
void leaky_relu_gpu(const int n, const float_t epsilon, const float_t* in, float_t* out); // Leaky ReLU
void d_leaky_relu_gpu(const int n, const float_t epsilon,
                      const float_t* in_diff, const float_t* data,
                      float_t* out_diff); // Leaky ReLU derivative
void dropout_gpu(int n, float scale, float drop_rate, const float* in, mask_t* masks, float* out); // dropout
void d_dropout_gpu(int n, float scale, const float* in, const mask_t* masks, float* out); // dropout derivative
void sgemm_gpu(const CBLAS_TRANSPOSE TransA, const CBLAS_TRANSPOSE TransB,
               const int M, const int N, const int K, const float alpha,
               const float* A, const float* B, const float beta, float* C);
void csrmm_gpu(const int M, const int N, const int K, const int nnz,
               const float alpha, const float* A_nonzeros, const int* A_idx_ptr,
               const int* A_nonzero_idx, const float* B, const float beta,
               float* trans_C, float* C);
void softmax_cross_entropy_gpu(int len, int begin, int end,
                               const float_t* in_data, const mask_t* masks,
                               const label_t* labels, float_t* loss,
                               float_t* out_data);
void d_softmax_cross_entropy_gpu(int len, int bengin, int end,
                                 const mask_t* masks, const label_t* labels,
                                 const float_t* out_data, float_t* diff);
void sigmoid_cross_entropy_gpu(int len, int begin, int end,
                               const float_t* in_data, const mask_t* masks,
                               const label_t* labels, float_t* loss,
                               float_t* out_data);
void d_sigmoid_cross_entropy_gpu(int len, int bengin, int end,
                                 const mask_t* masks, const label_t* labels,
                                 const float_t* out_data, float_t* diff);
void scal_gpu(const int n, const float alpha, float* X);
void add_scalar_gpu(const int n, const float_t alpha, float_t* Y);
void rng_uniform_gpu(size_t n, const float_t a, const float_t b, float_t* r);
void l2_norm_gpu(size_t x, size_t y, const float_t* in, float_t* out);
void d_l2_norm_gpu(size_t x, size_t y, const float_t* in_data, float_t* in_diff, float_t* out_diff);
acc_t l2_norm_gpu(int n, const float_t* in);
acc_t masked_avg_loss_gpu(int begin, int end, int count, mask_t* masks, float_t* loss);

bool is_allocated_device(float_t* data);
void copy_masks_device(int n, mask_t* h_masks, mask_t*& d_masks);
void float_malloc_device(int n, float_t*& ptr);
void float_free_device(float_t*& ptr);
void copy_float_device(int n, float* h_ptr, float* d_ptr);
void copy_float_host(int n, const float* d_ptr, float* h_ptr);
void uint_malloc_device(int n, uint32_t*& ptr);
void uint_free_device(uint32_t*& ptr);
void copy_uint_device(int n, uint32_t* h_ptr, uint32_t* d_ptr);
void uint8_malloc_device(int n, uint8_t*& ptr);
void uint8_free_device(uint8_t*& ptr);
void copy_uint8_device(int n, uint8_t* h_ptr, uint8_t* d_ptr);
void gpu_rng_uniform(size_t n, float* r);


#include <stdio.h>
#include <string.h>
#include <time.h>
#include <stdlib.h>
#include <unistd.h>
#include <stdbool.h>
#include <omp.h>
#include <cblas.h>

//Threshold for testing validity of matrix matrix multiplication
#define ERROR_THRESHOLD 0.0001

//For measuring wall time using omp_get_wtime()
static double start;
static double end;

//Serial version. Do not change this!
void serial_mxm(const double *A, const double *B, double *C, int m, int n, int k)
{
  clock_t start, end;
  start = clock();
  for (int i = 0; i < m; i++)
  {
    for (int j = 0; j < n; j++)
    {
      C[i * n + j] = 0;
      for (int l = 0; l < k; l++)
      {
        C[i * n + j] += A[i * k + l] * B[l * n + j];
      }
    }
  }
  end = clock();
  printf("\nUsage: %f", (((double)end - (double)start) / (double)CLOCKS_PER_SEC));
}

void omp_mxm(double *A, double *B, double *C, int m, int n, int k)
{
  clock_t start, end;
  start = clock();
#pragma omp parallell for
  for (int i = 0; i < m; i++)
  {
    for (int j = 0; j < n; j++)
    {
      C[i * n + j] = 0;
      for (int l = 0; l < k; l++)
      {
        C[i * n + j] += A[i * k + l] * B[l * n + j];
      }
    }
  }
  end = clock();
  printf("\nUsage: %f", (((double)end - (double)start) / (double)CLOCKS_PER_SEC));
}

void blas_mxm(double *A, double *B, double *C, int m, int n, int k)
{
  clock_t start, end;
  start = clock();
  cblas_dgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, m, n, k, 1.0, A, k, B, n, 1.0, C, n);
  end = clock();
  printf("\nUsage: %f", (((double)end - (double)start) / (double)CLOCKS_PER_SEC));
}

int main(const unsigned int argc, char **argv)
{
  if (argc <= optind)
  {
    printf("Please provide version:\n");
    printf("\ts(serial),\n");
    printf("\to(penmp) or\n");
    printf("\tb(las)\n");
    return 0;
  }
  char input = argv[optind][0];
  optind++;
  //Simple assumptions that any additional arguments means we want to test the results
  bool test = !(argc <= optind);

  int m = 2000;
  int n = 1000;
  int k = 200;

  double *A = (double *)malloc(m * k * sizeof(double));
  double *B = (double *)malloc(k * n * sizeof(double));
  double *C = (double *)malloc(m * n * sizeof(double));

  //Intializing matrix data
  for (int i = 0; i < (m * k); i++)
  {
    A[i] = (double)(i + 1);
  }

  for (int i = 0; i < (k * n); i++)
  {
    B[i] = (double)(-i - 1);
  }

  for (int i = 0; i < (m * n); i++)
  {
    C[i] = 0.0;
  }
  switch (input)
  {
  case 's':
    serial_mxm(A, B, C, m, n, k);
    break;
  case 'o':
    omp_mxm(A, B, C, m, n, k);
    break;
  case 'b':
    blas_mxm(A, B, C, m, n, k);
    break;
  default:
    printf("Please provide version:\n");
    printf("\ts(serial),\n");
    printf("\to(penmp) or\n");
    printf("\tb(las)\n");
    return 0;
  }

  printf("\nTop left of A:\n");
  for (int i = 0; i < 5; i++)
  {
    for (int j = 0; j < 5; j++)
    {
      printf("%8.2f\t", A[i * k + j]);
    }
    printf("\n");
  }

  printf("\nTop left of B:\n");
  for (int i = 0; i < 5; i++)
  {
    for (int j = 0; j < 5; j++)
    {
      printf("%8.2f\t", B[i * n + j]);
    }
    printf("\n");
  }

  printf("\nTop left of C:\n");
  for (int i = 0; i < 5; i++)
  {
    for (int j = 0; j < 5; j++)
    {
      printf("%8.2f\t", C[i * n + j]);
    }
    printf("\n");
  }

  if (test)
  {
    double *C2 = (double *)malloc(m * n * sizeof(double));
    serial_mxm(A, B, C2, m, n, k);
    bool correct = true;
    for (int i = 0; i < (m * n); i++)
    {
      if (abs(C[i] - C2[i]) > ERROR_THRESHOLD)
      {
        correct = false;
      }
    }
    if (correct)
    {
      printf("\nMatrix multiplication succeeded\n");
    }
    else
    {
      printf("\nMatrix multiplication failed\n");
      printf("Top left of correct C:\n");
      for (int i = 0; i < 5; i++)
      {
        for (int j = 0; j < 5; j++)
        {
          printf("%8.2f\t", C2[i * n + j]);
        }
        printf("\n");
      }
    }
  }

  printf("\nVersion: %c, time: %.4f\n", input, end - start);

  return 0;
}


/* specfunc/hyperg.h
 *
 * 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 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.
 */

/* Author:  G. Jungman */

/* Miscellaneous implementations of use
 * for evaluation of hypergeometric functions.
 */
#ifndef _HYPERG_H_
#define _HYPERG_H_

#include <gsl/gsl_sf_result.h>

/* Direct implementation of 1F1 series.
 */
int
gsl_sf_hyperg_1F1_series_e(const double a, const double b, const double x, gsl_sf_result * result);

/* Implementation of the 1F1 related to the
 * incomplete gamma function: 1F1(1,b,x), b >= 1.
 */
int
gsl_sf_hyperg_1F1_1_e(double b, double x, gsl_sf_result * result);

/* 1F1(1,b,x) for integer b >= 1
 */
int
gsl_sf_hyperg_1F1_1_int_e(int b, double x, gsl_sf_result * result);

/* Implementation of large b asymptotic.
 * [Bateman v. I, 6.13.3 (18)]
 * [Luke, The Special Functions and Their Approximations v. I, p. 129, 4.8 (4)]
 *
 * a^2 << b, |x|/|b| < 1 - delta
 */
int
gsl_sf_hyperg_1F1_large_b_e(const double a, const double b, const double x, gsl_sf_result * result);

/* Implementation of large b asymptotic.
 *
 * Assumes a > 0 is small, x > 0, and |x|<|b|.
 */
int
gsl_sf_hyperg_U_large_b_e(const double a, const double b, const double x,
                             gsl_sf_result * result,
                             double * ln_multiplier
                             );

/* Implementation of 2F0 asymptotic series.
 */
int
gsl_sf_hyperg_2F0_series_e(const double a, const double b, const double x, int n_trunc,
                              gsl_sf_result * result);

#endif /* !_HYPERG_H_ */


/*
 *  SpecialFunctions.h
 *  PSCAcoustic
 *
 *  Special functions and more.
 *
 *
 * Copyright 2014 Pierre-David Letourneau
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to
 * deal in the Software without restriction, including without limitation the
 * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
 * sell copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
*/

#ifndef SPECIALFUNCTIONS_H
#define SPECIALFUNCTIONS_H

#include "General.h"

// GSL library
//#include "../gsl-1.16/include/gsl/gsl_sf_bessel.h" //spherical bessel functions
#include <gsl/gsl_sf_bessel.h> //spherical bessel functions
#include "./gsl-1.16/include/gsl/gsl_sf_legendre.h" //associated Legendre polynomials (for spherical harmonics)
#include "./gsl-1.16/include/gsl/gsl_complex_math.h"
#include "./gsl-1.16/include/gsl/gsl_sf_gamma.h"
#include "./gsl-1.16/include/gsl/gsl_sf_hyperg.h"
#include "./gsl-1.16/include/gsl/gsl_sf_coupling.h"
#include "./gsl-1.16/include/gsl/gsl_integration.h"
#include "./gsl-1.16/include/gsl/gsl_complex.h"
#include "./gsl-1.16/include/gsl/gsl_sf_coupling.h"
#include "./gsl-1.16/include/gsl/gsl_math.h"
#include "./gsl-1.16/include/gsl/gsl_ieee_utils.h"
#include "./gsl-1.16/include/gsl/gsl_errno.h"
#include "./gsl-1.16/include/gsl/gsl_integration.h"

// Complex Bessel functions
//#include "../AmosBessel/libAmosBessel.h"

//GSL wrappers (Cannot handle complex  arguments)

//Hankel function of the 1st kind (gsl)
inline complex  gsl_sf_hankel_1( int l, double x){
  return gsl_sf_bessel_jl(l,x) + gsl_sf_bessel_yl(l,x) * CI;
}

inline double gsl_sf_legendre( int l, int m, double x){
  double C = 1;
  if( m < 0 )
    C *= pow(-1, std::abs(m));

  return (C * gsl_sf_legendre_sphPlm ( l, std::abs(m), x) );
}

// Y_lm (spherical harmonics)
//\sqrt{(2l+1)/(4\pi)} \sqrt{(l-m)!/(l+m)!} P_l^m(x) exp(I*m*phi) )
inline complex  gsl_sf_harmonic( int l, int m, double theta, double phi){
  if( l >= std::abs(m) ){
    return gsl_sf_legendre(l, m, cos(theta)) * exp(m*phi*CI);
  } else {
    return 0.;
  }
}

inline complex  gsl_sf_clebsch(int l1, int l2, int l3, int m1, int m2, int m3){
  return pow(-1., (double) l1-l2+m3) * sqrt(2.*l3+1.) * gsl_sf_coupling_3j(2*l1, 2*l2, 2*l3, 2*m1, 2*m2, -2*m3);
}

inline double gsl_sf_gaunt(int n, int m, int nu, int mu, int q){
  return pow(-1., (double) (m+mu) ) * sqrt((2*n+1)*(2*nu+1)*(2*q+1)/(4*PI)) *
    gsl_sf_coupling_3j(2.*n,2.*nu,2.*q,0,0,0) * gsl_sf_coupling_3j(2.*n,2.*nu,2.*q,2.*m,2.*mu,-2.*(m+mu));
}

// TODO: FIX AMOS
//Amos wrappers for Bessel functions (these can handle complex  arguments)
inline complex  Amos_sf_bessel_jl(int l, complex  z){
  	complex  value;
	//AmosBessel('j', z, (double) l, 1, 0, &value);
	value = gsl_sf_bessel_jl(l,std::real(z));
	return value;
}

inline complex  Amos_sf_bessel_yl(int l, complex  z){
  	complex  value;
	//AmosBessel('y', z, (double) l, 1, 0, &value);
	value = gsl_sf_bessel_yl(l,std::real(z));
	return value;
}

inline complex  Amos_sf_hankel_1(int l, complex  z){
  //complex  valueR;
	//AmosBessel('j', z, (double) l, 1, 0, &valueR);
  	//complex  valueI;
	//AmosBessel('y', z, (double) l, 1, 0, &valueI);
	//return (valueR + valueI*CI);

	complex value = gsl_sf_hankel_1(l, std::real(z));
	return value;
}

// Derivatives

inline double gsl_sf_bessel_jl_prime(int l, double z){
  return l/z * gsl_sf_bessel_jl(l,z) - gsl_sf_bessel_jl(l+1,z);
}

inline complex  gsl_sf_hankel_l_prime(int l, double z){
  return l/z * gsl_sf_hankel_1( l, z) - gsl_sf_hankel_1(l+1, z);
}

inline complex  Amos_sf_bessel_jl_prime(double l, complex  z){
  return l/z * Amos_sf_bessel_jl(l,z) - Amos_sf_bessel_jl(l+1,z);
}

inline complex  Amos_sf_hankel_l_prime(double l, complex  z){
  return l/z * Amos_sf_hankel_1( l, z) - Amos_sf_hankel_1(l+1, z);
}

// Chebyshev polynomial evaluation
inline double Cheb(int n, double x){
  return cos(n*acos(x));
}

// Gauss-Legendre quadrature
inline void getGaussLegendreQuad(int N, std::vector<double>& x, std::vector<double>& w)
{
  x.resize(N);
  w.resize(N);

  gsl_integration_glfixed_table* GLTable = gsl_integration_glfixed_table_alloc(N);

  for( int n = 0; n < (N+1)/2; ++n ) {
    x[n + N/2] = GLTable->x[n];
    w[n + N/2] = GLTable->w[n];
  }
  for( int n = 0; n < N/2; ++n ) {
    x[n] = -x[N-1-n];
    w[n] =  w[N-1-n];
  }

  gsl_integration_glfixed_table_free( GLTable );
}

// Gauss-Chebyshev quadrature
inline void getGaussChebyshevQuad(int N, std::vector<double>& x, std::vector<double>& w)
{
  x.resize(N);
  w.resize(N);

  for( int n = 0; n < N; n++ ) {
    x[n] = cos( (2.*n+1.)/(2.*N)*PI  );
    w[n] = 2./N;
  }
}

//---------General Functions--------//

#define ISODD(x) ((x) & 1)
#define ISNAN(x) ((x) != (x))

//TODO : still needed?
inline int choose( int n, int r ){

  if( (n < r) || (r < 0) || (n < 0) ){
    return 0;
  } else {
    return gsl_sf_gamma(n+1) / (gsl_sf_gamma(r+1) * gsl_sf_gamma(n-r+1));
  }

}

// Implementation os Sparse Matrix class
// TODO:  this takes a lot of memory?
class SpMatrix
{
  map<std::pair<int,int>, complex > M;

 public:

  // Empty constructor
  SpMatrix(){}

  ~SpMatrix(){}

  void add(int i, int j, complex  val){
    assert( i >= 0 );
    assert( j >= 0 );

    std::pair<int,int> coord = std::make_pair(i,j);
    if( M.count(coord) ){
      M.insert(std::make_pair(coord, val));
    } else {
      M[coord] = val;
    }

  }

  complex  operator () (int i, int j){
    std::pair<int,int> coord = std::make_pair(i,j);
    complex  val = 0.;

    ( M.count(coord) ) ? val = M.find(coord)->second : val = 0.;

    return val;
  }

  inline int size() { return M.size(); }

};

#endif