#include "light_matrix.h"
#include "helpler_funtions.h"
#include <math.h>



/************************************************************************/
/*                          Private Function                            */
/************************************************************************/

static void swap(int *a, int *b) {
  int m;
  m = *a;
  *a = *b;
  *b = m;
}

static void perm(int list[], int k, int m, int *p, Mat *mat, MAT_TYPE *det) {
  int i;

  if (k > m) {
    MAT_TYPE res = mat->element[0][list[0]];

    for (i = 1; i < mat->row; i++) {
      res *= mat->element[i][list[i]];
    }

    if (*p % 2) {
      // odd is negative
      *det -= res;
    } else {
      // even is positive
      *det += res;
    }
  } else {
    // if the element is 0, we don't need to calculate the value for this
    // permutation
    //		if(!equal(mat->element[k][list[k]], 0.0f))
    //			perm(list, k + 1, m, p, mat, det);
    perm(list, k + 1, m, p, mat, det);

    for (i = k + 1; i <= m; i++) {
      //			if(equal(mat->element[k][list[i]], 0.0f))
      //				continue;
      swap(&list[k], &list[i]);
      *p += 1;
      perm(list, k + 1, m, p, mat, det);
      swap(&list[k], &list[i]);
      *p -= 1;
    }
  }
}

/************************************************************************/
/*                           Public Function                            */
/************************************************************************/

/**
 * @brief creat a matrix instance by using dynamic memory
 *
 * @param row rows number of matrix
 * @param col cols number of matrix
 * @return Mat* pointer to the new matrix instance
 */
Mat *MatNew(int row, int col) {
  /* make sure row and column num is positive */
  MAT_ASSERT(row > 0);
  MAT_ASSERT(col > 0);

  Mat *mat = NULL;
  mat = (Mat *)MAT_MALLOC(sizeof(Mat));

  if (mat) {
    mat->element = (MAT_TYPE **)MAT_MALLOC(row * sizeof(MAT_TYPE *));
    mat->buffer = (MAT_TYPE *)MAT_MALLOC(row * col * sizeof(MAT_TYPE));
  }

  if (mat->element == NULL || mat->buffer == NULL) {
    MatDelete(mat);
    return NULL;
  } else {
    for (int i = 0; i < row; i++) {
      mat->element[i] = &mat->buffer[i * col];
    }
    mat->row = row;
    mat->col = col;
  }

  return mat;
}

/**
 * @brief delete a matrix instance created by using dynamic memory
 *
 * @param mat pointer to the matrix instance
 */
void MatDelete(Mat *mat) {
  if (mat == NULL) {
    return;
  }

  if (mat->element) {
    MAT_FREE(mat->element);
  }
  if (mat->buffer) {
    MAT_FREE(mat->buffer);
  }
  MAT_FREE(mat);
}

/**
 * @brief intialise a matrix, don't need dynamic memory allocation
 *
 * @param mat pointer to matrix instance
 * @param row matrix rows number
 * @param col matrix cols number
 * @param pdata matrix data buffer
 */
void MatInit(Mat *mat, int row, int col, MAT_TYPE *pdata) {
  MAT_ASSERT(mat != NULL);
  MAT_ASSERT(pdata != NULL);
  MAT_ASSERT((row > 0 && col > 0));

  mat->row = row;
  mat->col = col;

  mat->buffer = pdata;

  for (int i = 0; i < row; i++) {
    mat->element[i] = &mat->buffer[i * col];
  }
}

/**
 * @brief swap rows of the matrix
 *
 * @param mat matrix pointer
 * @param i row number
 * @param j row number
 */
void MatSwapRows(Mat *mat, size_t i, size_t j) {
  MAT_ASSERT(mat != NULL);
  MAT_ASSERT(i < mat->row);
  MAT_ASSERT(j < mat->row);

  if (i == j) {
    return;
  }

  if (mat->element) {
    for (size_t k = 0; k < mat->col; k++) {
      MAT_TYPE tmp = mat->element[i][k];
      mat->element[i][k] = mat->element[j][k];
      mat->element[j][k] = tmp;
    }
  }
}

/**
 * @brief swap cols of the matrix
 *
 * @param mat matrix pointer
 * @param i col number
 * @param j col number
 */
void MatSwapCols(Mat *mat, size_t i, size_t j) {
  MAT_ASSERT(mat != NULL);
  MAT_ASSERT(i < mat->col);
  MAT_ASSERT(j < mat->col);

  if (i == j) {
    return;
  }

  if (mat->element) {
    for (size_t k = 0; k < mat->row; k++) {
      MAT_TYPE tmp = mat->element[k][i];
      mat->element[k][i] = mat->element[k][j];
      mat->element[k][j] = tmp;
    }
  }
}

/**
 * @brief set matrix element value
 *
 * @param mat matrix instance
 * @param val value buffer
 * @return error status
 */
int MatSetVal(Mat *mat, MAT_TYPE *val) {
  int ret = MAT_EOK;
  MAT_ASSERT(mat != NULL);
  MAT_ASSERT(val != NULL);

  int row, col;

  if (mat->buffer == NULL || mat->element == NULL) {
    ret = -MAT_ERROR;
  }

  if (MAT_EOK == ret) {
    for (row = 0; row < mat->row; row++) {
      for (col = 0; col < mat->col; col++) {
        mat->element[row][col] = val[col + row * mat->col];
      }
    }
  }

  return ret;
}

/**
 * @brief dump the matrix data to a file to std out
 *
 * @param mat matrix pointer
 */
void MatDump(const Mat *mat) {
  int row, col;

  if (mat == NULL) {
    return;
  }

  MAT_PRINTF("Mat %zux%zu:\n", mat->row, mat->col);

  for (row = 0; row < mat->row; row++) {
    for (col = 0; col < mat->col; col++) {
      MAT_PRINTF("%.4f\t", mat->element[row][col]);
    }
    MAT_PRINTF("\n");
  }
}

/**
 * @brief set all the element of the matrix to be zero
 *
 * @param mat matrix pointer
 * @return int MAT_ERROR code
 */
int MatZeros(Mat *mat) {
  MAT_ASSERT(mat != NULL);

  int ret = MAT_EOK;
  int row, col;

  if (mat->buffer == NULL || mat->element == NULL) {
    ret = -MAT_ERROR;
  }

  if (MAT_EOK == ret) {
    for (row = 0; row < mat->row; row++) {
      for (col = 0; col < mat->col; col++) {
        mat->element[row][col] = 0;
      }
    }
  }

  return ret;
}

/**
 * @brief set the matrix to be a identity matrix
 *
 * @param mat
 * @return int
 */
int MatEye(Mat *mat) {
  MAT_ASSERT(mat != NULL);

  int ret = MAT_EOK;

  ret = MatZeros(mat);

  if (MAT_EOK == ret) {
    for (int i = 0; i < min(mat->row, mat->col); i++) {
      mat->element[i][i] = 1.0f;
    }
  }

  return ret;
}

/**
 * @brief add two matrix by element
 *
 * @param src1 matrix to be added
 * @param src2 matrix to be added
 * @param dst dst = src1 + src2
 * @return int MAT_ERROR code
 */
int MatAdd(Mat *src1, Mat *src2, Mat *dst) {
  MAT_ASSERT(src1 != NULL);
  MAT_ASSERT(src2 != NULL);
  MAT_ASSERT(dst != NULL);

  int ret = MAT_EOK;

  /* the two matrixs's column and row should be the same */
  if (!(src1->row == src2->row && src2->row == dst->row &&
        src1->col == src2->col && src2->col == dst->col)) {
    ret = -MAT_ERROR;
  }

  if (MAT_EOK == ret) {
    for (size_t row = 0; row < src1->row; row++) {
      for (size_t col = 0; col < src1->col; col++) {
        dst->element[row][col] =
            src1->element[row][col] + src2->element[row][col];
      }
    }
  }

  return ret;
}

/* dst = src1 - src2 */
int MatSub(Mat *src1, Mat *src2, Mat *dst) {
  MAT_ASSERT(src1 != NULL);
  MAT_ASSERT(src2 != NULL);
  MAT_ASSERT(dst != NULL);

  int ret = MAT_EOK;

  if (!(src1->row == src2->row && src2->row == dst->row &&
        src1->col == src2->col && src2->col == dst->col)) {
    ret = -MAT_ERROR;
  }

  if (MAT_EOK == ret) {
    for (int row = 0; row < src1->row; row++) {
      for (int col = 0; col < src1->col; col++) {
        dst->element[row][col] =
            src1->element[row][col] - src2->element[row][col];
      }
    }
  }

  return ret;
}

/* dst = src1 * src2 */
int MatMul(Mat *src1, Mat *src2, Mat *dst) {
  MAT_ASSERT(src1 != NULL);
  MAT_ASSERT(src2 != NULL);
  MAT_ASSERT(dst != NULL);

  int ret = MAT_EOK;

  if (src1->col != src2->row || src1->row != dst->row ||
      src2->col != dst->col) {
    ret = -MAT_ERROR;
  }

  if (MAT_EOK == ret) {
    for (int row = 0; row < dst->row; row++) {
      for (int col = 0; col < dst->col; col++) {
        MAT_TYPE temp = 0.0f;

        for (int i = 0; i < src1->col; i++) {
          temp += src1->element[row][i] * src2->element[i][col];
        }

        dst->element[row][col] = temp;
      }
    }
  }

  return ret;
}

/* dst = src' */
int MatTrans(Mat *src, Mat *dst) {
  MAT_ASSERT(src != NULL);
  MAT_ASSERT(dst != NULL);

  int rslt = MAT_EOK;

  if (src->row != dst->col || src->col != dst->row) {
    rslt = -MAT_ERROR;
  }

  if (MAT_EOK == rslt) {
    for (int row = 0; row < dst->row; row++) {
      for (int col = 0; col < dst->col; col++) {
        dst->element[row][col] = src->element[col][row];
      }
    }
  }

  return rslt;
}

/**
 * @brief calculate determinant of the matrix
 *
 * @param mat matrix instance
 * @return the determinant
 */
MAT_TYPE MatDet(Mat *mat) {
  MAT_ASSERT(mat != NULL);

  int rslt = MAT_EOK;
  MAT_TYPE det = 0.0f;
  int plarity = 0;
  int *list;
  int i;

  /* the mat should be a square matrix */
  if (mat->row != mat->col) {
    rslt = -MAT_ERROR;
  }

  if (MAT_EOK == rslt) {
    list = (int *)MAT_MALLOC(sizeof(int) * mat->col);
    if (list == NULL) {
      rslt = -MAT_ENOMEM;
    } else {
      for (i = 0; i < mat->col; i++)
        list[i] = i;

      perm(list, 0, mat->row - 1, &plarity, mat, &det);
      MAT_FREE(list);
    }
  }

  return det;
}

/* dst = adj(src) */
int MatAdj(Mat *src, Mat *dst) {
  int rslt = MAT_EOK;

  MAT_TYPE det = 0;

  if (src->row != src->col || src->row != dst->row || src->col != dst->col) {
    rslt = -MAT_ERROR;
  }

  if (MAT_EOK == rslt) {
    Mat *smat = NULL;
    smat = MatNew(src->row - 1, src->col - 1);

    if (smat) {
      for (int row = 0; row < src->row; row++) {
        for (int col = 0; col < src->col; col++) {
          int r = 0;

          for (int i = 0; i < src->row; i++) {
            if (i == row)
              continue;

            int c = 0;

            for (int j = 0; j < src->col; j++) {
              if (j == col)
                continue;

              smat->element[r][c] = src->element[i][j];
              c++;
            }

            r++;
          }

          det = MatDet(smat);

          if ((row + col) % 2)
            det = -det;

          dst->element[col][row] = det;
        }
      }
      MatDelete(smat);
    } else {
      rslt = -MAT_ENOMEM;
    }
  }

  return rslt;
}

int MatCopy(Mat *src, Mat *dst) {
  int ret = MAT_EOK;

  if (src->row != dst->row || src->col != dst->col) {
    ret = -MAT_ERROR;
  }

  if (MAT_EOK == ret) {
    for (int row = 0; row < src->row; row++) {
      for (int col = 0; col < src->col; col++)
        dst->element[row][col] = src->element[row][col];
    }
  }

  return ret;
}

int MatEig(Mat *mat, MAT_TYPE *eig_val, Mat *eig_vec, MAT_TYPE eps, int njt) {
  MAT_ASSERT(mat != NULL);
  MAT_ASSERT(eig_val != NULL);
  MAT_ASSERT(eig_vec != NULL);

  int rslt = MAT_EOK;

  int i, j;
  int nDim = mat->row;

  /* mat should be a squre matrix */
  if (mat->row != mat->col) {
    rslt = -MAT_ERROR;
  }

  Mat *temp_mat = NULL;
  if (MAT_EOK == rslt) {
    temp_mat = MatNew(mat->row, mat->col);

    if (temp_mat != NULL) {
      MatCopy(mat, temp_mat);

      for (i = 0; i < nDim; i++) {
        eig_vec->element[i][i] = 1.0f;

        for (int j = 0; j < nDim; j++) {
          if (i != j)
            eig_vec->element[i][j] = 0.0f;
        }
      }

      int nCount = 0; // iteration count

      while (1) {
        // find maxial element in non-diagram of mat
        MAT_TYPE dbMax = temp_mat->element[0][1];
        int nRow = 0;
        int nCol = 1;

        for (i = 0; i < nDim; i++) {   // row
          for (j = 0; j < nDim; j++) { // col
            MAT_TYPE d = fabs(temp_mat->element[i][j]);

            if ((i != j) && (d > dbMax)) {
              dbMax = d;
              nRow = i;
              nCol = j;
            }
          }
        }

        if (dbMax < eps) {
          break;
        }

        if (nCount > njt) {
          break;
        }

        nCount++;

        MAT_TYPE dbApp = temp_mat->element[nRow][nRow];
        MAT_TYPE dbApq = temp_mat->element[nRow][nCol];
        MAT_TYPE dbAqq = temp_mat->element[nCol][nCol];

        // calculate rotation angle
        MAT_TYPE dbAngle = 0.5 * atan2(-2 * dbApq, dbAqq - dbApp);
        MAT_TYPE dbSinTheta = sin(dbAngle);
        MAT_TYPE dbCosTheta = cos(dbAngle);
        MAT_TYPE dbSin2Theta = sin(2 * dbAngle);
        MAT_TYPE dbCos2Theta = cos(2 * dbAngle);

        temp_mat->element[nRow][nRow] = dbApp * dbCosTheta * dbCosTheta +
                                        dbAqq * dbSinTheta * dbSinTheta +
                                        2 * dbApq * dbCosTheta * dbSinTheta;
        temp_mat->element[nCol][nCol] = dbApp * dbSinTheta * dbSinTheta +
                                        dbAqq * dbCosTheta * dbCosTheta -
                                        2 * dbApq * dbCosTheta * dbSinTheta;
        temp_mat->element[nRow][nCol] =
            0.5f * (dbAqq - dbApp) * dbSin2Theta + dbApq * dbCos2Theta;
        temp_mat->element[nCol][nRow] = temp_mat->element[nRow][nCol];

        for (i = 0; i < nDim; i++) {
          if ((i != nCol) && (i != nRow)) {
            dbMax = temp_mat->element[i][nRow];
            temp_mat->element[i][nRow] =
                temp_mat->element[i][nCol] * dbSinTheta + dbMax * dbCosTheta;
            temp_mat->element[i][nCol] =
                temp_mat->element[i][nCol] * dbCosTheta - dbMax * dbSinTheta;
          }
        }

        for (j = 0; j < nDim; j++) {
          if ((j != nCol) && (j != nRow)) {
            dbMax = temp_mat->element[nRow][j];
            temp_mat->element[nRow][j] =
                temp_mat->element[nCol][j] * dbSinTheta + dbMax * dbCosTheta;
            temp_mat->element[nCol][j] =
                temp_mat->element[nCol][j] * dbCosTheta - dbMax * dbSinTheta;
          }
        }

        // calculate eigen vector
        for (i = 0; i < nDim; i++) {
          dbMax = eig_vec->element[i][nRow];
          eig_vec->element[i][nRow] =
              eig_vec->element[i][nCol] * dbSinTheta + dbMax * dbCosTheta;
          eig_vec->element[i][nCol] =
              eig_vec->element[i][nCol] * dbCosTheta - dbMax * dbSinTheta;
        }
      }

      // calculate eigen value
      for (i = 0; i < nDim; i++) {
        eig_val[i] = temp_mat->element[i][i];
      }

      // set sign
      for (i = 0; i < nDim; i++) {
        MAT_TYPE dSumVec = 0;

        for (j = 0; j < nDim; j++)
          dSumVec += eig_vec->element[j][i];

        if (dSumVec < 0) {
          for (j = 0; j < nDim; j++)
            eig_vec->element[j][i] *= -1;
        }
      }
      MatDelete(temp_mat);
    } else {
      rslt = -MAT_ENOMEM;
    }
  }

  return rslt;
}

MAT_TYPE MatNorm(Mat *mat) {
  MAT_TYPE max_eig_val = 0.0f;
  int rslt = MAT_EOK;

  /* mat should be a squre matrix */
  if (mat->row != mat->col) {
    rslt = -MAT_ERROR;
  }

  if (MAT_EOK == rslt) {
    MAT_TYPE *eig_val = NULL;
    Mat *eig_vec = NULL;
    eig_val = (MAT_TYPE *)MAT_MALLOC(mat->row * sizeof(MAT_TYPE));
    eig_vec = MatNew(mat->row, mat->col);

    if (eig_val != NULL && eig_vec != NULL) {
      MatEig(mat, eig_val, eig_vec, 1e-6, 100);

      for (int i = 0; i < mat->row; i++) {
        if (eig_val[i] > max_eig_val)
          max_eig_val = eig_val[i];
      }
    } else {
      rslt = -MAT_ENOMEM;
    }

    MAT_FREE(eig_val);
    MatDelete(eig_vec);
  }

  return max_eig_val;
}

int MatInvByLu(Mat *const A, Mat *inv) {
  MAT_ASSERT(A != NULL);
  MAT_ASSERT(inv != NULL);

  size_t N = A->col;
  Mat *L = NULL, *U = NULL, *P = NULL;
  int ret = MAT_EOK;

  if (A->col != inv->col || A->row != inv->row) {
    ret = -MAT_ERROR;
  }

  if (ret == MAT_EOK) {
    L = MatNew(N, N);
    U = MatNew(N, N);
    P = MatNew(N, N);

    if (L == NULL || U == NULL || P == NULL) {
      ret = -MAT_ENOMEM;
    }
  }

  if (ret == MAT_EOK) {
    MatEye(L);
    MatEye(P);
    MatCopy(A, U);

    for (size_t n = 0; n < N; n++) {
      if (fabs(U->element[n][n]) < (MAT_TYPE)FLT_EPSILON) {
        for (size_t i = n + 1; i < N; i++) {
          if (fabs(U->element[i][n]) > (MAT_TYPE)FLT_EPSILON) {
            MatSwapRows(U, i, n);
            MatSwapRows(P, i, n);
            MatSwapRows(L, i, n);
            MatSwapCols(L, i, n);
          }
        }
      }

      if ((float)(fabs(U->element[n][n])) < FLT_EPSILON) {
        ret = -MAT_ERROR;
        break;
      }

      for (size_t i = n + 1; i < N; i++) {
        L->element[i][n] = U->element[i][n] / U->element[n][n];
        for (size_t k = 0; k < N; k++) {
          U->element[i][k] -= L->element[i][n] * U->element[n][k];
        }
      }
    }

    if (ret == MAT_EOK) {
      for (size_t col = 0; col < N; col++) {
        for (size_t i = 0; i < N; i++) {
          for (size_t j = 0; j < i; j++) {
            P->element[i][col] -= L->element[i][j] * P->element[j][col];
          }
        }
      }

      for (size_t col = 0; col < N; col++) {
        for (size_t k = 0; k < N; k++) {
          size_t i = N - k - 1;
          for (size_t j = i + 1; j < N; j++) {
            P->element[i][col] -= U->element[i][j] * P->element[j][col];
          }
          P->element[i][col] /= U->element[i][i];
        }
      }

      for (size_t i = 0; i < N; i++) {
        for (size_t j = 0; j < N; j++) {
          if (!isfinite(P->element[i][j])) {
            ret = -MAT_ERROR;
          }
        }
      }
    }
  }

  if (ret == MAT_EOK) {
    MatCopy(P, inv);
  }

  MatDelete(L);
  MatDelete(U);
  MatDelete(P);

  return ret;
}

// dst = src^(-1)
int MatInv(Mat *src, Mat *dst) {

  int ret = MAT_EOK;

  MAT_TYPE det;
  int row, col;

  /* the mat should be a square matrix */
  if (src->row != src->col || src->row != dst->row || src->col != dst->col) {
    ret = -MAT_EINVAL;
  }

  if (ret == MAT_EOK) {
    Mat *adj_mat = MatNew(src->row, src->col);
    if (adj_mat) {
      if (MatAdj(src, adj_mat) == MAT_EOK) {
        det = MatDet(src);
        if (!is_zerof(det)) {
          for (row = 0; row < src->row; row++) {
            for (col = 0; col < src->col; col++)
              dst->element[row][col] = adj_mat->element[row][col] / det;
          }
          ret = MAT_EOK;
        } else {
          ret = -MAT_EINVAL;
        }
      } else {
        ret = -MAT_ERROR;
      }
      MatDelete(adj_mat);
    } else {
      ret = -MAT_ENOMEM;
    }
  }

  return ret;
}