#include <mpi.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <gsl/gsl_math.h>
#include <sys/stat.h>
#include <sys/types.h>

#include "allvars.h"
#include "proto.h"

#ifdef SUBFIND

#include "fof.h"
#include "subfind.h"

/*! Structure for communication during the density computation. Holds data that is sent to other processors.
 */
static struct densdata_in
{
  MyDouble Pos[3];
  MyFloat Hsml;
  int NodeList[NODELISTLENGTH];
}
 *DensDataIn, *DensDataGet;

static struct densdata_out
{
  MyFloat Rho;
  MyFloat VelDisp, Vx, Vy, Vz;
  int Ngb;
}
 *DensDataResult, *DensDataOut;

static MyFloat *DM_Vx, *DM_Vy, *DM_Vz;
static long long Ntotal;

#ifdef DENSITY_SPLIT_BY_TYPE
void subfind_density(int j_in)
#else
void subfind_density(void)
#endif
{
  long long ntot;
  int i, j, ndone, ndone_flag, npleft, dummy, iter = 0;
  MyFloat *Left, *Right;
  char *Todo;
  int ngrp, sendTask, recvTask, place, nexport, nimport;
  double vel_to_phys, dmax1, dmax2, t0, t1;

#ifdef DENSITY_SPLIT_BY_TYPE
  int j_target;

  if(j_in < 0)
    j_target = -j_in - 1;
  else
    j_target = j_in;
#endif

  if(ThisTask == 0)
    {
#ifdef DENSITY_SPLIT_BY_TYPE
      printf("finding densities, mode %d \n", j_in);
#else
      printf("finding densities for all particles\n");
#endif
      fflush(stdout);
    }

  /* allocate buffers to arrange communication */

  Ngblist = (int *) mymalloc(NumPart * sizeof(int));
  Dist2list = (double *) mymalloc(NumPart * sizeof(double));

  All.BunchSize =
    (int) ((All.BufferSize * 1024 * 1024) / (sizeof(struct data_index) + sizeof(struct data_nodelist) +
					     sizeof(struct densdata_in) + sizeof(struct densdata_out) +
					     sizemax(sizeof(struct densdata_in),
						     sizeof(struct densdata_out))));
  DataIndexTable = (struct data_index *) mymalloc(All.BunchSize * sizeof(struct data_index));
  DataNodeList = (struct data_nodelist *) mymalloc(All.BunchSize * sizeof(struct data_nodelist));

  Left = mymalloc(sizeof(MyFloat) * NumPart);
  Right = mymalloc(sizeof(MyFloat) * NumPart);
  Todo = mymalloc(sizeof(char) * NumPart);

  DM_Vx = mymalloc(sizeof(MyFloat) * NumPart);
  DM_Vy = mymalloc(sizeof(MyFloat) * NumPart);
  DM_Vz = mymalloc(sizeof(MyFloat) * NumPart);

  for(i = 0; i < NumPart; i++)
    {
      Left[i] = Right[i] = 0;
      P[i].DM_NumNgb = 0;
      Todo[i] = 1;
    }

  /* we will repeat the whole thing for those particles where we didn't find enough neighbours */
  do
    {
      t0 = second();

      i = 0;			/* begin with this index */

      do
	{
	  for(j = 0; j < NTask; j++)
	    {
	      Send_count[j] = 0;
	      Exportflag[j] = -1;
	    }

	  /* do local particles and prepare export list */

	  for(nexport = 0; i < NumPart; i++)
	    {
	      if(Todo[i])
#ifdef DENSITY_SPLIT_BY_TYPE
		if(P[i].Type == j_target)
		  {
		    if(subfind_density_evaluate(i, 0, &nexport, Send_count, j_in) < 0)
		      break;
		  }
#else
		if(((1 << P[i].Type) & (FOF_PRIMARY_LINK_TYPES)))
		  {
		    if(subfind_density_evaluate(i, 0, &nexport, Send_count) < 0)
		      break;
		  }
#endif
	    }

	  qsort(DataIndexTable, nexport, sizeof(struct data_index), data_index_compare);

	  MPI_Allgather(Send_count, NTask, MPI_INT, Sendcount_matrix, NTask, MPI_INT, MPI_COMM_WORLD);

	  for(j = 0, nimport = 0, Recv_offset[0] = 0, Send_offset[0] = 0; j < NTask; j++)
	    {
	      Recv_count[j] = Sendcount_matrix[j * NTask + ThisTask];
	      nimport += Recv_count[j];

	      if(j > 0)
		{
		  Send_offset[j] = Send_offset[j - 1] + Send_count[j - 1];
		  Recv_offset[j] = Recv_offset[j - 1] + Recv_count[j - 1];
		}
	    }

	  DensDataGet = (struct densdata_in *) mymalloc(nimport * sizeof(struct densdata_in));
	  DensDataIn = (struct densdata_in *) mymalloc(nexport * sizeof(struct densdata_in));

	  /* prepare particle data for export */
	  for(j = 0; j < nexport; j++)
	    {
	      place = DataIndexTable[j].Index;

	      DensDataIn[j].Pos[0] = P[place].Pos[0];
	      DensDataIn[j].Pos[1] = P[place].Pos[1];
	      DensDataIn[j].Pos[2] = P[place].Pos[2];
	      DensDataIn[j].Hsml = P[place].DM_Hsml;

	      memcpy(DensDataIn[j].NodeList,
		     DataNodeList[DataIndexTable[j].IndexGet].NodeList, NODELISTLENGTH * sizeof(int));
	    }

	  /* exchange particle data */
	  for(ngrp = 1; ngrp < (1 << PTask); ngrp++)
	    {
	      sendTask = ThisTask;
	      recvTask = ThisTask ^ ngrp;

	      if(recvTask < NTask)
		{
		  if(Send_count[recvTask] > 0 || Recv_count[recvTask] > 0)
		    {
		      /* get the particles */
		      MPI_Sendrecv(&DensDataIn[Send_offset[recvTask]],
				   Send_count[recvTask] * sizeof(struct densdata_in), MPI_BYTE,
				   recvTask, TAG_DENS_A,
				   &DensDataGet[Recv_offset[recvTask]],
				   Recv_count[recvTask] * sizeof(struct densdata_in), MPI_BYTE,
				   recvTask, TAG_DENS_A, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
		    }
		}
	    }

	  myfree(DensDataIn);
	  DensDataResult = (struct densdata_out *) mymalloc(nimport * sizeof(struct densdata_out));
	  DensDataOut = (struct densdata_out *) mymalloc(nexport * sizeof(struct densdata_out));

	  /* now do the particles that were sent to us */
	  for(j = 0; j < nimport; j++)
#ifdef DENSITY_SPLIT_BY_TYPE
	    subfind_density_evaluate(j, 1, &dummy, &dummy, j_in);
#else
	    subfind_density_evaluate(j, 1, &dummy, &dummy);
#endif
	  if(i >= NumPart)
	    ndone_flag = 1;
	  else
	    ndone_flag = 0;

	  MPI_Allreduce(&ndone_flag, &ndone, 1, MPI_INT, MPI_SUM, MPI_COMM_WORLD);

	  /* get the result */
	  for(ngrp = 1; ngrp < (1 << PTask); ngrp++)
	    {
	      sendTask = ThisTask;
	      recvTask = ThisTask ^ ngrp;
	      if(recvTask < NTask)
		{
		  if(Send_count[recvTask] > 0 || Recv_count[recvTask] > 0)
		    {
		      /* send the results */
		      MPI_Sendrecv(&DensDataResult[Recv_offset[recvTask]],
				   Recv_count[recvTask] * sizeof(struct densdata_out),
				   MPI_BYTE, recvTask, TAG_DENS_B,
				   &DensDataOut[Send_offset[recvTask]],
				   Send_count[recvTask] * sizeof(struct densdata_out),
				   MPI_BYTE, recvTask, TAG_DENS_B, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
		    }
		}
	    }

	  /* add the result to the local particles */
	  for(j = 0; j < nexport; j++)
	    {
	      place = DataIndexTable[j].Index;

	      P[place].DM_NumNgb += DensDataOut[j].Ngb;
	      P[place].u.DM_Density += DensDataOut[j].Rho;
#ifdef DENSITY_SPLIT_BY_TYPE
	      if(j_in >= 0)	/* sum up VelDisp only for own species */
#endif
		P[place].v.DM_VelDisp += DensDataOut[j].VelDisp;
	      DM_Vx[place] += DensDataOut[j].Vx;
	      DM_Vy[place] += DensDataOut[j].Vy;
	      DM_Vz[place] += DensDataOut[j].Vz;
	    }

	  myfree(DensDataOut);
	  myfree(DensDataResult);
	  myfree(DensDataGet);
	}
      while(ndone < NTask);

#ifdef DENSITY_SPLIT_BY_TYPE
      if(ThisTask == 0)
	printf("checking smoothing lenghs ... \n");
#endif

      /* do final operations on results */
      for(i = 0, npleft = 0; i < NumPart; i++)
	{
	  /* now check whether we had enough neighbours */
#ifdef DENSITY_SPLIT_BY_TYPE
	  if(j_in < 0)
	    Todo[i] = 0;
	  if((P[i].Type != j_target) || (j_in < 0))
	    continue;
#else
	  if(!((1 << P[i].Type) & (FOF_PRIMARY_LINK_TYPES)))
	    continue;
#endif

	  if(Todo[i])
	    {
	      if(P[i].DM_NumNgb != All.DesNumNgb &&
		 ((Right[i] - Left[i]) > 1.0e-4 * Left[i] || Left[i] == 0 || Right[i] == 0))
		{
		  /* need to redo this particle */
		  npleft++;

		  if(P[i].DM_NumNgb < All.DesNumNgb)
		    Left[i] = DMAX(P[i].DM_Hsml, Left[i]);
		  else
		    {
		      if(Right[i] != 0)
			{
			  if(P[i].DM_Hsml < Right[i])
			    Right[i] = P[i].DM_Hsml;
			}
		      else
			Right[i] = P[i].DM_Hsml;
		    }

		  if(iter >= MAXITER - 10)
		    {
		      printf
			("i=%d task=%d ID=%d Hsml=%g Left=%g Right=%g Ngbs=%g Right-Left=%g\n   pos=(%g|%g|%g)\n",
			 i, ThisTask, (int) P[i].ID, P[i].DM_Hsml, Left[i], Right[i],
			 (double) P[i].DM_NumNgb, Right[i] - Left[i], P[i].Pos[0], P[i].Pos[1], P[i].Pos[2]);
		      fflush(stdout);
		    }

		  if(Right[i] > 0 && Left[i] > 0)
		    P[i].DM_Hsml = pow(0.5 * (pow(Left[i], 3) + pow(Right[i], 3)), 1.0 / 3);
		  else
		    {
		      if(Right[i] == 0 && Left[i] == 0)
			endrun(8187);	/* can't occur */

		      if(Right[i] == 0 && Left[i] > 0)
			P[i].DM_Hsml *= 1.26;

		      if(Right[i] > 0 && Left[i] == 0)
			P[i].DM_Hsml /= 1.26;
		    }
		}
	      else
		Todo[i] = 0;
	    }
	}

      sumup_large_ints(1, &npleft, &ntot);

      t1 = second();

      if(ntot > 0)
	{
	  iter++;

	  if(iter > 0 && ThisTask == 0)
	    {
	      printf("ngb iteration %d: need to repeat for %d%09d particles. (took %g sec)\n", iter,
		     (int) (ntot / 1000000000), (int) (ntot % 1000000000), timediff(t0, t1));
	      fflush(stdout);
	    }

	  if(iter > MAXITER)
	    {
	      printf("failed to converge in neighbour iteration in density()\n");
	      fflush(stdout);
	      endrun(1155);
	    }
	}
    }
  while(ntot > 0);

  vel_to_phys = 1.0 / All.Time;

#ifdef DENSITY_SPLIT_BY_TYPE
  if(ThisTask == 0)
    printf("final operations ... \n");
#endif

  for(i = 0; i < NumPart; i++)
#ifdef DENSITY_SPLIT_BY_TYPE
    if(P[i].Type == j_target && j_in >= 0)
#else
    if(((1 << P[i].Type) & (FOF_PRIMARY_LINK_TYPES)))
#endif
      {
#ifdef DENSITY_SPLIT_BY_TYPE
	if(P[i].DM_NumNgb == 0)
	  endrun(5621);
#endif
	DM_Vx[i] /= P[i].DM_NumNgb;
	DM_Vy[i] /= P[i].DM_NumNgb;
	DM_Vz[i] /= P[i].DM_NumNgb;
	P[i].v.DM_VelDisp /= P[i].DM_NumNgb;

	P[i].v.DM_VelDisp = vel_to_phys * sqrt(P[i].v.DM_VelDisp -
					       DM_Vx[i] * DM_Vx[i] -
					       DM_Vy[i] * DM_Vy[i] - DM_Vz[i] * DM_Vz[i]);
      }

  myfree(DM_Vz);
  myfree(DM_Vy);
  myfree(DM_Vx);

  myfree(Todo);
  myfree(Right);
  myfree(Left);

  myfree(DataNodeList);
  myfree(DataIndexTable);

  myfree(Dist2list);
  myfree(Ngblist);
}

/*! This function represents the core of the SPH density computation. The
 *  target particle may either be local, or reside in the communication
 *  buffer.
 */
#ifdef DENSITY_SPLIT_BY_TYPE
int subfind_density_evaluate(int target, int mode, int *nexport, int *nsend_local, int tp)
#else
int subfind_density_evaluate(int target, int mode, int *nexport, int *nsend_local)
#endif
{
  int j, n;
  int startnode, numngb, ngb, listindex = 0;
  double hmax;
  double h, h2, hinv, hinv3;
  double rho, wk;
  double r, r2, u, mass_j, v2, vx, vy, vz;
  MyDouble *pos;

  rho = 0;
  numngb = 0;
  v2 = vx = vy = vz = 0;

  if(mode == 0)
    {
      pos = P[target].Pos;
      h = P[target].DM_Hsml;
    }
  else
    {
      pos = DensDataGet[target].Pos;
      h = DensDataGet[target].Hsml;
    }

#ifdef DENSITY_SPLIT_BY_TYPE
  if(h == 0)
    endrun(5622);
#endif

  h2 = h * h;
  hinv = 1.0 / h;
  hinv3 = hinv * hinv * hinv;

  if(mode == 0)
    {
      startnode = All.MaxPart;	/* root node */
    }
  else
    {
      startnode = DensDataGet[target].NodeList[0];
      startnode = Nodes[startnode].u.d.nextnode;	/* open it */
    }

  numngb = 0;

  while(startnode >= 0)
    {
      while(startnode >= 0)
	{
#ifdef DENSITY_SPLIT_BY_TYPE
	  if(tp < 0)
	    ngb =
	      subfind_ngb_treefind_linkpairs(pos, h, target, &startnode, mode, &hmax, nexport, nsend_local);
	  else
	    ngb = subfind_ngb_treefind_linkngb(pos, h, target, &startnode, mode, &hmax, nexport, nsend_local);
#else
	  ngb = subfind_ngb_treefind_linkngb(pos, h, target, &startnode, mode, &hmax, nexport, nsend_local);
#endif
	  if(ngb < 0)
	    return -1;

#ifdef DENSITY_SPLIT_BY_TYPE
	  if(tp >= 0)
#endif
	    if(mode == 0 && hmax > 0)
	      {
		P[target].DM_Hsml = hmax;
		h = hmax;
		h2 = h * h;
		hinv = 1.0 / h;
		hinv3 = hinv * hinv * hinv;

		if(ngb != All.DesNumNgb)
		  endrun(121);
	      }

	  numngb += ngb;

	  for(n = 0; n < ngb; n++)
	    {
	      j = Ngblist[n];

	      r2 = Dist2list[n];

#ifdef DENSITY_SPLIT_BY_TYPE
	      if(tp < 0)
		{
		  h = P[j].DM_Hsml;
		  if(h == 0)
		    endrun(5622);
		  h2 = h * h;
		  hinv = 1.0 / h;
		  hinv3 = hinv * hinv * hinv;
		}
#endif

	      if(r2 < h2)
		{
		  r = sqrt(r2);

		  u = r * hinv;

		  if(u < 0.5)
		    wk = hinv3 * (KERNEL_COEFF_1 + KERNEL_COEFF_2 * (u - 1) * u * u);
		  else
		    wk = hinv3 * KERNEL_COEFF_5 * (1.0 - u) * (1.0 - u) * (1.0 - u);

		  mass_j = P[j].Mass;

		  rho += (mass_j * wk);
		}

	      vx += P[j].Vel[0];
	      vy += P[j].Vel[1];
	      vz += P[j].Vel[2];

	      v2 += P[j].Vel[0] * P[j].Vel[0] + P[j].Vel[1] * P[j].Vel[1] + P[j].Vel[2] * P[j].Vel[2];
	    }
	}

      if(mode == 1)
	{
	  listindex++;
	  if(listindex < NODELISTLENGTH)
	    {
	      startnode = DensDataGet[target].NodeList[listindex];
	      if(startnode >= 0)
		startnode = Nodes[startnode].u.d.nextnode;	/* open it */
	    }
	}
    }

  if(mode == 0)
    {
      P[target].DM_NumNgb = numngb;
      P[target].u.DM_Density = rho;
#ifdef DENSITY_SPLIT_BY_TYPE
      if(tp >= 0)		/* sum up VelDisp only for own species */
#endif
	P[target].v.DM_VelDisp = v2;
      DM_Vx[target] = vx;
      DM_Vy[target] = vy;
      DM_Vz[target] = vz;
    }
  else
    {
      DensDataResult[target].Ngb = numngb;
      DensDataResult[target].Rho = rho;
      DensDataResult[target].VelDisp = v2;
      DensDataResult[target].Vx = vx;
      DensDataResult[target].Vy = vy;
      DensDataResult[target].Vz = vz;
    }

  return 0;
}

#ifdef DENSITY_SPLIT_BY_TYPE
void subfind_setup_smoothinglengths(int j)
#else
void subfind_setup_smoothinglengths(void)
#endif
{
  int i, no, p;

  for(i = 0; i < NumPart; i++)
    {
#ifdef DENSITY_SPLIT_BY_TYPE
      if(P[i].Type == j)
#else
      if(((1 << P[i].Type) & (FOF_PRIMARY_LINK_TYPES)))
#endif
	{
	  no = Father[i];

	  /* Not a good guess for gas/stars component, need more thought ! */
	  while(10 * All.DesNumNgb * P[i].Mass > Nodes[no].u.d.mass)
	    {
	      p = Nodes[no].u.d.father;

	      if(p < 0)
		break;

	      no = p;
	    }

#ifdef DENSITY_SPLIT_BY_TYPE
	  if(P[i].Type == 0)
	    P[i].DM_Hsml = PPP[i].Hsml;
	  else
#ifdef KEEP_DM_HSML_AS_GUESS
	  if(P[i].DM_Hsml < 0)
#endif
	    P[i].DM_Hsml =
	      pow(3.0 / (4 * M_PI) * All.DesNumNgb * P[i].Mass / Nodes[no].u.d.mass, 1.0 / 3) * Nodes[no].len;
#else
	  P[i].DM_Hsml =
	    pow(3.0 / (4 * M_PI) * All.DesNumNgb * P[i].Mass / Nodes[no].u.d.mass, 1.0 / 3) * Nodes[no].len;
#endif
	}
    }
}

static int Nhsml;

static struct hsml_data
{
  float Hsml;
  float Density;
  float VelDisp;
  MyIDType ID;
}
 *Hsml_list;

int subfind_compare_hsml_data(const void *a, const void *b)
{
  if(((struct hsml_data *) a)->ID < ((struct hsml_data *) b)->ID)
    return -1;

  if(((struct hsml_data *) a)->ID > ((struct hsml_data *) b)->ID)
    return +1;

  return 0;
}

void subfind_save_densities(int num)
{
  int i, nprocgroup, masterTask, groupTask;
  char buf[1000];
  double t0, t1;

  if(ThisTask == 0)
    {
      printf("start saving smoothing lengths and densities\n");
      fflush(stdout);
    }

  for(i = 0, Nhsml = 0; i < NumPart; i++)
#ifdef DENSITY_SPLIT_BY_TYPE
    if(((1 << P[i].Type) & (DENSITY_SPLIT_BY_TYPE)))
#else
    if(((1 << P[i].Type) & (FOF_PRIMARY_LINK_TYPES)))
#endif
      Nhsml++;

  MPI_Allgather(&Nhsml, 1, MPI_INT, Send_count, 1, MPI_INT, MPI_COMM_WORLD);
  for(i = 1, Send_offset[0] = 0; i < NTask; i++)
    Send_offset[i] = Send_offset[i - 1] + Send_count[i - 1];

  sumup_large_ints(1, &Nhsml, &Ntotal);

  Hsml_list = mymalloc(Nhsml * sizeof(struct hsml_data));

  for(i = 0, Nhsml = 0; i < NumPart; i++)
#ifdef DENSITY_SPLIT_BY_TYPE
    if(((1 << P[i].Type) & (DENSITY_SPLIT_BY_TYPE)))
#else
    if(((1 << P[i].Type) & (FOF_PRIMARY_LINK_TYPES)))
#endif
      {
	Hsml_list[Nhsml].Hsml = P[i].DM_Hsml;
	Hsml_list[Nhsml].Density = P[i].u.DM_Density;
	Hsml_list[Nhsml].VelDisp = P[i].v.DM_VelDisp;
#ifdef SAVE_HSML_IN_IC_ORDER
	Hsml_list[Nhsml].ID = P[i].ID_ic_order;
#else
	Hsml_list[Nhsml].ID = P[i].ID;
#endif
	Nhsml++;
      }

  t0 = second();
  parallel_sort(Hsml_list, Nhsml, sizeof(struct hsml_data), subfind_compare_hsml_data);
  t1 = second();

  if(ThisTask == 0)
    {
      printf("Sorting of densities in ID sequence took = %g sec\n", timediff(t0, t1));
      fflush(stdout);
    }

  if(ThisTask == 0)
    {
      sprintf(buf, "%s/hsmldir_%03d", All.OutputDir, num);
      mkdir(buf, 02755);
    }
  MPI_Barrier(MPI_COMM_WORLD);

  if(NTask < All.NumFilesWrittenInParallel)
    {
      printf
	("Fatal error.\nNumber of processors must be a smaller or equal than `NumFilesWrittenInParallel'.\n");
      endrun(241931);
    }

  nprocgroup = NTask / All.NumFilesWrittenInParallel;
  if((NTask % All.NumFilesWrittenInParallel))
    nprocgroup++;
  masterTask = (ThisTask / nprocgroup) * nprocgroup;
  for(groupTask = 0; groupTask < nprocgroup; groupTask++)
    {
      if(ThisTask == (masterTask + groupTask))	/* ok, it's this processor's turn */
	subfind_save_local_densities(num);
      MPI_Barrier(MPI_COMM_WORLD);	/* wait inside the group */
    }

  myfree(Hsml_list);

}

void subfind_save_local_densities(int num)
{
  char fname[1000];
  int i;
  float *tmp;
  FILE *fd;

  sprintf(fname, "%s/hsmldir_%03d/%s_%03d.%d", All.OutputDir, num, "hsml", num, ThisTask);
  if(!(fd = fopen(fname, "w")))
    {
      printf("can't open file `%s`\n", fname);
      endrun(1183);
    }

  my_fwrite(&Nhsml, sizeof(int), 1, fd);
  my_fwrite(&Send_offset[ThisTask], sizeof(int), 1, fd);	/* this is the number of IDs in previous files */
  my_fwrite(&Ntotal, sizeof(long long), 1, fd);
  my_fwrite(&NTask, sizeof(int), 1, fd);

  tmp = mymalloc(Nhsml * sizeof(float));

  for(i = 0; i < Nhsml; i++)
    tmp[i] = Hsml_list[i].Hsml;
  my_fwrite(tmp, sizeof(float), Nhsml, fd);

  for(i = 0; i < Nhsml; i++)
    tmp[i] = Hsml_list[i].Density;
  my_fwrite(tmp, sizeof(float), Nhsml, fd);

  for(i = 0; i < Nhsml; i++)
    tmp[i] = Hsml_list[i].VelDisp;
  my_fwrite(tmp, sizeof(float), Nhsml, fd);

  myfree(tmp);

  fclose(fd);
}

#endif


// many-small-dgemms.c:
// Mini-app which investigates different ways to perform many small dgemm operations.
// Looping over thousands of rocblas calls is slow, so looking into other solutions.
// This functionality is very interesting to several applications, including nuccor.
// Date: Oct 15th, 2019
// Author: Justin Gage Lietz
// Contact: lietzjg@ornl.gov

#include <stdio.h>
#include <stdlib.h>
#include <time.h>
#include <cblas.h>
/* #include <cublas_v2.h> */
#include <hip/hip_runtime.h>
#include <rocblas.h>
#include <omp.h>
#include <string.h> //for memcpy

#include <magma_v2.h>

// tolerance as 1.e-15 fails
#define TOL 1.e-14

// Pulled from magma test code
#define TESTING_CHECK( err )                                                 \
    do {                                                                     \
        magma_int_t err_ = (err);                                            \
        if ( err_ != 0 ) {                                                   \
            fprintf( stderr, "Error: %s\nfailed at %s:%d: error %lld: %s\n", \
                     #err, __FILE__, __LINE__,                               \
                     (long long) err_, magma_strerror(err_) );               \
            exit(1);                                                         \
        }                                                                    \
    } while( 0 )

size_t idx(size_t i, size_t j, size_t nrows, size_t ncols){
  return j*nrows + i;
}

size_t cidx(size_t i, size_t j, size_t nrows, size_t ncols){
  return i*ncols + j;
}

void print_c_matrix(double *mat, size_t nrows, size_t ncols){
  for(int i = 0; i < nrows; i++){
    for(int j = 0; j < ncols; j++){
      printf("%f ", mat[ cidx(i,j,nrows,ncols) ]);
    }
    printf("\n");
  }
  printf("\n");
}

void print_matrix(double *mat, size_t nrows, size_t ncols){
  for(int i = 0; i < nrows; i++){
    for(int j = 0; j < ncols; j++){
      printf("%f ", mat[ idx(i,j,nrows,ncols) ]);
    }
    printf("\n");
  }
  printf("\n");
}

void print_matrix_array(double *mat_A, size_t *Aoffsets, size_t *nrows, size_t *ncols, size_t nBlocks){
  for(size_t iBlock = 0; iBlock < nBlocks; iBlock++){
    size_t Aoffset = Aoffsets[iBlock];
    printf("iBlock: %zu\n",iBlock);
    print_matrix( &mat_A[Aoffset], nrows[iBlock], ncols[iBlock] );
  }
}

void cblas_wrapper(double *A, double* B, double* C, size_t M_in, size_t N_in, size_t K_in){
  int M = M_in;
  int K = K_in;
  int N = N_in;
  int LDA = M;
  int LDB = K;
  int LDC = M;
  double alpha = 1.0;
  double beta = 0.0;

  cblas_dgemm(/*CBLAS*/ CblasColMajor,
      /* TRANS A */ CblasNoTrans,
      /* TRANS B */ CblasNoTrans,
      /* M */M,
      /* N */N,
      /* K */K,
      /* alpha */alpha,
      /* A */A,
      /* LDA */LDA,
      /* B  */B,
      /* LDB */LDB,
      /* BETA */beta,
      /* C */C,
      /* LDC */LDC);
}

void device_dgemm(rocblas_handle handle, double *A, double* B, double* C, size_t M_in, size_t N_in, size_t K_in){
  int M = M_in;
  int K = K_in;
  int N = N_in;
  int LDA = M;
  int LDB = K;
  int LDC = M;
  double alpha = 1.0;
  double beta = 0.0;

  size_t ABytes = sizeof(double)*M*K;
  size_t BBytes = sizeof(double)*K*N;
  size_t CBytes = sizeof(double)*M*N;

  double *d_A;
  double *d_B;
  double *d_C;
  const rocblas_operation transA = rocblas_operation_none;
  const rocblas_operation transB = rocblas_operation_none;

  hipMalloc(&d_A, ABytes);
  hipMalloc(&d_B, BBytes);
  hipMalloc(&d_C, CBytes);

  hipMemcpy(d_A, A, ABytes, hipMemcpyHostToDevice);
  hipMemcpy(d_B, B, BBytes, hipMemcpyHostToDevice);

  rocblas_dgemm(/*rocblas handle*/ handle,
      /* TRANS A */ transA,
      /* TRANS B */ transB,
      /* M */M,
      /* N */N,
      /* K */K,
      /* alpha */&alpha,
      /* A */d_A,
      /* LDA */LDA,
      /* B  */d_B,
      /* LDB */LDB,
      /* BETA */&beta,
      /* C */d_C,
      /* LDC */LDC);

  hipMemcpy(C, d_C, CBytes, hipMemcpyDeviceToHost);

  hipFree(d_A);
  hipFree(d_B);
  hipFree(d_C);
}

void rocblas_array_dgemm(
    rocblas_handle handle,
    const double *A, const double* B, double* C,
    const size_t *Ms, const size_t *Ns, const size_t *Ks,
    const size_t *Aoffsets, const size_t *Boffsets, const size_t *Coffsets, const size_t nBlocks){

  const rocblas_operation transA = rocblas_operation_none;
  const rocblas_operation transB = rocblas_operation_none;

  double *d_A, *d_B, *d_C;
  size_t ABytes = Aoffsets[nBlocks]*sizeof(double);
  size_t BBytes = Boffsets[nBlocks]*sizeof(double);
  size_t CBytes = Coffsets[nBlocks]*sizeof(double);

  hipMalloc(&d_A, ABytes);
  hipMalloc(&d_B, BBytes);
  hipMalloc(&d_C, CBytes);

  hipMemcpy(d_A, A, ABytes, hipMemcpyHostToDevice);
  hipMemcpy(d_B, B, BBytes, hipMemcpyHostToDevice);

  double start = omp_get_wtime();

  for(size_t iBlock = 0; iBlock < nBlocks; iBlock++){
    size_t Aoffset = Aoffsets[iBlock];
    size_t Boffset = Boffsets[iBlock];
    size_t Coffset = Coffsets[iBlock];

    int M = Ms[iBlock];
    int N = Ns[iBlock];
    int K = Ks[iBlock];
    int LDA = M;
    int LDB = K;
    int LDC = M;
    double alpha = 1.0;
    double beta = 0.0;
    /* char transa = 'N'; */
    /* char transb = 'N'; */

    rocblas_dgemm(/*rocblas handle*/ handle,
        /* TRANS A */ transA,
        /* TRANS B */ transB,
        /* M */       M,
        /* N */       N,
        /* K */       K,
        /* alpha */   &alpha,
        /* A */       &d_A[Aoffset],
        /* LDA */     LDA,
        /* B  */      &d_B[Boffset],
        /* LDB */     LDB,
        /* BETA */    &beta,
        /* C */       &d_C[Coffset],
        /* LDC */     LDC);
  }
  double stop = omp_get_wtime();
  printf("time of rocblas_dgemm loop: %f\n", stop - start);
  hipMemcpy(C, d_C, CBytes, hipMemcpyDeviceToHost);

  hipFree(d_A);
  hipFree(d_B);
  hipFree(d_C);
}

void magma_array_dgemm(
    const double *A, const double* B, double* C,
    const size_t *Ms, const size_t *Ns, const size_t *Ks,
    const size_t *Aoffsets, const size_t *Boffsets, const size_t *Coffsets, const size_t nBlocks){
  magma_trans_t transA = MagmaNoTrans;
  magma_trans_t transB = MagmaNoTrans;
  double const* * dA_array;
  double const* * dB_array;
  double **dC_array;
  double 	alpha;
  double 	beta;
  magma_int_t* d_lddb;
  magma_int_t* d_ldda;
  magma_int_t* d_lddc;
  magma_int_t* d_m;
  magma_int_t* d_n;
  magma_int_t* d_k;
  int* h_m;
  int* h_n;
  int* h_k;

  double **hA_array;
  double **hB_array;
  double **hC_array;
  double *d_A_elems;
  double *d_B_elems;
  double *d_C_elems;

  magma_int_t batchCount;
  magma_queue_t queue;
  magma_device_t device;

  magma_getdevice( &device );
  magma_queue_create( device, &queue );

  batchCount = nBlocks;
  // Magma needs larger arrays
  TESTING_CHECK( magma_malloc((void**)&d_m, (batchCount+1)*sizeof(magma_int_t)) );
  TESTING_CHECK( magma_malloc((void**)&d_n, (batchCount+1)*sizeof(magma_int_t)) );
  TESTING_CHECK( magma_malloc((void**)&d_k, (batchCount+1)*sizeof(magma_int_t)) );

  TESTING_CHECK( magma_malloc((void**)&d_ldda, (batchCount+1)*sizeof(magma_int_t) ) );
  TESTING_CHECK( magma_malloc((void**)&d_lddb, (batchCount+1)*sizeof(magma_int_t) ) );
  TESTING_CHECK( magma_malloc((void**)&d_lddc, (batchCount+1)*sizeof(magma_int_t) ) );

  // Looks like magma_malloc works with regular c int as well.
  TESTING_CHECK( magma_malloc_cpu((void**)&h_m, (batchCount+1)*sizeof(int) ) );
  TESTING_CHECK( magma_malloc_cpu((void**)&h_n, (batchCount+1)*sizeof(int) ) );
  TESTING_CHECK( magma_malloc_cpu((void**)&h_k, (batchCount+1)*sizeof(int) ) );

  // dA_array is the array of pointers need by dgemm
  // d_A_elems are the actual mtx elements being pointed to
  // hA_array is the host side pointers that will get passed to dA_array
  TESTING_CHECK( magma_malloc( (void**)&dA_array, sizeof(double*)*batchCount ) );
  TESTING_CHECK( magma_malloc( (void**)&dB_array, sizeof(double*)*batchCount ) );
  TESTING_CHECK( magma_malloc( (void**)&dC_array, sizeof(double*)*batchCount ) );

  TESTING_CHECK( magma_malloc( (void**)&d_A_elems, sizeof(double)*Aoffsets[nBlocks] ) );
  TESTING_CHECK( magma_malloc( (void**)&d_B_elems, sizeof(double)*Boffsets[nBlocks] ) );
  TESTING_CHECK( magma_malloc( (void**)&d_C_elems, sizeof(double)*Coffsets[nBlocks] ) );

  TESTING_CHECK( magma_malloc_cpu( (void**)&hA_array, sizeof(double*)*batchCount ) );
  TESTING_CHECK( magma_malloc_cpu( (void**)&hB_array, sizeof(double*)*batchCount ) );
  TESTING_CHECK( magma_malloc_cpu( (void**)&hC_array, sizeof(double*)*batchCount ) );

  // magma needs 32-bit int, while main was using 64
  for(size_t iBlock = 0; iBlock < nBlocks; iBlock++){
    h_m[iBlock] = Ms[iBlock];
    h_n[iBlock] = Ns[iBlock];
    h_k[iBlock] = Ks[iBlock];
  }

  magma_setvector(batchCount, sizeof(magma_int_t), h_m, 1, d_m, 1, queue);
  magma_setvector(batchCount, sizeof(magma_int_t), h_n, 1, d_n, 1, queue);
  magma_setvector(batchCount, sizeof(magma_int_t), h_k, 1, d_k, 1, queue);
  magma_setvector(batchCount, sizeof(magma_int_t), h_m, 1, d_ldda, 1, queue);
  magma_setvector(batchCount, sizeof(magma_int_t), h_k, 1, d_lddb, 1, queue);
  magma_setvector(batchCount, sizeof(magma_int_t), h_m, 1, d_lddc, 1, queue);

  magma_setvector(Aoffsets[nBlocks], sizeof(double), A, 1, d_A_elems, 1, queue);
  magma_setvector(Boffsets[nBlocks], sizeof(double), B, 1, d_B_elems, 1, queue);

  size_t Aoffset;
  size_t Boffset;
  size_t Coffset;

  for(int iBlock = 0; iBlock < nBlocks; iBlock++){
    Aoffset = Aoffsets[iBlock];
    Boffset = Boffsets[iBlock];
    Coffset = Coffsets[iBlock];

    hA_array[iBlock] = d_A_elems + Aoffset;
    hB_array[iBlock] = d_B_elems + Boffset;
    hC_array[iBlock] = d_C_elems + Coffset;
  }

  magma_setvector(batchCount, sizeof(double*), hA_array, 1, dA_array, 1, queue);
  magma_setvector(batchCount, sizeof(double*), hB_array, 1, dB_array, 1, queue);
  magma_setvector(batchCount, sizeof(double*), hC_array, 1, dC_array, 1, queue);

  alpha = 1.0;
  beta = 0.0;
  batchCount = nBlocks;

  double start,stop;

  start = omp_get_wtime();

  magmablas_dgemm_vbatched(	      transA,
      /* magma_trans_t */ 	      transB,
      /* magma_int_t * */         d_m,
      /* magma_int_t * */	        d_n,
      /* magma_int_t * */	        d_k,
      /* double */	              alpha,
      /* double const *const * */	dA_array,
      /* magma_int_t * */	        d_ldda,
      /* double const *const * */	dB_array,
      /* magma_int_t * */	        d_lddb,
      /* double */	              beta,
      /* double ** */	            dC_array,
      /* magma_int_t * */	        d_lddc,
      /* magma_int_t */	          batchCount,
      /* magma_queue_t */	        queue);
  stop = omp_get_wtime();

  printf("time of magmablas_dgemm_vbatched: %f\n", stop - start);

  magma_getvector(Coffsets[nBlocks], sizeof(double), d_C_elems, 1, C, 1, queue);

  // Clean up
  magma_queue_destroy( queue );

  TESTING_CHECK( magma_free(d_m) );
  TESTING_CHECK( magma_free(d_n) );
  TESTING_CHECK( magma_free(d_k) );

  TESTING_CHECK( magma_free(d_ldda) );
  TESTING_CHECK( magma_free(d_lddb) );
  TESTING_CHECK( magma_free(d_lddc) );

  TESTING_CHECK( magma_free_cpu(h_m) );
  TESTING_CHECK( magma_free_cpu(h_n) );
  TESTING_CHECK( magma_free_cpu(h_k) );

  TESTING_CHECK( magma_free(dA_array) );
  TESTING_CHECK( magma_free(dB_array) );
  TESTING_CHECK( magma_free(dC_array) );

  TESTING_CHECK( magma_free(d_A_elems) );
  TESTING_CHECK( magma_free(d_B_elems) );
  TESTING_CHECK( magma_free(d_C_elems) );

  TESTING_CHECK( magma_free_cpu(hA_array) );
  TESTING_CHECK( magma_free_cpu(hB_array) );
  TESTING_CHECK( magma_free_cpu(hC_array) );
}

void host_matmul(const double *A, const double* B, double* C, const size_t M, const size_t N, const size_t K){
  for(size_t i = 0; i < M; i++){
    for(size_t j = 0; j < N; j++){
      C[ idx(i,j,M,N) ] = 0.0;
      for(size_t k = 0; k < K; k++){
        C[ idx(i,j,M,N) ] += A[ idx(i,k,M,K) ] * B[ idx(k,j,K,N) ];
      }
    }
  }
}

void host_array_matmul(const double *A, const double* B, double* C,
    const size_t *Ms, const size_t *Ns, const size_t *Ks,
    const size_t *Aoffsets, const size_t *Boffsets, const size_t *Coffsets, const size_t nBlocks){

  for(size_t iBlock = 0; iBlock < nBlocks; iBlock++){
    size_t Aoffset = Aoffsets[iBlock];
    size_t Boffset = Boffsets[iBlock];
    size_t Coffset = Coffsets[iBlock];

    host_matmul(&A[Aoffset], &B[Boffset], &C[Coffset], Ms[iBlock], Ns[iBlock], Ks[iBlock]);
  }
}

void host_array_dgemm(const double *A, const double* B, double* C,
    const size_t *Ms, const size_t *Ns, const size_t *Ks,
    const size_t *Aoffsets, const size_t *Boffsets, const size_t *Coffsets, const size_t nBlocks){

  for(size_t iBlock = 0; iBlock < nBlocks; iBlock++){
    size_t Aoffset = Aoffsets[iBlock];
    size_t Boffset = Boffsets[iBlock];
    size_t Coffset = Coffsets[iBlock];

    int M = Ms[iBlock];
    int N = Ns[iBlock];
    int K = Ks[iBlock];
    int LDA = M;
    int LDB = K;
    int LDC = M;
    double alpha = 1.0;
    double beta = 0.0;
    /* char transa = 'N'; */
    /* char transb = 'N'; */

    cblas_dgemm(/*CBLAS*/ CblasColMajor,
        /* TRANS A */ CblasNoTrans,
        /* TRANS B */ CblasNoTrans,
        /* M */       M,
        /* N */       N,
        /* K */       K,
        /* alpha */   alpha,
        /* A */       &A[Aoffset],
        /* LDA */     LDA,
        /* B  */      &B[Boffset],
        /* LDB */     LDB,
        /* BETA */    beta,
        /* C */       &C[Coffset],
        /* LDC */     LDC);
  }
}

void host_transpose(double *mat_in, double *mat_out, size_t nrows_in, size_t ncols_in){
  for(size_t i = 0; i < nrows_in; i++){
    for(size_t j = 0; j < ncols_in; j++){
      mat_out[ idx(j,i,ncols_in,nrows_in) ] = mat_in[ idx(i,j,nrows_in, ncols_in) ];
    }
  }
}

int check_equal(double *mat_A, double *mat_B, size_t nrows, size_t ncols){
  int flag = 1;
  for(size_t i = 0; i < nrows; i++){
    for(size_t j = 0; j < ncols; j++){
      if( fabs( mat_A[ idx(i,j,nrows,ncols) ] - mat_B[ idx(i,j,nrows,ncols) ]) > TOL ){
        flag = 0;
        return flag;
      }
    }
  }
  return flag;
}

int check_matrix_array_equal(double *mat_A, double *mat_B, size_t *Aoffsets, size_t *Boffsets, size_t *nrows, size_t *ncols, size_t nBlocks){
  int flag = 1;
  for(size_t iBlock = 0; iBlock < nBlocks; iBlock++){
    size_t Aoffset = Aoffsets[iBlock];
    size_t Boffset = Boffsets[iBlock];

    flag *= check_equal( &mat_A[Aoffset], &mat_B[Boffset], nrows[iBlock], ncols[iBlock] );
    if( !check_equal(&mat_A[Aoffset], &mat_B[Boffset], nrows[iBlock], ncols[iBlock]) ){
      printf("error in block: %zu\n");
      printf("A\n");
      print_matrix( &mat_A[Aoffset], nrows[iBlock], ncols[iBlock] );
      printf("B\n");
      print_matrix( &mat_B[Aoffset], nrows[iBlock], ncols[iBlock] );
      return flag;
    }
  }
  return flag;
}

void generate_random_array(size_t* array, size_t nBlocks, size_t max_size){
  size_t rando;

  for(size_t iBlock = 0; iBlock < nBlocks; iBlock++){
    rando = rand()%max_size + 1;
    array[iBlock] = rando;
  }
}

void load_random_matrix(double* mat_in, size_t nrows, size_t ncols){
  double rando;

  for(size_t i = 0; i < nrows; i++){
    for(size_t j = 0; j < ncols; j++){
      rando= (double) rand()/RAND_MAX*2.0 - 1.0;
      mat_in[ idx(i,j,nrows,ncols) ] = rando;
    }
  }
}

int main(){
  // Start magma
  TESTING_CHECK( magma_init() );
  magma_print_environment();

  // magma seems to break at 2^16, but 2^15 blocks is fine
  size_t nBlocks = 1<<15;
  size_t max_block_size = 10;
  // Small test for debugging
  // size_t nBlocks = 3;
  // size_t max_block_size = 4;

  size_t *Ms_array;
  size_t *Ns_array;
  size_t *Ks_array;
  size_t *Aoffsets;
  size_t *Boffsets;
  size_t *Coffsets;

  Ms_array = (size_t*)malloc(sizeof(size_t)*nBlocks);
  Ns_array = (size_t*)malloc(sizeof(size_t)*nBlocks);
  Ks_array = (size_t*)malloc(sizeof(size_t)*nBlocks);

  Aoffsets = (size_t*)malloc(sizeof(size_t)*(nBlocks+1));
  Boffsets = (size_t*)malloc(sizeof(size_t)*(nBlocks+1));
  Coffsets = (size_t*)malloc(sizeof(size_t)*(nBlocks+1));

  /* srand( time(NULL) ); */
  // Reproducible seed
  srand( 100 );

  generate_random_array(Ms_array,nBlocks,max_block_size);
  generate_random_array(Ns_array,nBlocks,max_block_size);
  generate_random_array(Ks_array,nBlocks,max_block_size);

  Aoffsets[0] = 0;
  Boffsets[0] = 0;
  Coffsets[0] = 0;

  // Offset array elements indicate starting location of each block
  // Final array element indicates where the next block WOULD start
  // This is needed to count the total number of elements
  for(size_t iBlock = 1; iBlock < nBlocks + 1; iBlock++){
    Aoffsets[iBlock] = Aoffsets[iBlock-1] + Ms_array[iBlock-1] * Ks_array[iBlock-1];
    Boffsets[iBlock] = Boffsets[iBlock-1] + Ks_array[iBlock-1] * Ns_array[iBlock-1];
    Coffsets[iBlock] = Coffsets[iBlock-1] + Ms_array[iBlock-1] * Ns_array[iBlock-1];
  }

  size_t ABytes = Aoffsets[nBlocks]*sizeof(double);
  size_t BBytes = Boffsets[nBlocks]*sizeof(double);
  size_t CBytes = Coffsets[nBlocks]*sizeof(double);

  double *h_A;
  double *h_B;
  double *h_C;
  double *h_C_host;
  double *h_C_device;
  double *h_C_magma;

  h_A = (double*)malloc(ABytes);
  h_B = (double*)malloc(BBytes);
  h_C = (double*)malloc(CBytes);
  h_C_host = (double*)malloc(CBytes);
  h_C_device = (double*)malloc(CBytes);
  h_C_magma = (double*)malloc(CBytes);

  for(size_t iBlock = 0; iBlock < nBlocks; iBlock++){
    size_t Aoffset = Aoffsets[iBlock];
    size_t Boffset = Boffsets[iBlock];

    load_random_matrix(&h_A[Aoffset], Ms_array[iBlock], Ks_array[iBlock]);
    load_random_matrix(&h_B[Boffset], Ks_array[iBlock], Ns_array[iBlock]);
  }

  double start,stop;
  int check;
  //////////////////////////////////////////////////////////////////////
  // Tests on host
  //////////////////////////////////////////////////////////////////////

  start = omp_get_wtime();
  host_array_matmul(h_A, h_B, h_C, Ms_array, Ns_array, Ks_array, Aoffsets, Boffsets, Coffsets, nBlocks);
  stop = omp_get_wtime();
  printf("naive loop time: %f\n", stop - start);

  start = omp_get_wtime();
  host_array_dgemm(h_A, h_B, h_C_host, Ms_array, Ns_array, Ks_array, Aoffsets, Boffsets, Coffsets, nBlocks);
  stop = omp_get_wtime();
  printf("host blas time: %f\n", stop - start);

  check = check_matrix_array_equal(h_C_host, h_C, Coffsets, Coffsets, Ms_array, Ns_array, nBlocks);
  if( check == 0 ){
    printf("check host dgemm failed: %d\n",check);
  }

  //////////////////////////////////////////////////////////////////////
  // Run tests on device
  //////////////////////////////////////////////////////////////////////

  rocblas_handle handle;
  rocblas_create_handle(&handle);

  start = omp_get_wtime();
  rocblas_array_dgemm(handle, h_A, h_B, h_C_device, Ms_array, Ns_array, Ks_array, Aoffsets, Boffsets, Coffsets, nBlocks);
  stop = omp_get_wtime();
  printf("rocblas wrapper time: %f\n", stop - start);

  check = check_matrix_array_equal(h_C_device, h_C, Coffsets, Coffsets, Ms_array, Ns_array, nBlocks);
  if( check == 0 ){
    printf("check device dgemm failed: %d\n",check);
  }

  start = omp_get_wtime();
  magma_array_dgemm(h_A, h_B, h_C_magma, Ms_array, Ns_array, Ks_array, Aoffsets, Boffsets, Coffsets, nBlocks);
  stop = omp_get_wtime();
  printf("magma blas wrapper time: %f\n", stop - start);

  check = check_matrix_array_equal(h_C_magma, h_C, Coffsets, Coffsets, Ms_array, Ns_array, nBlocks);
  if( check == 0 ){
    printf("magma dgemm failed: %d\n",check);
  }

  // Clean up
  rocblas_destroy_handle(handle);
  free(h_A);
  free(h_B);
  free(h_C);
  free(h_C_host);
  free(h_C_device);
  free(h_C_magma);
  magma_finalize();

  return 0;
}