socket right,up,down,left;

import timer;

// Routing plan (interior, computing cores):
//   c1: C>R
//   c2: C>U
//   c3: D>C
//   c4: L>C
//   anti-red:   R>L
//   anti-blue:  R>L
//   anti-green:  U>D
//   anti-yellow: U>D
//
//   Core coord x,y:
//     c1 = (x%2 == 0) ? red : blue
//     c2 = (y%2 == 0) ? green : yellow
//     c3 = (y%2 == 0) ? blue : red
//     c4 = (x%2 == 0) ? yellow : green
//
//   Color swapping moats:
//      left,right: horizontal red <-> anti-red, blue <-> anti-blue
//      top,bottom: vertical green <-> anti-green, yellow <-> anti-yellow
//
//   This means that if I send out the right edge, color swaps, goes
//   all the way back to the left edge, swaps again, and enters first
//   column on the same color it was sent out over the right edge.

xp param n,p,pad;
xp const w = 2*p;
xp const pm1 = p - 1;

sp fpmem[2*(n+w)*(n+w+pad)+2];
sp x[n=n+w][m=n+w+pad] extern address(&fpmem);
sp y[n=n+w][m=n+w+pad] extern address(&x+2*((n+w)*(n+w+pad)+1));

struct paramstruct {
  sp origin;
  sp h;  // grid spacing
  sp dt; // time step
  sp D;  // diffusion coefficient
  sp niter,ninner,nsample;
};
struct paramstruct params;

sp a1,a11,a4;
sp avec[9];

function setweights() {
  sp h,dt,D,hinv,alpha;
  h  = params.h;
  dt = params.dt;
  D  = params.D;

  hinv = 1.0/h;
  alpha = 0.5*D*dt*hinv*hinv;
  a1 = alpha;
  a11 = 0.5*alpha;
  a4 = 1.0 - 6.0*alpha;
  avec[0] = a11;  avec[1] = a1;  avec[2] = a11;
  avec[3] = a1;   avec[4] = a4;  avec[5] = a1;
  avec[6] = a11;  avec[7] = a1;  avec[8] = a11;
}

function sendrecv() {
  let ir ∈ [w,n+w);
  let jr ∈ [n,n+w);
  let iu ∈ [n,n+w);
  let ju ∈ [w,n+w);
  let iul ∈ [n,n+w);
  let jul ∈ [0,w);

  let il ∈ [w,n+w);
  let jl ∈ [0,w);
  let id ∈ [0,w);
  let jd ∈ [w,n+w);
  let idl ∈ [0,w);
  let jdl ∈ [0,w);

  parallel {
    ∀ir ∀jr right[] <- x[ir][jr];
    ∀iu ∀ju up[] <- x[iu][ju];
    ∀il ∀jl x[il][jl] <- left[];
    ∀id ∀jd x[id][jd] <- down[];
  }

  parallel {
    ∀iul ∀jul up[] <- x[iul][jul];
    ∀idl ∀jdl x[idl][jdl] <- down[];
  }
}

function compute() {
  let i ∈ [0,n);
  let j ∈ [0,n);

  ∀i ∀j y[i+w][j+w] = 0.0;
  ∀di in [0,3) ∀dj in [0,3) {
    sp a = avec[3*di + dj];
    ∀i ∀j y[i+w][j+w] += a*x[i + di + pm1][j + dj + pm1];
  }
      
  ∀i ∀j x[i+w][j+w] = y[i+w][j+w];
}

function initfield() {
  /* Compute initial conditions */
  let ii ∈ [0,n+w);
  let jj ∈ [0,n+w);
  let i1 ∈ [w,w+1);
  let j1 ∈ [w,w+1);
  sp oval = params.origin;
  ∀ii ∀jj x[ii][jj] = 0.0;
  ∀i1 ∀j1 x[i1][j1] = 1.0*oval;
}

main {
  setweights();
  initfield();
  timerinit();
  
  sp
    niter = params.niter,
    kmax = params.ninner,
    nsample = params.nsample,
    iter = 0.0,
    isample = 0.0;

  while(iter < niter) {
    sp k = 0.0;
    while(k < kmax) {
      sendrecv();
      compute();
      k += 1.0;
    }
    isample += 1.0;
    if(isample == nsample) {
      recordtime();
      isample = 0.0;
    }
    iter += 1.0;
  }

}