// Copyright 2021 DeepMind Technologies Limited
//
// Licensed 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.

#include "engine/engine_util_blas.h"

#include <string.h>

#include <mujoco/mjmacro.h>
#include <mujoco/mjtnum.h>

#ifdef mjUSEPLATFORMSIMD
  #if defined(__AVX__) && defined(mjUSEDOUBLE)
    #define mjUSEAVX
    #include "immintrin.h"
  #endif
#endif



//------------------------------ 3D vector and matrix-vector operations ----------------------------

// res = 0
void mju_zero3(mjtNum res[3]) {
  res[0] = 0;
  res[1] = 0;
  res[2] = 0;
}



// res = vec
void mju_copy3(mjtNum res[3], const mjtNum data[3]) {
  res[0] = data[0];
  res[1] = data[1];
  res[2] = data[2];
}



// res = vec*scl
void mju_scl3(mjtNum res[3], const mjtNum vec[3], mjtNum scl) {
  res[0] = vec[0] * scl;
  res[1] = vec[1] * scl;
  res[2] = vec[2] * scl;
}



// res = vec1 + vec2
void mju_add3(mjtNum res[3], const mjtNum vec1[3], const mjtNum vec2[3]) {
  res[0] = vec1[0] + vec2[0];
  res[1] = vec1[1] + vec2[1];
  res[2] = vec1[2] + vec2[2];
}



// res = vec1 - vec2
void mju_sub3(mjtNum res[3], const mjtNum vec1[3], const mjtNum vec2[3]) {
  res[0] = vec1[0] - vec2[0];
  res[1] = vec1[1] - vec2[1];
  res[2] = vec1[2] - vec2[2];
}



// res += vec
void mju_addTo3(mjtNum res[3], const mjtNum vec[3]) {
  res[0] += vec[0];
  res[1] += vec[1];
  res[2] += vec[2];
}



// res -= vec
void mju_subFrom3(mjtNum res[3], const mjtNum vec[3]) {
  res[0] -= vec[0];
  res[1] -= vec[1];
  res[2] -= vec[2];
}



// res += vec*scl
void mju_addToScl3(mjtNum res[3], const mjtNum vec[3], mjtNum scl) {
  res[0] += vec[0] * scl;
  res[1] += vec[1] * scl;
  res[2] += vec[2] * scl;
}



// res = vec1 + vec2*scl
void mju_addScl3(mjtNum res[3], const mjtNum vec1[3], const mjtNum vec2[3], mjtNum scl) {
  res[0] = vec1[0] + scl*vec2[0];
  res[1] = vec1[1] + scl*vec2[1];
  res[2] = vec1[2] + scl*vec2[2];
}



// normalize vector, return length before normalization
mjtNum mju_normalize3(mjtNum vec[3]) {
  mjtNum norm = mju_sqrt(vec[0]*vec[0] + vec[1]*vec[1] + vec[2]*vec[2]);

  if (norm < mjMINVAL) {
    vec[0] = 1;
    vec[1] = 0;
    vec[2] = 0;
  } else {
    mjtNum normInv = 1/norm;
    vec[0] *= normInv;
    vec[1] *= normInv;
    vec[2] *= normInv;
  }

  return norm;
}



// compute vector length (without normalizing)
mjtNum mju_norm3(const mjtNum vec[3]) {
  return mju_sqrt(vec[0]*vec[0] + vec[1]*vec[1] + vec[2]*vec[2]);
}



// vector dot-product
mjtNum mju_dot3(const mjtNum vec1[3], const mjtNum vec2[3]) {
  return vec1[0]*vec2[0] + vec1[1]*vec2[1] + vec1[2]*vec2[2];
}



// Cartesian distance between 3D vectors
mjtNum mju_dist3(const mjtNum pos1[3], const mjtNum pos2[3]) {
  mjtNum dif[3] = {pos1[0]-pos2[0], pos1[1]-pos2[1], pos1[2]-pos2[2]};
  return mju_sqrt(dif[0]*dif[0] + dif[1]*dif[1] + dif[2]*dif[2]);
}



// multiply vector by 3D rotation matrix
void mju_rotVecMat(mjtNum res[3], const mjtNum vec[3], const mjtNum mat[9]) {
  mjtNum tmp[3] = {
    mat[0]*vec[0] + mat[1]*vec[1] + mat[2]*vec[2],
    mat[3]*vec[0] + mat[4]*vec[1] + mat[5]*vec[2],
    mat[6]*vec[0] + mat[7]*vec[1] + mat[8]*vec[2]
  };
  res[0] = tmp[0];
  res[1] = tmp[1];
  res[2] = tmp[2];
}



// multiply vector by transposed 3D rotation matrix
void mju_rotVecMatT(mjtNum res[3], const mjtNum vec[3], const mjtNum mat[9]) {
  mjtNum tmp[3] = {
    mat[0]*vec[0] + mat[3]*vec[1] + mat[6]*vec[2],
    mat[1]*vec[0] + mat[4]*vec[1] + mat[7]*vec[2],
    mat[2]*vec[0] + mat[5]*vec[1] + mat[8]*vec[2]
  };
  res[0] = tmp[0];
  res[1] = tmp[1];
  res[2] = tmp[2];
}



//------------------------------ 4D vector and matrix-vector operations ----------------------------

// res = 0
void mju_zero4(mjtNum res[4]) {
  res[0] = 0;
  res[1] = 0;
  res[2] = 0;
  res[3] = 0;
}



// res = (1,0,0,0)
void mju_unit4(mjtNum res[4]) {
  res[0] = 1;
  res[1] = 0;
  res[2] = 0;
  res[3] = 0;
}


// res = vec
void mju_copy4(mjtNum res[4], const mjtNum data[4]) {
  res[0] = data[0];
  res[1] = data[1];
  res[2] = data[2];
  res[3] = data[3];
}



// normalize vector, return length before normalization
mjtNum mju_normalize4(mjtNum vec[4]) {
  mjtNum norm = mju_sqrt(vec[0]*vec[0] + vec[1]*vec[1] + vec[2]*vec[2] + vec[3]*vec[3]);

  if (norm < mjMINVAL) {
    vec[0] = 1;
    vec[1] = 0;
    vec[2] = 0;
    vec[3] = 0;
  } else {
    mjtNum normInv = 1/norm;
    vec[0] *= normInv;
    vec[1] *= normInv;
    vec[2] *= normInv;
    vec[3] *= normInv;
  }

  return norm;
}



//------------------------------ vector operations -------------------------------------------------

// res = 0
void mju_zero(mjtNum* res, int n) {
  if (n > 0) {
    memset(res, 0, n*sizeof(mjtNum));
  }
}



// res = val
void mju_fill(mjtNum* res, mjtNum val, int n) {
  for (int i=0; i < n; i++) {
    res[i] = val;
  }
}



// res = vec
void mju_copy(mjtNum* res, const mjtNum* vec, int n) {
  if (n > 0) {
    memcpy(res, vec, n*sizeof(mjtNum));
  }
}



// sum(vec)
mjtNum mju_sum(const mjtNum* vec, int n) {
  mjtNum res = 0;

  for (int i=0; i < n; i++) {
    res += vec[i];
  }

  return res;
}



// sum(abs(vec))
mjtNum mju_L1(const mjtNum* vec, int n) {
  mjtNum res = 0;

  for (int i=0; i < n; i++) {
    res += mju_abs(vec[i]);
  }

  return res;
}



// res = vec*scl
void mju_scl(mjtNum* res, const mjtNum* vec, mjtNum scl, int n) {
  int i = 0;

#ifdef mjUSEAVX
  int n_4 = n - 4;

  // vector part
  if (n_4 >= 0) {
    __m256d sclpar, val1, val1scl;

    // init
    sclpar = _mm256_set1_pd(scl);

    // parallel computation
    while (i <= n_4) {
      val1 = _mm256_loadu_pd(vec+i);
      val1scl = _mm256_mul_pd(val1, sclpar);
      _mm256_storeu_pd(res+i, val1scl);
      i += 4;
    }
  }

  // process remaining
  int n_i = n - i;
  if (n_i == 3) {
    res[i] = vec[i]*scl;
    res[i+1] = vec[i+1]*scl;
    res[i+2] = vec[i+2]*scl;
  } else if (n_i == 2) {
    res[i] = vec[i]*scl;
    res[i+1] = vec[i+1]*scl;
  } else if (n_i == 1) {
    res[i] = vec[i]*scl;
  }

#else
  for (; i < n; i++) {
    res[i] = vec[i]*scl;
  }
#endif
}



// res = vec1 + vec2
void mju_add(mjtNum* res, const mjtNum* vec1, const mjtNum* vec2, int n) {
  int i = 0;

#ifdef mjUSEAVX
  int n_4 = n - 4;

  // vector part
  if (n_4 >= 0) {
    __m256d sum, val1, val2;

    // parallel computation
    while (i <= n_4) {
      val1 = _mm256_loadu_pd(vec1+i);
      val2 = _mm256_loadu_pd(vec2+i);
      sum = _mm256_add_pd(val1, val2);
      _mm256_storeu_pd(res+i, sum);
      i += 4;
    }
  }

  // process remaining
  int n_i = n - i;
  if (n_i == 3) {
    res[i] = vec1[i] + vec2[i];
    res[i+1] = vec1[i+1] + vec2[i+1];
    res[i+2] = vec1[i+2] + vec2[i+2];
  } else if (n_i == 2) {
    res[i] = vec1[i] + vec2[i];
    res[i+1] = vec1[i+1] + vec2[i+1];
  } else if (n_i == 1) {
    res[i] = vec1[i] + vec2[i];
  }

#else
  for (; i < n; i++) {
    res[i] = vec1[i] + vec2[i];
  }
#endif
}



// res = vec1 - vec2
void mju_sub(mjtNum* res, const mjtNum* vec1, const mjtNum* vec2, int n) {
  int i = 0;

#ifdef mjUSEAVX
  int n_4 = n - 4;

  // vector part
  if (n_4 >= 0) {
    __m256d dif, val1, val2;

    // parallel computation
    while (i <= n_4) {
      val1 = _mm256_loadu_pd(vec1+i);
      val2 = _mm256_loadu_pd(vec2+i);
      dif = _mm256_sub_pd(val1, val2);
      _mm256_storeu_pd(res+i, dif);
      i += 4;
    }
  }

  // process remaining
  int n_i = n - i;
  if (n_i == 3) {
    res[i] = vec1[i] - vec2[i];
    res[i+1] = vec1[i+1] - vec2[i+1];
    res[i+2] = vec1[i+2] - vec2[i+2];
  } else if (n_i == 2) {
    res[i] = vec1[i] - vec2[i];
    res[i+1] = vec1[i+1] - vec2[i+1];
  } else if (n_i == 1) {
    res[i] = vec1[i] - vec2[i];
  }

#else
  for (; i < n; i++) {
    res[i] = vec1[i] - vec2[i];
  }
#endif
}



// res += vec
void mju_addTo(mjtNum* res, const mjtNum* vec, int n) {
  int i = 0;

#ifdef mjUSEAVX
  int n_4 = n - 4;

  // vector part
  if (n_4 >= 0) {
    __m256d sum, val1, val2;

    // parallel computation
    while (i <= n_4) {
      val1 = _mm256_loadu_pd(res+i);
      val2 = _mm256_loadu_pd(vec+i);
      sum = _mm256_add_pd(val1, val2);
      _mm256_storeu_pd(res+i, sum);
      i += 4;
    }
  }

  // process remaining
  int n_i = n - i;
  if (n_i == 3) {
    res[i] += vec[i];
    res[i+1] += vec[i+1];
    res[i+2] += vec[i+2];
  } else if (n_i == 2) {
    res[i] += vec[i];
    res[i+1] += vec[i+1];
  } else if (n_i == 1) {
    res[i] += vec[i];
  }

#else
  for (; i < n; i++) {
    res[i] += vec[i];
  }
#endif
}



// res -= vec
void mju_subFrom(mjtNum* res, const mjtNum* vec, int n) {
  int i = 0;

#ifdef mjUSEAVX
  int n_4 = n - 4;

  // vector part
  if (n_4 >= 0) {
    __m256d dif, val1, val2;

    // parallel computation
    while (i <= n_4) {
      val1 = _mm256_loadu_pd(res+i);
      val2 = _mm256_loadu_pd(vec+i);
      dif = _mm256_sub_pd(val1, val2);
      _mm256_storeu_pd(res+i, dif);
      i += 4;
    }
  }

  // process remaining
  int n_i = n - i;
  if (n_i == 3) {
    res[i] -= vec[i];
    res[i+1] -= vec[i+1];
    res[i+2] -= vec[i+2];
  } else if (n_i == 2) {
    res[i] -= vec[i];
    res[i+1] -= vec[i+1];
  } else if (n_i == 1) {
    res[i] -= vec[i];
  }

#else
  for (; i < n; i++) {
    res[i] -= vec[i];
  }
#endif
}



// res += vec*scl
void mju_addToScl(mjtNum* res, const mjtNum* vec, mjtNum scl, int n) {
  int i = 0;

#ifdef mjUSEAVX
  int n_4 = n - 4;

  // vector part
  if (n_4 >= 0) {
    __m256d sclpar, sum, val1, val2, val2scl;

    // init
    sclpar = _mm256_set1_pd(scl);

    // parallel computation
    while (i <= n_4) {
      val1 = _mm256_loadu_pd(res+i);
      val2 = _mm256_loadu_pd(vec+i);
      val2scl = _mm256_mul_pd(val2, sclpar);
      sum = _mm256_add_pd(val1, val2scl);
      _mm256_storeu_pd(res+i, sum);
      i += 4;
    }
  }

  // process remaining
  int n_i = n - i;
  if (n_i == 3) {
    res[i] += vec[i]*scl;
    res[i+1] += vec[i+1]*scl;
    res[i+2] += vec[i+2]*scl;
  } else if (n_i == 2) {
    res[i] += vec[i]*scl;
    res[i+1] += vec[i+1]*scl;
  } else if (n_i == 1) {
    res[i] += vec[i]*scl;
  }

#else
  for (; i < n; i++) {
    res[i] += vec[i]*scl;
  }
#endif
}

// res = vec1 + vec2*scl
void mju_addScl(mjtNum* res, const mjtNum* vec1, const mjtNum* vec2, mjtNum scl, int n) {
  int i = 0;

#if defined(__AVX__) && defined(mjUSEAVX)  && defined(mjUSEDOUBLE)
  int n_4 = n - 4;

  // vector part
  if (n_4 >= 0) {
    __m256d sclpar, sum, val1, val2, val2scl;

    // init
    sclpar = _mm256_set1_pd(scl);

    // parallel computation
    while (i <= n_4) {
      val1 = _mm256_loadu_pd(vec1+i);
      val2 = _mm256_loadu_pd(vec2+i);
      val2scl = _mm256_mul_pd(val2, sclpar);
      sum = _mm256_add_pd(val1, val2scl);
      _mm256_storeu_pd(res+i, sum);
      i += 4;
    }
  }

  // process remaining
  int n_i = n - i;
  if (n_i == 3) {
    res[i] = vec1[i] + vec2[i]*scl;
    res[i+1] = vec1[i+1] + vec2[i+1]*scl;
    res[i+2] = vec1[i+2] + vec2[i+2]*scl;
  } else if (n_i == 2) {
    res[i] = vec1[i] + vec2[i]*scl;
    res[i+1] = vec1[i+1] + vec2[i+1]*scl;
  } else if (n_i == 1) {
    res[i] = vec1[i] + vec2[i]*scl;
  }

#else
  for (; i < n; i++) {
    res[i] = vec1[i] + vec2[i]*scl;
  }
#endif
}



// normalize vector, return length before normalization
mjtNum mju_normalize(mjtNum* res, int n) {
  mjtNum norm = (mjtNum)mju_sqrt(mju_dot(res, res, n));
  mjtNum normInv;

  if (norm < mjMINVAL) {
    res[0] = 1;
    for (int i=1; i < n; i++) {
      res[i] = 0;
    }
  } else {
    normInv = 1/norm;
    for (int i=0; i < n; i++) {
      res[i] *= normInv;
    }
  }

  return norm;
}



// compute vector length (without normalizing)
mjtNum mju_norm(const mjtNum* res, int n) {
  return mju_sqrt(mju_dot(res, res, n));
}



// vector dot-product
mjtNum mju_dot(const mjtNum* vec1, const mjtNum* vec2, int n) {
  mjtNum res = 0;
  int i = 0;
  int n_4 = n - 4;
#ifdef mjUSEAVX

  // vector part
  if (n_4 >= 0) {
    __m256d sum, prod, val1, val2;
    __m128d vlow, vhigh, high64;

    // init
    val1 = _mm256_loadu_pd(vec1);
    val2 = _mm256_loadu_pd(vec2);
    sum = _mm256_mul_pd(val1, val2);
    i = 4;

    // parallel computation
    while (i <= n_4) {
      val1 = _mm256_loadu_pd(vec1+i);
      val2 = _mm256_loadu_pd(vec2+i);
      prod = _mm256_mul_pd(val1, val2);
      sum = _mm256_add_pd(sum, prod);
      i += 4;
    }

    // reduce
    vlow = _mm256_castpd256_pd128(sum);
    vhigh = _mm256_extractf128_pd(sum, 1);
    vlow = _mm_add_pd(vlow, vhigh);
    high64 = _mm_unpackhi_pd(vlow, vlow);
    res = _mm_cvtsd_f64(_mm_add_sd(vlow, high64));
  }

#else
  // do the same order of additions as the AVX intrinsics implementation.
  // this is faster than the simple for loop you'd expect for a dot product,
  // and produces exactly the same results.
  mjtNum res0 = 0;
  mjtNum res1 = 0;
  mjtNum res2 = 0;
  mjtNum res3 = 0;

  for (; i <= n_4; i+=4) {
    res0 += vec1[i] * vec2[i];
    res1 += vec1[i+1] * vec2[i+1];
    res2 += vec1[i+2] * vec2[i+2];
    res3 += vec1[i+3] * vec2[i+3];
  }
  res = (res0 + res2) + (res1 + res3);
#endif

  // process remaining
  int n_i = n - i;
  if (n_i == 3) {
    res += vec1[i]*vec2[i] + vec1[i+1]*vec2[i+1] + vec1[i+2]*vec2[i+2];
  } else if (n_i == 2) {
    res += vec1[i]*vec2[i] + vec1[i+1]*vec2[i+1];
  } else if (n_i == 1) {
    res += vec1[i]*vec2[i];
  }
  return res;
}

//------------------------------ matrix-vector operations ------------------------------------------

// multiply matrix and vector
void mju_mulMatVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int nr, int nc) {
  for (int r=0; r < nr; r++) {
    res[r] = mju_dot(mat + r*nc, vec, nc);
  }
}



// multiply transposed matrix and vector
void mju_mulMatTVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int nr, int nc) {
  mjtNum tmp;
  mju_zero(res, nc);

  for (int r=0; r < nr; r++) {
    if ((tmp = vec[r])) {
      mju_addToScl(res, mat+r*nc, tmp, nc);
    }
  }
}



// multiply square matrix with vectors on both sides: return vec1'*mat*vec2
mjtNum mju_mulVecMatVec(const mjtNum* vec1, const mjtNum* mat, const mjtNum* vec2, int n) {
  mjtNum res = 0;
  for (int i=0; i < n; i++) {
    res += vec1[i] * mju_dot(mat + i*n, vec2, n);
  }
  return res;
}



//------------------------------ matrix operations -------------------------------------------------

// transpose matrix
void mju_transpose(mjtNum* res, const mjtNum* mat, int nr, int nc) {
  for (int i=0; i < nr; i++) {
    for (int j=0; j < nc; j++) {
      res[j*nr+i] = mat[i*nc+j];
    }
  }
}



// symmetrize square matrix res = (mat + mat')/2
void mju_symmetrize(mjtNum* res, const mjtNum* mat, int n) {
  for (int i=0; i < n; i++) {
    res[i*(n+1)] = mat[i*(n+1)];
    for (int j=0; j < i; j++) {
      res[i*n+j] = res[j*n+i] = 0.5 * (mat[i*n+j] + mat[j*n+i]);
    }
  }
}



// identity matrix
void mju_eye(mjtNum* mat, int n) {
  mju_zero(mat, n*n);
  for (int i=0; i < n; i++) {
    mat[i*(n + 1)] = 1;
  }
}



//------------------------------ matrix-matrix operations ------------------------------------------

// multiply matrices, exploit sparsity of mat1
void mju_mulMatMat(mjtNum* res, const mjtNum* mat1, const mjtNum* mat2,
                   int r1, int c1, int c2) {
  mjtNum tmp;

  mju_zero(res, r1*c2);

  for (int i=0; i < r1; i++) {
    for (int k=0; k < c1; k++) {
      if ((tmp = mat1[i*c1+k])) {
        mju_addToScl(res+i*c2, mat2+k*c2, tmp, c2);
      }
    }
  }
}



// multiply matrices, second argument transposed
void mju_mulMatMatT(mjtNum* res, const mjtNum* mat1, const mjtNum* mat2,
                    int r1, int c1, int r2) {
  for (int i=0; i < r1; i++) {
    for (int j=0; j < r2; j++) {
      res[i*r2+j] = mju_dot(mat1+i*c1, mat2+j*c1, c1);
    }
  }
}



// compute M'*diag*M (diag=NULL: compute M'*M)
void mju_sqrMatTD(mjtNum* res, const mjtNum* mat, const mjtNum* diag, int nr, int nc) {
  mjtNum tmp;

  // half of MatMat routine: only lower triangle
  mju_zero(res, nc*nc);
  if (diag) {
    for (int j=0; j < nr; j++) {
      if (diag[j]) {
        for (int i=0; i < nc; i++) {
          if ((tmp = mat[j*nc+i])) {
            mju_addToScl(res+i*nc, mat+j*nc, tmp*diag[j], i+1);
          }
        }
      }
    }
  } else {
    for (int i=0; i < nc; i++) {
      for (int j=0; j < nr; j++) {
        if ((tmp = mat[j*nc+i])) {
          mju_addToScl(res+i*nc, mat+j*nc, tmp, i+1);
        }
      }
    }
  }

  // make symmetric
  for (int i=0; i < nc; i++) {
    for (int j=i+1; j < nc; j++) {
      res[i*nc+j] = res[j*nc+i];
    }
  }
}



// multiply matrices, first argument transposed
void mju_mulMatTMat(mjtNum* res, const mjtNum* mat1, const mjtNum* mat2,
                    int r1, int c1, int c2) {
  mjtNum tmp;

  mju_zero(res, c1*c2);

  for (int i=0; i < r1; i++) {
    for (int j=0; j < c1; j++) {
      if ((tmp = mat1[i*c1+j])) {
        mju_addToScl(res+j*c2, mat2+i*c2, tmp, c2);
      }
    }
  }
}