/*
 * =======================================================================================
 *
 *      Author:   Jan Eitzinger (je), jan.eitzinger@fau.de
 *      Copyright (c) 2019 RRZE, University Erlangen-Nuremberg
 *
 *      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,
 *      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.
 *
 * =======================================================================================
 */

/* Compile with
 * icx -O3 -xHost -qopenmp -ffast-math -qopt-streaming-stores=always
 * -qopt-zmm-usage=high -qopt-report=3 mfcg.c -o a.out
 */

#define _GNU_SOURCE
#include <float.h>
#include <limits.h>
#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <time.h>
#include <unistd.h>

#define ARRAY_ALIGNMENT 64
#define HLINE                                                                  \
  "--------------------------------------------------------------------------" \
  "--\n"

#ifdef TIMING
#define T(TYPE) timer[TYPE].val +=
#define TIMER_START S = getTimeStamp()
#define TIMER_STOP E = getTimeStamp()
#else
#define T(TYPE)
#define TIMER_START
#define TIMER_STOP
#endif

#define KNRM "\x1B[0m"
#define KRED "\x1B[31m"
#define KGRN "\x1B[32m"

#define CHECK_LESS_THAN(a, b, OP)                                              \
  if (a >= b) {                                                                \
    printf("%sTest %s failed\n", KRED, OP);                                    \
    printf("%s", KNRM);                                                        \
  } else {                                                                     \
    printf("%sTest %s success\n", KGRN, OP);                                   \
    printf("%s", KNRM);                                                        \
  }

typedef enum component {
  AXPBY = 0,
  DOT_PRODUCT,
  APPLY_STENCIL,
  CG_SOLVER,
  NUMTIMERS
} component;

typedef struct {
  char *label;
  double val;
} timerType;

typedef struct {
  double *ptr;
  int size_x; // number of grid points in x dimension
  int size_y; // number of grid points in y dimension
} gridType;

typedef struct {
  double h_x;
  double h_y;
  int niter;
  double tolerance;
} solverType;

/********************/
/* Helper functions */
/********************/
double getTimeStamp() {
  struct timespec ts;
  clock_gettime(CLOCK_MONOTONIC, &ts);
  return (double)ts.tv_sec + (double)ts.tv_nsec * 1.e-9;
}

void allocGrid(int size_x, int size_y, size_t bytesPerWord, gridType *grid) {
  posix_memalign((void **)&(grid->ptr), ARRAY_ALIGNMENT,
                 size_x * size_y * bytesPerWord);
  grid->size_x = size_x;
  grid->size_y = size_y;
}

void copyGrid(gridType *src, gridType *dst) {
  int size = (src->size_x * src->size_y);
  double *restrict S = src->ptr;
  double *restrict D = dst->ptr;

#pragma omp parallel for
  for (int i = 0; i < size; i++) {
    D[i] = S[i];
  }
}

void freeGrid(gridType *grid) { free(grid->ptr); }

void init(gridType *x, gridType *u_sine, gridType *rhs_sine,
          solverType *solver) {
  int size = (x->size_x * x->size_y);
  int size_x = x->size_x;
  int size_y = x->size_y;
  double *P = x->ptr;
  unsigned int seed = 0;

// initialize x
#pragma omp parallel for
  for (int i = 0; i < size; i++) {
    P[i] = i / (double)size;
  }

  P = u_sine->ptr;
// initialize u_sine
#pragma omp parallel for
  for (int j = 0; j < size_y; j++) {
    for (int i = 0; i < size_x; i++) {
      P[j * size_x + i] =
          sin(M_PI * i * solver->h_x) * sin(M_PI * j * solver->h_y);
    }
  }

  P = rhs_sine->ptr;
// initialize rhs_sine
#pragma omp parallel for
  for (int j = 0; j < size_y; j++) {
    for (int i = 0; i < size_x; i++) {
      P[j * size_x + i] = 2.0 * M_PI * M_PI * sin(M_PI * i * solver->h_x) *
                          sin(M_PI * j * solver->h_y);
    }
  }
}

/********************/
/* Basic operations */
/********************/
/* acpby - Calculates res[:] = a*x[:] + b*y[:] */
double axpby(gridType *res, double a, gridType *x, double b, gridType *y) {
  double S = 0.0, E = 0.0;
  double *restrict R =
      (double *)__builtin_assume_aligned(res->ptr, ARRAY_ALIGNMENT);
  double *restrict X =
      (double *)__builtin_assume_aligned(x->ptr, ARRAY_ALIGNMENT);
  double *restrict Y =
      (double *)__builtin_assume_aligned(y->ptr, ARRAY_ALIGNMENT);
  int size = x->size_x * x->size_y;

  TIMER_START;
#pragma omp parallel for
  for (int i = 0; i < size; i++) {
    R[i] = (a * X[i]) + (b * Y[i]);
  }
  TIMER_STOP;

  return E - S;
}

double dotProduct(gridType *x, gridType *y, double *l2sq) {
  double S = 0.0, E = 0.0;
  double l2_sq = 0;
  int size = x->size_x * x->size_y;
  double *restrict X =
      (double *)__builtin_assume_aligned(x->ptr, ARRAY_ALIGNMENT);
  double *restrict Y =
      (double *)__builtin_assume_aligned(y->ptr, ARRAY_ALIGNMENT);

  TIMER_START;
#pragma omp parallel for reduction(+ : l2_sq)
  for (int i = 0; i < size; i++) {
    l2_sq += X[i] * Y[i];
  }
  TIMER_STOP;

  (*l2sq) = l2_sq;
  // printf("DP = %.15lf\n",l2_sq);

  return E - S;
}

/* Applies stencil operation on to u */
/* i.e., res = A*u */
double applyStencil(gridType *res, gridType *u, solverType *solver) {
  double S = 0.0, E = 0.0;
  double *restrict R =
      (double *)__builtin_assume_aligned(res->ptr, ARRAY_ALIGNMENT);
  double *restrict U =
      (double *)__builtin_assume_aligned(u->ptr, ARRAY_ALIGNMENT);
  int size_x = res->size_x;
  int size_y = res->size_y;

  double w_x = 1.0 / (solver->h_x * solver->h_x);
  double w_y = 1.0 / (solver->h_y * solver->h_y);
  double w_c = 2.0 * (w_x + w_y);

  TIMER_START;
#pragma omp parallel for
  for (int j = 1; j < size_y - 1; j++) {
    for (int i = 1; i < size_x - 1; i++) {
      R[j * size_x + i] =
          w_c * U[j * size_x + i] -
          w_y * (U[(j + 1) * size_x + i] + U[(j - 1) * size_x + i]) -
          w_x * (U[j * size_x + i + 1] + U[j * size_x + i - 1]);
    }
  }
  TIMER_STOP;

  return E - S;
}

/* Fused: r[:] = r[:] - lambda*v[:] and returns alpha = r[:].r[:] in one pass.
 * Replaces axpby(r,1,r,-lambda,v) + dotProduct(r,r) -> saves a full re-read of
 * r. */
double axpyDot(gridType *r, double lambda, gridType *v, double *l2sq) {
  double S = 0.0, E = 0.0;
  double alpha = 0.0;
  double *restrict R =
      (double *)__builtin_assume_aligned(r->ptr, ARRAY_ALIGNMENT);
  double *restrict V =
      (double *)__builtin_assume_aligned(v->ptr, ARRAY_ALIGNMENT);
  int size = r->size_x * r->size_y;

  TIMER_START;
#pragma omp parallel for reduction(+ : alpha)
  for (int i = 0; i < size; i++) {
    double ri = R[i] - lambda * V[i];
    R[i] = ri;
    alpha += ri * ri;
  }
  TIMER_STOP;

  (*l2sq) = alpha;
  return E - S;
}

/* Fused: v[:] = A*p[:] (interior stencil) and returns lambda = v[:].p[:] in one
 * pass. Replaces applyStencil(v,p) + dotProduct(v,p) -> saves a re-read of v
 * and p. */
double stencilDot(gridType *res, gridType *u, solverType *solver,
                  double *l2sq) {
  double S = 0.0, E = 0.0;
  double lambda = 0.0;
  double *restrict R =
      (double *)__builtin_assume_aligned(res->ptr, ARRAY_ALIGNMENT);
  double *restrict U =
      (double *)__builtin_assume_aligned(u->ptr, ARRAY_ALIGNMENT);
  int size_x = res->size_x;
  int size_y = res->size_y;

  double w_x = 1.0 / (solver->h_x * solver->h_x);
  double w_y = 1.0 / (solver->h_y * solver->h_y);
  double w_c = 2.0 * (w_x + w_y);

  TIMER_START;
#pragma omp parallel for reduction(+ : lambda)
  for (int j = 1; j < size_y - 1; j++) {
    for (int i = 1; i < size_x - 1; i++) {
      int idx = j * size_x + i;
      double rv = w_c * U[idx] - w_y * (U[idx + size_x] + U[idx - size_x]) -
                  w_x * (U[idx + 1] + U[idx - 1]);
      R[idx] = rv;
      lambda += rv * U[idx];
    }
  }
  TIMER_STOP;

  (*l2sq) = lambda;
  return E - S;
}

/* compute residual */
void computeResidual(gridType *residual, gridType *x, gridType *rhs_sine,
                     solverType *solver, double *res_start) {
  applyStencil(residual, x, solver);
  axpby(residual, 1.0, rhs_sine, -1.0, residual);
  dotProduct(residual, residual, res_start);
}

double CG(gridType *x, gridType *b, solverType *solver, int *iter_end,
          timerType *timer) {
  double S = 0.0, E = 0.0;
  size_t bytesPerWord = sizeof(double);
  gridType p;
  gridType v;
  gridType r;
  int iter = 0;
  double lambda = 0;
  double alpha_0 = 0, alpha_1 = 0;
  double tolSquared = solver->tolerance * solver->tolerance;

  allocGrid(x->size_x, x->size_y, bytesPerWord, &p);
  allocGrid(x->size_x, x->size_y, bytesPerWord, &v);
  allocGrid(x->size_x, x->size_y, bytesPerWord, &r);

  /* NUMA first-touch: initialize the full extent of p, v and r in parallel
   * with the same static distribution used by the compute kernels, so their
   * pages are placed across all NUMA domains. This also zeroes the boundary
   * cells of v/p that applyStencil never writes. */
  {
    int size = x->size_x * x->size_y;
    double *restrict P = p.ptr;
    double *restrict V = v.ptr;
    double *restrict R = r.ptr;
#pragma omp parallel for
    for (int i = 0; i < size; i++) {
      P[i] = 0.0;
      V[i] = 0.0;
      R[i] = 0.0;
    }
  }

  computeResidual(&p, x, b, solver, &alpha_0);
  copyGrid(&p, &r);
  // printf("Before loop: iter = %d, res = %.15e\n", iter, alpha_0);

  S = getTimeStamp();
  while ((iter < solver->niter) && (alpha_0 > tolSquared)) {
    // v = A*p and lambda = v.p in one streaming pass
    T(APPLY_STENCIL) stencilDot(&v, &p, solver, &lambda);
    lambda = alpha_0 / lambda;
    // printf("lambda = %.15e\n", lambda);
    // Update x
    T(AXPBY) axpby(x, 1.0, x, lambda, &p);
    // r = r - lambda*v and alpha_1 = r.r in one streaming pass
    T(AXPBY) axpyDot(&r, lambda, &v, &alpha_1);
    // Update p
    T(AXPBY) axpby(&p, 1.0, &r, alpha_1 / alpha_0, &p);
    alpha_0 = alpha_1;
#ifdef DEBUG
    printf("iter = %d, res = %.15e\n", iter, alpha_0);
#endif
    ++iter;
  }
  E = getTimeStamp();

  freeGrid(&p);
  freeGrid(&v);
  freeGrid(&r);
  (*iter_end) = iter;

  return E - S;
}

int main(int argc, char **argv) {
  size_t bytesPerWord = sizeof(double);
  int size_x = 0;
  int size_y = 0;
  double res_start, err_start;
  double res_sine_cg, err_sine_cg;
  int itermax = 50;
  int iter_sine_cg;
  gridType u_sine, rhs_sine, x, residual;
  timerType timer[NUMTIMERS] = {{"axpby         ", 0.0},
                                {"dot_product   ", 0.0},
                                {"apply_stencil ", 0.0},
                                {"cg            ", 0.0}};

  if (argc > 2) {
    size_x = atoi(argv[1]);
    size_y = atoi(argv[2]);
    if (argc == 4) {
      itermax = atoi(argv[3]);
    }
  } else {
    printf("Usage: %s <outer dimension y> <inner dimension x>\n", argv[0]);
    exit(EXIT_SUCCESS);
  }

  solverType solver = {(double)1.0 / (size_x - 1.0),
                       (double)1.0 / (size_y - 1.0), itermax, 1e-8};

  allocGrid(size_x, size_y, bytesPerWord, &u_sine);
  allocGrid(size_x, size_y, bytesPerWord, &rhs_sine);
  allocGrid(size_x, size_y, bytesPerWord, &x);
  allocGrid(size_x, size_y, bytesPerWord, &residual);

  init(&x, &u_sine, &rhs_sine, &solver);

  /* compute start error and residual */
  axpby(&residual, 1.0, &u_sine, -1.0, &x);
  dotProduct(&residual, &residual, &err_start);
  computeResidual(&residual, &x, &rhs_sine, &solver, &res_start);

  timer[CG_SOLVER].val = CG(&x, &rhs_sine, &solver, &iter_sine_cg, timer);
  printf("CG iterations = %d\n", iter_sine_cg);
  printf("Performance CG = %f [MLUP/s]\n",
         (double)(iter_sine_cg)*size_x * size_y * 1e-6 / timer[CG_SOLVER].val);

  /* compute end error and residual */
  axpby(&residual, 1.0, &u_sine, -1.0, &x);
  dotProduct(&residual, &residual, &err_sine_cg);
  computeResidual(&residual, &x, &rhs_sine, &solver, &res_sine_cg);

  // CHECK_LESS_THAN(res_sine_cg,res_start,"Solver::CG - residual check")
  // CHECK_LESS_THAN(err_sine_cg,err_start,"Solver::CG - error check")

#ifdef TIMING
  printf(HLINE);
  for (int i = 0; i < NUMTIMERS; i++) {
    printf("%s%11.2fs\n", timer[i].label, timer[i].val);
  }
  printf(HLINE);
#endif

#ifdef DEBUG
  printf("Initial residual = %.9e, curr residual CG = %.9e\n", sqrt(res_start),
         sqrt(res_sine_cg));
  printf("Initial error = %.9e, curr error CG = %.9e\n", sqrt(err_start),
         sqrt(err_sine_cg));
#endif
}