#include "sctl.hpp"

// TODO: show how templates improve performance
// TODO: roofline model, arithmetic intensity etc.

template <class Real, long B1, long B2, long B3> void GEMM_Vec(Real* C, const Real* A, const Real* B, const long LDA, const long LDB, const long LDC) {
  using Vec = sctl::Vec<Real,B3>;
  Vec B_[B2], C_[B1];

  //#pragma GCC unroll (10)
  for (long i = 0; i < B1; i++) {
    C_[i] = Vec::Load(C+i*LDC);
  }

  //#pragma GCC unroll (40)
  for (long i = 0; i < B2; i++) {
    B_[i] = Vec::Load(B+i*LDB);
  }

  //#pragma GCC unroll (40)
  for (long j = 0; j <B2; j++) {
    const Real* A_ = A+j;
    //#pragma GCC unroll (10)
    for (long i = 0; i < B1; i++) {
      C_[i] = FMA(A_[i*LDA], B_[j], C_[i]);
    }
  }

  //#pragma GCC unroll (10)
  for (long i = 0; i < B1; i++) {
    C_[i].Store(C+i*LDC);
  }
}

template <class Real, long B1, long B2, long B3> void GEMM_Blocked(Real* C, const Real* A, const Real* B, const long LDA, const long LDB, const long LDC) {
  if (B3 == sctl::DefaultVecLen<Real>())
    return GEMM_Vec<Real, B1, B2, B3>(C, A, B, LDA, LDB, LDC);

  for (long i = 0; i < B1; i++) {
    for (long k = 0; k < B3; k++) {
      for (long j = 0; j < B2; j++) {
        C[i*LDC+k] += A[i*LDA+j] * B[j*LDB+k];
      }
    }
  }
}
template <class Real, long B1, long B2, long B3, long B1_, long B2_, long B3_, long... NN> void GEMM_Blocked(Real* C, const Real* A, const Real* B, const long LDA, const long LDB, const long LDC) {
  static_assert(B1 % B1_ == 0);
  static_assert(B2 % B2_ == 0);
  static_assert(B3 % B3_ == 0);
  for (long i = 0; i < B1; i+=B1_) {
    for (long k = 0; k < B3; k+=B3_) {
      for (long j = 0; j < B2; j+=B2_) {
        GEMM_Blocked<Real, B1_, B2_, B3_, NN...>(C+i*LDC+k, A+i*LDA+j, B+j*LDB+k, LDA, LDB, LDC);
      }
    }
  }
}

template <class Real> void GEMM2(Real* C, const Real* A, const Real* B, const long LDA, const long LDB, const long LDC, const long N1, const long N2, const long N3) {
  for (long i = 0; i < N1; i++) {
    for (long k = 0; k < N3; k++) {
      for (long j = 0; j < N2; j++) {
        C[i*LDC+k] += A[i*LDA+j] * B[j*LDB+k];
      }
    }
  }
}
template <class Real, long B1,long B2,long B3, long... NN> void GEMM2(Real* C, const Real* A, const Real* B, const long LDA, const long LDB, const long LDC, const long N1, const long N2, const long N3) {
  if (N1*N2*N3 == 0) return;
  const long N1_ = (N1/B1)*B1;
  const long N2_ = (N2/B2)*B2;
  const long N3_ = (N3/B3)*B3;
  for (long i = 0; i < N1_; i+=B1) {
    for (long k = 0; k < N3_; k+=B3) {
      for (long j = 0; j < N2_; j+=B2) {
        GEMM_Blocked<Real, B1,B2,B3, NN...>(C+i*LDC+k, A+i*LDA+j, B+j*LDB+k, LDA, LDB, LDC);
      }
    }
  }

  GEMM2<Real, NN...>(C, A+N2_, B+N2_*LDB, LDA, LDB, LDC, N1, N2-N2_, N3);
  GEMM2<Real, NN...>(C+N3_, A, B+N3_, LDA, LDB, LDC, N1_, N2_, N3-N3_);
  GEMM2<Real, NN...>(C+N1_*LDC, A+N1_*LDA, B, LDA, LDB, LDC, N1-N1_, N2_, N3_);
  GEMM2<Real, NN...>(C+N1_*LDC+N3_, A+N1_*LDA, B+N3_, LDA, LDB, LDC, N1-N1_, N2_, N3-N3_);
}

template <class Real> void myGEMM(sctl::Matrix<Real>& C, const sctl::Matrix<Real>& A, const sctl::Matrix<Real>& B) {
  const long N1 = A.Dim(0);
  const long N2 = A.Dim(1);
  const long N3 = B.Dim(1);
  SCTL_ASSERT(B.Dim(0) == N2);
  sctl::Profile::Add_FLOP(2*N1*N2*N3);
  if (C.Dim(0) != N1 || C.Dim(1) != N3) C.ReInit(N1, N3);

  // no-blocking, no-vectorizing
  //GEMM2<Real>(C.begin(), A.begin(), B.begin(), N2, N3, N3, N1, N2, N3);

  // blocking + vectorizing
  //GEMM2<Real, 200,200,200, 40,40,40, 10,40,8>(C.begin(), A.begin(), B.begin(), N2, N3, N3, N1, N2, N3); // double
  //GEMM2<Real, 200,200,320, 40,40,80, 8,40,16>(C.begin(), A.begin(), B.begin(), N2, N3, N3, N1, N2, N3); // float


  //GEMM2<Real, 200,200,200, 40,40,40, 10,40,8>(C.begin(), A.begin(), B.begin(), N2, N3, N3, N1, N2, N3);
  //GEMM2<Real, 160,240,200, 16,40,40, 16,40,8>(C.begin(), A.begin(), B.begin(), N2, N3, N3, N1, N2, N3);
  //GEMM2<Real, 160,200,200, 16,40,40, 16,40,8>(C.begin(), A.begin(), B.begin(), N2, N3, N3, N1, N2, N3);
  //GEMM2<Real, 120,120,120, 12,60,8>(C.begin(), A.begin(), B.begin(), N2, N3, N3, N1, N2, N3);

  GEMM2<Real, 200,200,200, 40,40,40, 10,40,4>(C.begin(), A.begin(), B.begin(), N2, N3, N3, N1, N2, N3); ///////////////////////////
}

template <class Real> void TestMatrix() {
  constexpr long N1 = 200; //3000;
  constexpr long N2 = 200; //3000;
  constexpr long N3 = 200; //3000;

  sctl::Profile::Tic("TestMatrix");
  sctl::Matrix<Real> A(N1, N2);
  sctl::Matrix<Real> B(N2, N3);
  sctl::Matrix<Real> C(N1, N3);
  sctl::Matrix<Real> C_ref(N1, N3);

  sctl::Profile::Tic("Init");
  for (long i = 0; i < N1*N2; i++) A[0][i] = drand48();
  for (long i = 0; i < N2*N3; i++) B[0][i] = drand48();
  for (long i = 0; i < N1*N3; i++) C[0][i] = 0;
  for (long i = 0; i < N1*N3; i++) C_ref[0][i] = 0;
  sctl::Profile::Toc();

  //sctl::Profile::Tic("GEMM");
  //for (long i = 0; i < 2e10/N1/N2/N3; i++) sctl::Matrix<Real>::GEMM(C_ref, A, B, 1);
  //sctl::Profile::Toc();

  sctl::Profile::Tic("myGEMM");
  for (long i = 0; i < 2e10/N1/N2/N3; i++) myGEMM(C, A, B);
  sctl::Profile::Toc();

  //sctl::Profile::Tic("myGEMM");
  //for (long i = 0; i < 2e10/N1/N2/N3; i++) {
  //  //myGEMM(C, A, B);
  //  //GEMM_Vec<Real,N1,N2,N3>(C[0], A[0], B[0], N2, N3, N3);
  //  GEMM2<Real, N1,N2,N3, 320,200,200, 16,40,40, 16,40,8>(C[0], A[0], B[0], N2, N3, N3, N1, N2, N3);
  //}
  //sctl::Profile::Add_FLOP(ceil(2e10/N1/N2/N3) * 2*N1*N2*N3);
  //sctl::Profile::Toc();

  Real max_err = 0, max_val = 0;
  for (long i = 0; i < C.Dim(0) * C.Dim(1); i++) max_err = std::max<Real>(max_err, fabs(C[0][i] - C_ref[0][i]));
  for (long i = 0; i < C.Dim(0) * C.Dim(1); i++) max_val = std::max<Real>(max_val, fabs(C_ref[0][i]));
  std::cout<<max_err/max_val<<'\n';

  sctl::Profile::Toc();
}


void Transpose(double* A, long N1, long N2, double* buff) {
  double* B = (buff ? buff : (double*)malloc(N1*N2*sizeof(double)));
  for (long i = 0; i < N2; i++) { // TODO: loop reorder effect
    for (long j = 0; j < N1; j++) {
      B[i*N1+j] = A[j*N2+i];
    }
  }
  for (long i = 0; i < N1*N2; i++) {
    A[i] = B[i];
  }
  if (!buff) free(B);
}
void BlockedTranspose(double* A, long N1, long N2, double* buff) { // not faster
  double* B = (buff ? buff : (double*)malloc(N1*N2*sizeof(double)));
  constexpr long BLOCK_SIZE = 40;
  const long N1_ = (N1/BLOCK_SIZE)*BLOCK_SIZE;
  const long N2_ = (N2/BLOCK_SIZE)*BLOCK_SIZE;
  for (long i0 = 0; i0 < N2_; i0+=BLOCK_SIZE) {
    for (long j0 = 0; j0 < N1_; j0+=BLOCK_SIZE) {
      const double* A_ = A + j0*N2 + i0;
      for (long i = 0; i < BLOCK_SIZE; i++) {
        double* B_ = B + (i0+i)*N1+j0;
        for (long j = 0; j < BLOCK_SIZE; j++) {
          B_[j] = A_[j*N2+i];
        }
      }
    }

    const double* A_ = A + i0;
    for (long i = 0; i < BLOCK_SIZE; i++) {
      double* B_ = B + (i0+i)*N1;
      for (long j = N1_; j < N1; j++) {
        B_[j] = A_[j*N2+i];
      }
    }
  }
  for (long i = N2_; i < N2; i++) {
    for (long j = 0; j < N1; j++) {
      B[i*N1+j] = A[j*N2+i];
    }
  }
  for (long i = 0; i < N1*N2; i++) {
    A[i] = B[i];
  }
  if (!buff) free(B);
}
void ProfileTranspose() {
  long np = omp_get_max_threads();
  long N1 = 1024*3+1, N2 = 1024*3+1;
  double* A = new double[np*N1*N2];
  double* B = new double[np*N1*N2];
  double tt;

  #pragma omp parallel for schedule(static)
  for (long i = 0; i < np; i++) {
    for (long j = 0; j < N1*N2; j++) {
      A[i*N1*N2+j] = j;
      B[i*N1*N2+j] = 0;
    }
  }

  tt = -omp_get_wtime();
  #pragma omp parallel for schedule(static)
  for (long i = 0; i < np; i++) {
    long N1_ = N1;
    long N2_ = N2;
    for (long j = 0; j < 100; j++) {
      Transpose(A+i*N1*N2, N1_, N2_, B+i*N1*N2);
      long N_ = N1_;
      N1_ = N2_;
      N2_ = N_;
    }
  }
  tt += omp_get_wtime();
  std::cout<<"Bandwidth (no-blocking, no-alloc) : "<<N1*N2*np*8*3*100/tt/1e9<<'\n';

  tt = -omp_get_wtime();
  #pragma omp parallel for schedule(static)
  for (long i = 0; i < np; i++) {
    long N1_ = N1;
    long N2_ = N2;
    for (long j = 0; j < 100; j++) {
      Transpose(A+i*N1*N2, N1_, N2_, nullptr);
      long N_ = N1_;
      N1_ = N2_;
      N2_ = N_;
    }
  }
  tt += omp_get_wtime();
  std::cout<<"Bandwidth (no-blocking, alloc) : "<<N1*N2*np*8*3*100/tt/1e9<<'\n';

  tt = -omp_get_wtime();
  #pragma omp parallel for schedule(static)
  for (long i = 0; i < np; i++) {
    long N1_ = N1;
    long N2_ = N2;
    for (long j = 0; j < 100; j++) {
      BlockedTranspose(A+i*N1*N2, N1_, N2_, B+i*N1*N2);
      long N_ = N1_;
      N1_ = N2_;
      N2_ = N_;
    }
  }
  tt += omp_get_wtime();
  std::cout<<"Bandwidth (blocking, no-alloc) : "<<N1*N2*np*8*3*100/tt/1e9<<'\n';

  tt = -omp_get_wtime();
  #pragma omp parallel for schedule(static)
  for (long i = 0; i < np; i++) {
    long N1_ = N1;
    long N2_ = N2;
    for (long j = 0; j < 100; j++) {
      BlockedTranspose(A+i*N1*N2, N1_, N2_, nullptr);
      long N_ = N1_;
      N1_ = N2_;
      N2_ = N_;
    }
  }
  tt += omp_get_wtime();
  std::cout<<"Bandwidth (blocking, alloc) : "<<N1*N2*np*8*3*100/tt/1e9<<'\n';

  #pragma omp parallel for schedule(static)
  for (long i = 0; i < np; i++) {
    for (long j = 0; j < N1*N2; j++) {
      SCTL_ASSERT(A[i*N1*N2+j] == j);
    }
  }

  delete[] A;
  delete[] B;
}

int main(int argc, char** argv) {
  sctl::Comm::MPI_Init(&argc, &argv);
  sctl::Profile::Enable(true);

  TestMatrix<double>();


  //ProfileTranspose();

  //ProfileMemgr();
  //ProfileMemgr();

  // Print profiling results
  sctl::Profile::print();
  sctl::Comm::MPI_Finalize();
  return 0;
}
