/*
 * Copyright (c) 2025 Cerebras Systems Inc.
 * All rights reserved.
 */

import params;

xp param nlocmax,nmax,nintmax;

xp ymptr,y0ptr,ypptr,yend;
xp xmptr,x0ptr,xpptr,xend;
xp zpptr,zend;

xp icount1,icount2;
xp nint;

xp const range = 3;

xp const countsize = range + 1;
xp const firstsize = range + 1;
xp count1[5][countsize],count3[5][countsize],count4[5][countsize];
xp first1[5][firstsize],first3[5][firstsize],first4[5][firstsize];

xp const memsize = 1500;
sp fpmem[n=memsize];
xp xpmem[n=2*memsize] address(&fpmem);
xp memend[1] address(&xpmem + 2*memsize);


xp countblock[nd=5][n=countsize] extern address(&xpmem);
xp xmcount[n=countsize] address(&xpmem);
xp xpcount[n=countsize] address(&xmcount + countsize);
xp ymcount[n=countsize] address(&xpcount + countsize);
xp ypcount[n=countsize] address(&ymcount + countsize);
xp zpcount[n=countsize] address(&ypcount + countsize);
xp countend[1] address(&zpcount + countsize);

xp firstblock[nd=5][n=firstsize] extern address(&countend);
xp xmfirst[n=firstsize] address(&countend);
xp xpfirst[n=firstsize] address(&xmfirst + firstsize);
xp ymfirst[n=firstsize] address(&xpfirst + firstsize);
xp ypfirst[n=firstsize] address(&ymfirst + firstsize);
xp zpfirst[n=firstsize] address(&ypfirst + firstsize);
xp firstend[1] address(&zpfirst + firstsize);

sp fpalign[1] address(&firstend+(firstsize + countsize)%2);
sp vel[n=nlocmax][m=3] address(&fpalign);
sp vel1[n=3*nlocmax] extern address(&vel);
xp velend[1] address(&vel + 3*nlocmax*2);
sp pos[n=nmax][m=4] address(&velend);
sp pos0[n=nlocmax][m=4] extern address(&pos); // Ref. to local atoms for miset/miget
sp pos1[n=4*nmax] address(&pos);
xp posend[1] address(&pos + 4*nmax*2);
sp ff[n=nmax][m=3] address(&posend);
sp ff0[n=nlocmax][m=3] address(&ff); // Ref. to local atoms for miget
sp ff1[n=3*nmax] extern address(&ff);
xp ffend[1] address(&ff + 3*nmax*2);

sp dxij[nintmax][3] extern address(&ffend);
xp dxijend[1] address(&dxij + 3*nintmax*2);

sp rinv[n=nintmax] extern address(&dxijend);
sp tmp1[nintmax] address(&rinv + nintmax*2);
sp tmp2[nintmax] address(&tmp1 + nintmax*2);
xp tmpend[1] address(&tmp2 + nintmax*2);
xp varend[1] address(&tmpend);
// varend = end of variables defined in {xp,fp}mem.
// Check map file that varend <= memend.

sp dist2[n=nintmax];
sp df[n=nintmax];

xp ilist[n=nintmax];
xp jlist[n=nintmax];
xp klist[nintmax];
//xp klocal[nlocmax];
xp atomtype[nmax];

struct paramstruct params;
xp nlocal;
xp hix,hiy;

struct fparamstruct fparams;

sp invmass[2] address(&fparams + 2*6);
sp box[3] address(&invmass + 2*2);
sp boxinv[3] address(&box + 2*3);
sp cell[3][2] address(&boxinv + 2*3);
sp center[3] address(&cell + 2*6);
sp dt,rcut2,rscale,rhoscale;

sp const mhalf = -0.5;
sp const convmagic_const = 1024.0 * 1024.0 * 8.0; // 1<<23
sp convmagic = convmagic_const;


/*sp*/ socket s_xm,s_ym,s_zm; // x-, y-, z-
/*sp*/ socket s_xp,s_yp,s_zp; // x+, y+, z+
/*sp*/ socket ping;

// Multicast control wavelet sequence:
//struct wv { xp a; xp b; };

//                 (adv*,adv*,rst*), (adv*)
wv xycmds[n=2] = { {0x0001, 0x2b40}, {0x0000, 0x0140} };
//                 (adv*,rst*)     , (adv*)
wv zcmds[n=2]  = { {0x0000, 0x0d40}, {0x0000, 0x0140} };
xp const ncmd = 2;
// * = no-CE bit set.
wv adv[1] = { {0x9249,0x0940} };

xp const router_config_addr = 0x7c00; // SDR
xp const inq_config_addr = 0x7bc0; // SDR
xp const outq_config_addr = 0x7b80;  // SDR
xp router_config_regs[24] address(router_config_addr);
//xp switch_config_regs[24][4] address(0x7c80); // SDR
xp inq_config_regs[8][2] address(inq_config_addr);
xp outq_config_regs[8] address(outq_config_addr);


function set_iq_color(xp q,xp color) {
  xp t = inq_config_regs[q][0];
  t >>= 5;
  t = (t << 5) | (color & 31);
  inq_config_regs[q][0] = t;
}

function get_iq_color(xp q) {
  xp t = inq_config_regs[q][0];
  t = t & 31;
  return t;
}

function set_oq_color(xp q,xp color) {
  xp const colmask = 0xfe0f; // ~(31<<4)
  xp t = outq_config_regs[q];
  t = t & colmask;
  t = t | (color << 4);
  outq_config_regs[q] = t;
}

function get_oq_color(xp q) {
  xp t = outq_config_regs[q];
  t = (t >> 4) & 31;
  return t;
}


function clear_teardown(xp color) {
  xp const tdmask = 0x7fff; // ~(1<<15) (SDR)
  xp t = router_config_regs[color];
  t = t & tdmask;
  router_config_regs[color] = t;
}

function addressinit() {
  /*
  router_config_regs.base = router_config_addr;
  inq_config_regs.base = inq_config_addr;
  outq_config_regs.base = outq_config_addr;
  */
}


function initparams() {
  nlocal = params.nlocal;
  hix = params.hix;
  hiy = params.hiy;
  dt = fparams.dt;
  rcut2 = fparams.rcut * fparams.rcut;
  rscale = fparams.rscale;
  rhoscale = fparams.rhoscale;
}


function getcount() {
  let ic in [0,2);
  xp r0 = 0;
  //#pragma uthread enable off
  while(r0 < range) {
    // send in x+ direction
    s_xp[] <- (sp) (xmcount[r0]);
    ∀ic s_xp[] <- control(xycmds[ic]);
    // send in x- direction
    s_xm[] <- (sp) xpcount[r0];
    ∀ic s_xm[] <- control(xycmds[ic]);
    // send in y+ direction
    s_yp[] <- (sp) ymcount[r0];
    ∀ic s_yp[] <- control(xycmds[ic]);
    // send in y- direction
    s_ym[] <- (sp) ypcount[r0];
    ∀ic s_ym[] <- control(xycmds[ic]);

    // Receive one count from each direction
    sp tmp1 = s_xp[];
    sp tmp2 = s_xm[];
    sp tmp3 = s_yp[];
    sp tmp4 = s_ym[];
    xmcount[r0+1] <- tmp1;//(xp) ((sp) (s_xp[]));
    xpcount[r0+1] <- tmp2;
    ymcount[r0+1] <- tmp3;
    ypcount[r0+1] <- tmp4;

    r0 += 1;
  }

  { // send in z- direction
    xp c = 0,j;
    ∀j in [0,range+1) c += xmcount[j];
    ∀j in [1,range+1) c += xpcount[j];
    zpcount[0] = c;
    j = 0;
    while(j < range) { 
      s_zm[] <- (sp) zpcount[j];
      ∀ic s_zm[] <- control(zcmds[ic]);
      //zpcount[j+1] <- s_zm[];
      sp tmpz = s_zm[];
      zpcount[j+1] = tmpz;
      j += 1;
    }
  }
}

function exchange() {
  let ic in [0,2);
  xp r0 = 0;

  while(r0 < range) {
    xp fxm0 = xmfirst[r0],fxm1 = xmfirst[r0+1];
    xp cxm0 = xmcount[r0],cxm1 = xmcount[r0+1];
    xp fxp0 = xpfirst[r0],fxp1 = xpfirst[r0+1];
    xp cxp0 = xpcount[r0],cxp1 = xpcount[r0+1];
    xp fym0 = ymfirst[r0],fym1 = ymfirst[r0+1];
    xp cym0 = ymcount[r0],cym1 = ymcount[r0+1];
    xp fyp0 = ypfirst[r0],fyp1 = ypfirst[r0+1];
    xp cyp0 = ypcount[r0],cyp1 = ypcount[r0+1];

    // send in x+ direction
    parallel {
      ∀j ∈ [0,cxm1) pos1[fxm1+j] <- s_xp[];
      {
	∀j ∈ [0,cxm0) s_xp[] <- pos1[fxm0+j];
	∀ic s_xp[] <- control(xycmds[ic]);
      }
    }

    // send in x- direction
    parallel {
      ∀j ∈ [0,cxp1) pos1[fxp1+j] <- s_xm[];
      {
	∀j ∈ [0,cxp0) s_xm[] <- pos1[fxp0+j];
	∀ic s_xm[] <- control(xycmds[ic]);
      }
    }
    
    // send in y+ direction
    parallel {
      ∀j ∈ [0,cym1) pos1[fym1+j] <- s_yp[];
      {
	∀j ∈ [0,cym0) s_yp[] <- pos1[fym0+j];
	∀ic s_yp[] <- control(xycmds[ic]);
      }
    }

    // send in y- direction
    parallel {
      ∀j ∈ [0,cyp1) pos1[fyp1+j] <- s_ym[];
      {
	∀j ∈ [0,cyp0) s_ym[] <- pos1[fyp0+j];
	∀ic s_ym[] <- control(xycmds[ic]);
      }
    }

    r0 += 1;
  }

  // send in z- direction
  {
    xp r = 0;
    while(r < range) { 
      parallel {
          {
            xp f1 = zpfirst[r+1];
            xp c1 = zpcount[r+1];
	        ∀j ∈ [0,c1) pos1[f1+j] <- s_zm[];
          }
	{
      xp f0 = zpfirst[r];
      xp c0 = zpcount[r];
	  ∀j ∈ [0,c0) s_zm[] <- pos1[f0+j];
	  ∀ic s_zm[] <- control(zcmds[ic]);
	}
      }
      r <+- 1;
    }
  }
}


function reduce() {
  let ic in [0,2);
  xp r,f0,f1,c0,c1;

  xp const // Debug flags. A 'real' run have them all set to 1.
    zpflag = 1,
    xpflag = 1,
    xmflag = 1,
    ypflag = 1,
    ymflag = 1;
  
  // send in z+ direction
  if(zpflag > 0) {
    r = range-1;
    while(r >= 0) {
      f0 = zpfirst[r+1]; f1 = zpfirst[r];
      c0 = zpcount[r+1]; c1 = zpcount[r];
      parallel {
	∀j ∈ [0,c1) ff1[f1+j] <+- s_zp[];
	{
	  ∀j ∈ [0,c0) s_zp[] <- ff1[f0+j];
	  ∀ic s_zp[] <- control(zcmds[ic]);
	}
      }
      r -= 1;
    }
  }

  // Reduction along x+
  if(xpflag > 0) {
    r = range-1;
    while(r >= 0) {
      f0 = xpfirst[r+1]; f1 = xpfirst[r];
      c0 = xpcount[r+1]; c1 = xpcount[r];
      parallel {
	∀j ∈ [0,c1) ff1[f1+j] <+- s_xp[];
	{
	  ∀j ∈ [0,c0) s_xp[] <- ff1[f0+j];
	  ∀ic s_xp[] <- control(xycmds[ic]);
	}
      }
      r -= 1;
    }
  }

  // Reduction along x-
  if(xmflag > 0) {
    r = range-1;
    while(r >= 0) {
      f0 = xmfirst[r+1]; f1 = xmfirst[r];
      c0 = xmcount[r+1]; c1 = xmcount[r];
      parallel {
	∀j ∈ [0,c1) ff1[f1+j] <+- s_xm[];
	{
	  ∀j ∈ [0,c0) s_xm[] <- ff1[f0+j];
	  ∀ic s_xm[] <- control(xycmds[ic]);
	}
      }
      r -= 1;
    }
  }
  
  // Reduction along y+
  if(ypflag > 0) {
    r = range-1;
    while(r >= 0) {
      f0 = ypfirst[r+1]; f1 = ypfirst[r];
      c0 = ypcount[r+1]; c1 = ypcount[r];
      parallel {
	∀j ∈ [0,c1) ff1[f1+j] <+- s_yp[];
	{
	  ∀j ∈ [0,c0) s_yp[] <- ff1[f0+j];
	  ∀ic s_yp[] <- control(xycmds[ic]);
	}
      }
      r -= 1;
    }
  }
  
  // Reduction along y-
  if(ymflag > 0) {
    r = range-1;
    while(r >= 0) {
      f0 = ymfirst[r+1]; f1 = ymfirst[r];
      c0 = ymcount[r+1]; c1 = ymcount[r];
      parallel {
	∀j ∈ [0,c1) ff1[f1+j] <+- s_ym[];
	{
	  ∀j ∈ [0,c0) s_ym[] <- ff1[f0+j];
	  ∀ic s_ym[] <- control(xycmds[ic]);
	}
      }
      r -= 1;
    }
  }

}

function remap_positions() {
  sp one = 1.0;
  // Convert floating point id to integer type.
  let i0 in [ymptr,zend); 
  { ∀i0 tmp1[i0] = __builtin_copysign(one,pos[i0][3]); }
  { ∀i0 atomtype[i0] = tmp1[i0]; }
  
  // Center all coordinates on current box center. Then no
  // PBC need to be applied in distance computations.
  /*
  //∀k in [0,3) {
  xp k = 0;
  let i in [ymptr,zend);
  while(k < 3) {
    sp c = center[k];
    sp bi = boxinv[k];
    sp b = box[k];
    { ∀i pos[i][k] -= c; }
    { ∀i tmp1[i] = pos[i][k]*bi; }
    { ∀i tmp2[i] = __builtin_copysign(convmagic,tmp1[i]); }
    { ∀i tmp1[i] += tmp2[i]; }
    { ∀i tmp1[i] -= tmp2[i]; }
    { ∀i tmp1[i] *= b; }
    { ∀i pos[i][k] -= tmp1[i]; }
    k++;
  }
  */
}

function compute_dist2() {
  let m in [0,nint);
  let k in [0,3);
  ∀m {
    xp idx = ilist[m];
    ∀k dxij[m][k] = pos[idx][k];
  }
  ∀m {
    xp idx = jlist[m];
    ∀k dxij[m][k] -= pos[idx][k];
  }
  ∀m dist2[m] = dxij[m][0]*dxij[m][0];
  ∀m tmp1[m] = dxij[m][1]*dxij[m][1];
  ∀m dist2[m] += tmp1[m];
  ∀m tmp1[m] = dxij[m][2]*dxij[m][2];
  ∀m dist2[m] += tmp1[m];
}


function detect_platelet() {
  sp cm;
  xp ptr = 0;
  icount1 = 0;
  icount2 = 0;
  xp i,j;
  ∀i ∈ [ymptr,yend) {
    xp itype = atomtype[i];
    ∀j ∈ [xpptr,zend) {
      sp dx,d2 = 0.0;
      ∀k ∈ [0,3) {
	dx = pos[i][k] - pos[j][k];
	sp imgx = dx*boxinv[k];
	// Compute imgx = round(imgx)
	cm = __builtin_copysign(convmagic,imgx);
	imgx += cm;
	imgx -= cm;
	imgx = imgx*box[k];
	dx -= imgx;
	d2 += dx*dx;
	dxij[ptr][k] = dx;
      }
      if(d2 < rcut2) {
	xp jtype = atomtype[j];
	dist2[ptr] = d2;
	ilist[ptr] = i;
	jlist[ptr] = j;
	klist[ptr] = itype + jtype;
	ptr += 1;
	if(ptr >= nintmax) {
	  ptr -= 1;
	}
      }
      icount1 += 1;
    }
  }

  ∀i ∈ [y0ptr,ypptr) {
    xp itype = atomtype[i];
    ∀j ∈ [i+1,yend) {
      sp dx,d2 = 0.0;
      ∀k ∈ [0,3) {
	dx = pos[i][k] - pos[j][k];
	sp imgx = dx*boxinv[k];
	cm = __builtin_copysign(convmagic,imgx);
	imgx += cm;
	imgx -= cm;
	imgx = imgx*box[k];
	dx -= imgx;
	d2 += dx*dx;
	dxij[ptr][k] = dx;
      }
      if(d2 < rcut2) {
	xp jtype = atomtype[j];
	dist2[ptr] = d2;
	ilist[ptr] = i;
	jlist[ptr] = j;
	klist[ptr] = itype + jtype;
	ptr += 1;
	if(ptr >= nintmax) {
	  ptr -= 1;
	}
      }
      icount2 += 1;
    }
  }

  nint = ptr;

  // On input klist[i] is -2, 0, or 2.
  {
    let i in [0,nint);
    let i1 in [0,nlocal);
    ∀i klist[i] += 2;
    ∀i klist[i] >>= 1;
    // Now klist[i] is 0, 1, or 2.
    //∀i klist[i] *= rhotablen; // Offset to select appropriate table
    // Build 'klist' for local atoms
    ∀i1 atomtype[y0ptr+i1] += 1; // Atoms at y0ptr are local
    ∀i1 atomtype[y0ptr+i1] >>= 1;
    //∀i1 klocal[i1] = atomtype[y0ptr+i1];
    //∀i1 klocal[i1] *= Fptablen;
  }

}

/*
function detect_platelet_old() {
  xp ptr = 0;
  icount1 = 0;
  icount2 = 0;

  let j in [xpptr,zend);
  ∀i ∈ [ymptr,yend) {
    xp dptr = ptr + 4; // So that copying from dxij to itself is
		       // separated by more than SIMD length
    xp ptr0 = ptr;
    sp pix = pos[i][0];
    { ∀j dxij[j+dptr][0] = pix - pos[j][0]; }
    { ∀j tmp1[j] = dxij[j+dptr][0] * dxij[j+dptr][0]; }
    pix = pos[i][1];
    { ∀j dxij[j+dptr][1] = pix - pos[j][1]; }
    { ∀j tmp2[j] = dxij[j+dptr][1] * dxij[j+dptr][1]; }
    { ∀j tmp1[j] += tmp2[j]; }
    pix = pos[i][2];
    { ∀j dxij[j+dptr][2] = pix - pos[j][2]; }
    { ∀j tmp2[j] = dxij[j+dptr][2] * dxij[j+dptr][2]; }
    { ∀j tmp1[j] += tmp2[j]; }
    ∀j {
      if(tmp1[j] < rcut2) {
	∀k in [0,3) dxij[ptr][k] = dxij[j+dptr][k];
	dist2[ptr] = tmp1[j];
	jlist[ptr] = j;
	klist[ptr] = atomtype[j];
	ptr += 1;
	if(ptr == nintmax) { ptr -= 1; }
      }
    }

    {
      let m in [ptr0,ptr);
      xp itype = atomtype[i];
      { ∀m ilist[m] = i; }
      { ∀m klist[m] += itype; }
    }
    icount1 += zend;
    icount1 -= xpptr;
  }

  ∀i ∈ [y0ptr,ypptr) {
    xp ip1 = i+1;
    let j in [ip1,yend);

    xp dptr = ptr + 4; // So that copying from dxij to itself is
		       // separated by more than SIMD length
    xp ptr0 = ptr;
    sp pix = pos[i][0];
    { ∀j dxij[j+dptr][0] = pix - pos[j][0]; }
    { ∀j tmp1[j] = dxij[j+dptr][0] * dxij[j+dptr][0]; }
    pix = pos[i][1];
    { ∀j dxij[j+dptr][1] = pix - pos[j][1]; }
    { ∀j tmp2[j] = dxij[j+dptr][1] * dxij[j+dptr][1]; }
    { ∀j tmp1[j] += tmp2[j]; }
    pix = pos[i][2];
    { ∀j dxij[j+dptr][2] = pix - pos[j][2]; }
    { ∀j tmp2[j] = dxij[j+dptr][2] * dxij[j+dptr][2]; }
    { ∀j tmp1[j] += tmp2[j]; }
    ∀j {
      if(tmp1[j] < rcut2) {
	∀k in [0,3) dxij[ptr][k] = dxij[j+dptr][k];
	dist2[ptr] = tmp1[j];
	jlist[ptr] = j;
	klist[ptr] = atomtype[j];
	ptr += 1;
	if(ptr == nintmax) { ptr -= 1; }
      }
    }
    {
      let m in [ptr0,ptr);
      xp itype = atomtype[i];
      { ∀m ilist[m] = i; }
      { ∀m klist[m] += itype; }
    }

    icount2 += yend;
    icount2 -= ip1;
  }

  nint = ptr;     
  // On input klist[i] is -2, 0, or 2.
  {
    let i in [0,nint);
    let i1 in [0,nlocal);
    ∀i klist[i] += 2;
    ∀i klist[i] >>= 1;
    // Now klist[i] is 0, 1, or 2.
    ∀i klist[i] *= rhotablen; // Offset to select appropriate table
    // Build 'klist' for local atoms
    ∀i1 atomtype[y0ptr+i1] += 1; // Atoms at y0ptr are local
    ∀i1 atomtype[y0ptr+i1] >>= 1;
    ∀i1 klocal[i1] = atomtype[y0ptr+i1];
    ∀i1 klocal[i1] *= Fptablen;
  }
}
*/

xp id2[nintmax][2] address(&rinv);
sp y2[nintmax] address(&tmp1);
sp xm2[nintmax] address(&tmp2);
function rsqrt_rinv() {
  // Computes rinv[i] := 1/sqrt(rinv[i]), i=0,1,...,nint-1.
  let j ∈ [0,nint);
  sp x,y;

  xp const magic = 0x5f37;
  //sp const mhalf = -0.5;
  ∀j xm2[j] = dist2[j]*mhalf;
  ∀j rinv[j] = dist2[j];
  //u.v[1] <- magic - (u.v[1] >> 1);
  ∀j id2[j][1] >>= 1;
  ∀j id2[j][1] = -id2[j][1];
  ∀j id2[j][1] += magic;
  // y <- y*(y*y*xm2 + 1.5); 3 iterations
  ∀j y2[j] = rinv[j]*rinv[j];
  ∀j y2[j] = y2[j]*xm2[j] + 1.5;
  ∀j rinv[j] *= y2[j];

  ∀j y2[j] = rinv[j]*rinv[j];
  ∀j y2[j] = y2[j]*xm2[j] + 1.5;
  ∀j rinv[j] *= y2[j];

  ∀j y2[j] = rinv[j]*rinv[j];
  ∀j y2[j] = y2[j]*xm2[j] + 1.5;
  ∀j rinv[j] *= y2[j];
}

function zeroforce() {
  let px ∈ [0,3*zend);
  sp zero;
  zero = 0.0;
  ∀px ff1[px] = zero; 
}

sp eval = 0.0,keval = 0.0;
function ljforce() {
  let m in [0,nint);

  ∀m tmp1[m] = rinv[m]*rinv[m]; // tmp1 = r^(-2)
  ∀m tmp2[m] = tmp1[m]*tmp1[m]; // tmp2 = r^(-4)
  ∀m tmp1[m] *= tmp2[m]; // tmp1 = r^(-6)
  ∀m tmp2[m] *= tmp2[m]; // tmp2 = r^(-8)

  // Compute energy E = (1/2)*0.1*sum_i sum_j 1/r>ij^12
  ∀m df[m] = tmp1[m]*tmp1[m]; // df = r^(-12)
  eval = 0.0;
  ∀m eval += df[m];
  eval *= 0.1;
  
  ∀m tmp1[m] *= tmp2[m]; // tmp1 = r^(-14)
  /*
  sp e = 0.0 - 24.0*fparam.ljeps;
  ∀m df[m] = 2.0*tmp1[m]; // df = 2/r^14
  ∀m df[m] -= tmp2[m];    // df = 2/r^12 - 1/r^8
  ∀m df[m] *= e;          // df = -4*ljeps*(12/r^14 - 6/r^8)
  */
  sp e = 0.1*12.0;
  ∀m df[m] = e*tmp1[m]; // df = 0.1*12/r^14 -- 0.1/r^12 repulsive potential
  
  ∀m dxij[m][0] *= df[m];
  ∀m dxij[m][1] *= df[m];
  ∀m dxij[m][2] *= df[m];

  ∀m {
    xp idx = ilist[m];
    ∀k in [0,3) ff[idx][k] += dxij[m][k];
  }  
  ∀m {
    xp idx = jlist[m];
    ∀k in [0,3) ff[idx][k] -= dxij[m][k];
  }  
}

/* Timer functions */
xp timeridx;
xp const log2ntime = 7;
xp const ntime = 1<<log2ntime;
xp const tmask = ntime - 1;
xp tstamp[n=ntime][m=3];
sp evec[n=ntime],kevec[n=ntime];

/*

xp timerreg_sdr[3] address(0x7b30);
xp timercontrolreg_sdr[1] address(0x7b34);
*/
xp const tra1 = 0x7b30;
xp const tra2 = 0x7b34;
xp timerreg[3] address(tra1);
xp timercontrolreg[1] address(tra2);

function timerinit() {
  xp t;
  //timerreg.base = tra1;
  //timercontrolreg.base = tra2;
  
  t = timercontrolreg[0];
  t |= 12;
  timercontrolreg[0] = t;
  timeridx = 0;
}
function recordtime() {
  let it ∈ [0,3);
  xp j;
  j = timeridx & tmask;
  ∀it tstamp[j][it] = timerreg[it];
  evec[j] = eval; // Save energy when time is recorded
  kevec[j] = keval; // Save energy when time is recorded
  timeridx += 1;
}


xp ptrs[n=22];
main {
  timerinit();
  initparams();
  addressinit();
  /* Set correct colors in x+,x-,y+,y- sockets */ {
    xp q = 0;
    while(q < 8) {
      xp c = get_iq_color(q);
      if(hix) {
	if(c == 2 || c == 0) {
	  set_iq_color(q,c+1);
	}
      }
      if(hiy) {
	if(c == 4 || c == 6) {
	  set_iq_color(q,c+1);
	}
      }

      c = get_oq_color(q);
      if(hix) {
	if(c == 2 || c == 0) {
	  set_oq_color(q,c+1);
	}
      }
      if(hiy) {
	if(c == 4 || c == 6) {
	  set_oq_color(q,c+1);
	}
      }
      c = get_oq_color(q);
      q += 1;
    }
    // Unset teardown to allow communication
    ∀c ∈ [0,8) {
      clear_teardown(c);
    }
  }

  ptrs[17] = 0;
  ptrs[18] = 0;
  ptrs[19] = 0;
  ptrs[20] = 0;

  recordtime();
  xp ninfo[5];
  xp nxm = 0,nxp = 0,nym = 0,nyp = 0,nzp = 0;
  sp niter = fparams.niter,ninner = fparams.ninner;
  sp iter = 0.0,inner = 0.0;
  xp nsample = params.nsample;
  xp isample = 0;
  xp const rebuild = 1,rebuildnext = 1;

  while(iter < niter) {
    inner = 0.0;
    while(inner <= ninner) {
      let j5 in [0,5);

      if(rebuild > 0) {
	∀j5 countblock[j5][0] = nlocal;
	
	getcount();
	
	∀j5 ninfo[j5] = 0;
	∀j5 firstblock[j5][0] = 0;
	∀r ∈ [1,range+1) {
	  ∀j5 firstblock[j5][r] = ninfo[j5];
	  ∀j5 ninfo[j5] += countblock[j5][r];
	}
	nxm = ninfo[0];
	nxp = ninfo[1];
	nym = ninfo[2];
	nyp = ninfo[3];
	nzp = ninfo[4];

	xp ninfo1[5];
	∀j5 ninfo1[j5] = 0;
	
	ptrs[0] = nlocal;
	∀j5 ptrs[j5+1] = ninfo1[j5] + ninfo[j5];
	
	// Set up pointers to received atoms
	ymptr = nlocal;
	y0ptr = ymptr + nym;
	ypptr = y0ptr + nlocal;
	yend  = ypptr + nyp;
	xmptr = yend;
	x0ptr = xmptr + nxm;
	xpptr = x0ptr + nlocal;
	xend  = xpptr + nxp;
	zpptr = xend;
	zend  = zpptr + nzp;
	
	ptrs[ 6] = ymptr;
	ptrs[ 7] = y0ptr;
	ptrs[ 8] = ypptr;
	ptrs[ 9] = xmptr;
	ptrs[10] = x0ptr;
	ptrs[11] = xpptr;
	ptrs[12] = zpptr;
	ptrs[13] = zend;
	
	zpfirst[0] = xmptr;
	
	{
	  let r1 ∈ [0,range+1);
	  let r2 ∈ [1,range+1);
	  ∀r2 xmfirst[r2] += xmptr;
	  ∀r2 xpfirst[r2] += xpptr;
	  ∀r2 ymfirst[r2] += ymptr;
	  ∀r2 ypfirst[r2] += ypptr;
	  ∀r2 zpfirst[r2] += zpptr;
	  
	  ∀j5 ∀r1 count1[j5][r1] = countblock[j5][r1];
	  ∀j5 ∀r1 first1[j5][r1] = firstblock[j5][r1];
	  ∀j5 ∀r1 count3[j5][r1] = countblock[j5][r1] * 3;
	  ∀j5 ∀r1 first3[j5][r1] = firstblock[j5][r1] * 3;
	  ∀j5 ∀r1 count4[j5][r1] = countblock[j5][r1] << 2;
	  ∀j5 ∀r1 first4[j5][r1] = firstblock[j5][r1] << 2;
	}
      }
      
      /* Copy local atoms into x0 and y0 */ {
	xp n4,x0ptr4,y0ptr4;
	n4 = nlocal << 2;
	x0ptr4 = x0ptr << 2;
	y0ptr4 = y0ptr << 2;
	let j ∈ [0,n4);
	∀j pos1[x0ptr4+j] = pos1[j];
	∀j pos1[y0ptr4+j] = pos1[j];
      }
      // Set pointer for position (and id) exchange
      ∀j5 ∀r in [0,countsize) countblock[j5][r] = count4[j5][r];
      ∀j5 ∀r in [0,firstsize) firstblock[j5][r] = first4[j5][r];
      
      exchange();
      
      remap_positions();
      if(rebuild > 0) {
	nint = 0;
	detect_platelet();
      } // else { calc_dist2(); } // Needed if rebuild is ever 0
      
      rsqrt_rinv();
      zeroforce();
      ljforce();
      
      ∀j5 ∀r in [0,countsize) countblock[j5][r] = count3[j5][r];
      ∀j5 ∀r in [0,firstsize) firstblock[j5][r] = first3[j5][r];
      
      reduce();
      
      /* Add forces from y0 to local atoms */ {
	xp n3,x0ptr3,y0ptr3;
	n3 = 3*nlocal;
	x0ptr3 = 3*x0ptr;
	y0ptr3 = 3*y0ptr;
	let j ∈ [0,n3);
	∀j ff1[j] += ff1[x0ptr3+j];
	∀j ff1[j] += ff1[y0ptr3+j];
      }
      
      if(ninner > 0.0) { // Integrate
	xp n3 = nlocal*3;
	let j3 in [0,n3);
	let j in [0,nlocal);
	sp dtv = dt;
	if(inner == 0.0) {
	  dtv = 0.5*dt;
	} else if(inner == ninner) {
	  dtv = 0.5*dt;
	}
	
	∀j {
	  xp jtype = atomtype[y0ptr+j];
	  tmp1[j] = invmass[jtype];
	}
	∀j tmp1[j] *= dtv;
	/*
	  ∀j ff[x0ptr+j][0] = ff[j][0]*tmp1[j];
	  ∀j ff[x0ptr+j][1] = ff[j][1]*tmp1[j];
	  ∀j ff[x0ptr+j][2] = ff[j][2]*tmp1[j];
	  ∀j3 vel1[j3] += ff[x0ptr][j3];
	*/
	
	∀j tmp2[j] = ff[j][0]*tmp1[j];
	∀j vel[j][0] += tmp2[j];
	∀j tmp2[j] = ff[j][1]*tmp1[j];
	∀j vel[j][1] += tmp2[j];
	∀j tmp2[j] = ff[j][2]*tmp1[j];
	∀j vel[j][2] += tmp2[j];
	
	if(inner < ninner) {
	  ∀j pos[j][0] += vel[j][0]*dt;
	  ∀j pos[j][1] += vel[j][1]*dt;
	  ∀j pos[j][2] += vel[j][2]*dt;
	}
      }
    
      if(rebuildnext > 0) { // Atom migration
	xp k = 0;
	xp amflags[3] = {1,1,1};
	xp nloc0,j;
	xp xmrecv,xprecv;
	let ic in [0,2);
	while(k < 3) {
	  if(amflags[k] > 0) { // Atom migration, k direction
	    nloc0 = 0;
	    j = 0;
	    xmptr = nlocal;
	    xpptr = 3*nlocal;
	    nxm = 0;
	    nxp = 0;
	    while(j < nlocal) {
	      sp dx = pos[j][k],imgx;
	      sp cx = center[k];
	      sp cm;
	      dx -= cx;
	      imgx = dx*boxinv[k];
	      cm = __builtin_copysign(convmagic,imgx);
	      imgx += cm;
	      imgx -= cm;
	      imgx = imgx*box[k];
	      dx -= imgx;
	      if(dx < cell[k][0]) {
		∀m in [0,4) pos[xmptr+2*nxm][m] = pos[j][m];
		∀m in [0,3) pos[xmptr+2*nxm+1][m] = vel[j][m];
		nxm += 1;
	      } else if(dx > cell[k][1]) {
		∀m in [0,4) pos[xpptr+2*nxp][m] = pos[j][m];
		∀m in [0,3) pos[xpptr+2*nxp+1][m] = vel[j][m];
		nxp += 1;
	      } else {
		∀m in [0,4) pos[nloc0][m] = pos[j][m];
		∀m in [0,3) vel[nloc0][m] = vel[j][m];
		nloc0 += 1;
	      }
	      j += 1;
	    }
	    nlocal = nloc0;

	    //#pragma uthread enable off	    
	    { // Exchange counts and data
	      sp tmp1,tmp2;
	      if(k == 0) {
		s_xm[] <- (sp) nxm;
		∀ic s_xm[] <- control(xycmds[ic]);
		s_xp[] <- (sp) nxp;
		∀ic s_xp[] <- control(xycmds[ic]);
		//xmrecv <- s_xm[];
		//xprecv <- s_xp[];
		tmp1 <- s_xm[];
		tmp2 <- s_xp[];
		xmrecv = tmp1;
		xprecv = tmp2;
	      } else if(k == 1) {
		s_ym[] <- (sp) nxm;
		∀ic s_ym[] <- control(xycmds[ic]);
		s_yp[] <- (sp) nxp;
		∀ic s_yp[] <- control(xycmds[ic]);
		//xmrecv <- s_ym[];
		//xprecv <- s_yp[];
		tmp1 <- s_ym[];
		tmp2 <- s_yp[];
		xmrecv <- tmp1;
		xprecv <- tmp2;
	      } else {
		s_zm[] <- (sp) nxm;
		∀ic s_zm[] <- control(zcmds[ic]);
		s_zp[] <- (sp) nxp;
		∀ic s_zp[] <- control(zcmds[ic]);
		//xmrecv <- s_zm[];
		//xprecv <- s_zp[];
		tmp1 <- s_zm[];
		tmp2 <- s_zp[];
		xmrecv <- tmp1;
		xprecv <- tmp2;
	      }
	      
	      ymptr = xpptr + 2*nxp;
	      ypptr = ymptr + 2*xmrecv;
	      
	      xp
		nxm8 = 8*nxm,
		nxp8 = 8*nxp,
		xmrecv8 = 8*xmrecv,
		xprecv8 = 8*xprecv;
	      
	      if(k == 0) {
		parallel {
		  ∀j in [0,xmrecv8) pos[ymptr][j] <- s_xm[];
		  ∀j in [0,xprecv8) pos[ypptr][j] <- s_xp[];
		  {
		    ∀j in [0,nxm8) s_xm[] <- pos[xmptr][j];
		    ∀ic s_xm[] <- control(xycmds[ic]);
		  }
		  {
		    ∀j in [0,nxp8) s_xp[] <- pos[xpptr][j];
		    ∀ic s_xp[] <- control(xycmds[ic]);
		  }
		}
	      } else if(k == 1) {
		parallel {
		  ∀j in [0,xmrecv8) pos[ymptr][j] <- s_ym[];
		  ∀j in [0,xprecv8) pos[ypptr][j] <- s_yp[];
		  {
		    ∀j in [0,nxm8) s_ym[] <- pos[xmptr][j];
		    ∀ic s_ym[] <- control(xycmds[ic]);
		  }
		  {
		    ∀j in [0,nxp8) s_yp[] <- pos[xpptr][j];
		    ∀ic s_yp[] <- control(xycmds[ic]);
		  }
		}	
	      } else {
		parallel {
		  ∀j in [0,xmrecv8) pos[ymptr][j] <- s_zm[];
		  ∀j in [0,xprecv8) pos[ypptr][j] <- s_zp[];
		  {
		    ∀j in [0,nxm8) s_zm[] <- pos[xmptr][j];
		    ∀ic s_zm[] <- control(zcmds[ic]);
		  }
		  {
		    ∀j in [0,nxp8) s_zp[] <- pos[xpptr][j];
		    ∀ic s_zp[] <- control(zcmds[ic]);
		  }
		}
	      }
	    }
	    
	    { // Copy data from ymptr and ypptr to after nlocal
	      j = 0;
	      while(j < xmrecv) {
		∀kk in [0,4) pos[nlocal+j][kk] = pos[ymptr+2*j][kk];
		∀kk in [0,3) vel[nlocal+j][kk] = pos[ymptr+2*j+1][kk];
		j += 1;
	      }
	      nlocal += xmrecv;
	      j = 0;
	      while(j < xprecv) {
		∀kk in [0,4) pos[nlocal+j][kk] = pos[ypptr+2*j][kk];
		∀kk in [0,3) vel[nlocal+j][kk] = pos[ypptr+2*j+1][kk];
		j += 1;
	      }
	      nlocal += xprecv;
	    }
	  }
	  k += 1;
	}
      }
      
      ptrs[0] = nlocal;
      ptrs[14] = nint;
      ptrs[15] = icount1;
      ptrs[16] = icount2;
      ptrs[17] += 1; // Save iteration count
      // Max nlocal so far
      if(nlocal > ptrs[18]) { ptrs[18] = nlocal; }
      // Max nlocal+received atoms so far
      if(zend > ptrs[19]) { ptrs[19] = zend; }
      // Max nint so far
      if(nint > ptrs[20]) { ptrs[20] = nint; }
      params.nlocal = nlocal;
    
      inner += 1.0;
    }

    isample += 1;
    if(isample == nsample) {
      sp junk;

      xp n3 = 3*nlocal;
      keval = 0.0;
      ∀m in [0,n3) tmp1[m] = vel1[m]*vel1[m];
      ∀m in [0,n3) keval += tmp1[m];
      keval *= 0.5;

      recordtime();
      isample = 0;
      
      ping[] <- control(adv[0]); // Open channel to signal timestep done.
      junk <- ping[];         // Discard probe data.
      junk <- ping[];         // Ready to go when receive completes.
    }
    iter += 1.0;
  }
}