// Types

type range Range3D;
type partition <Range3D> PartRange3D [PartNP];
type swap <Range3D> SwapRange3D [PartNP];

type data Setup, Checksums, Int, Bool, Cplx;
type data Vector3Dcplx [PartRange3D];
type data <Cplx> ArrayCplx [PartNP];


// Init Data

const Setup sp[ALL1], Int niter[ALL1], fdir[ALL1], finv[ALL1], lay1d[ALL1];
const PartRange3D dom0[ALLn], dom1[ALLn], dom2[ALLn];
const SwapRange3D sw01[PPEn], sw12[PPEn], sw21[PPEn], sw10[PPEn], sw02[PPEn], sw20[PPEn];


// Decomposition Rules

part sw01 : dom0 -> dom1;
part sw12 : dom1 -> dom2;
part sw21 : dom2 -> dom1;
part sw10 : dom1 -> dom0;
part sw02 : dom0 -> dom2;
part sw20 : dom2 -> dom0;


// Parallel Computations

function fft (A -> D)
  Vector3Dcplx A[dom0], D[dom2]
{
  cffts1 (sp, fdir, A[dom0] -> B[dom0]);

  if (lay1d, B[dom0] -> C[dom2]) { 
    // 1D
    cffts2 (sp, fdir, B[dom0] -> C[dom0]);
  } else { 
    // 2D
    cffts2 (sp, fdir, B[dom1] -> C[dom1]); 
  }

  cffts3 (sp, fdir, C[dom2] -> D[dom2]);
}

function ifft (A -> D)
  Vector3Dcplx A[dom2], D[dom0]
{
  cffts3 (sp, finv, A[dom2] -> B[dom2]);

  if (lay1d, B[dom2] -> C[dom0]) {
    // 1D
    cffts2 (sp, finv, B[dom0] -> C[dom0]);
  } else {
    // 2D
    cffts2 (sp, finv, B[dom1] -> C[dom1]);
  }

  cffts1 (sp, finv, C[dom0] -> D[dom0]);
}

function checksum (iter, x, cksum, dom -> cksum2)
  Int iter[ALL1], PartRange3D dom[ALLn],
  Vector3Dcplx x[dom], Checksums cksum[ONE1], cksum2[ONE1]
{
  checksum_dom <ALL1> (x[dom] -> c1[PPEn]);
  sumproc      <ONE1> (c1[ONEn] -> c2[ONE1]);
  checksum_set <ONE1> (sp, c2[ONE1], iter, cksum[ONE1] -> cksum2[ONE1]);
}

function mainloop ( -> success) [gen]
  Bool success[ALL1]
{
  init_checksums (sp -> cksum);
  compute_init_cond (sp -> u0[dom0]);

  fft (u0[dom0] -> u1[dom2]);

  loop (iter, cond; cksum[ONE1] -> cksum2[ONE1])
  {
    evolve (sp, iter, u1[dom2] -> u2[dom2]);

    ifft (u2[dom2] -> u3[dom0]);

    checksum (iter, u3[dom0], cksum -> cksum2);
    econd (niter, iter -> cond);
  }

  verify <ONE1> (sp, cksum2[ONE1] -> success[ONE1]);
}



// Sequential Computations

function evolve (sp, idx, u0, dom -> u1)
  Setup sp, Int idx, Range3D dom,
  Vector3Dcplx u0[dom], u1[dom];

function cffts1 (sp, dir, A, dom -> A*)
  Setup sp, Int dir, Range3D dom, Vector3Dcplx A[dom], A*[dom];

function cffts2 (sp, dir, A, dom -> A*)
  Setup sp, Int dir, Range3D dom, Vector3Dcplx A[dom], A*[dom];

function cffts3 (sp, dir, A, dom -> A*)
  Setup sp, Int dir, Range3D dom, Vector3Dcplx A[dom], A*[dom];


function sumproc (arr, ra -> res)
  ArrayCplx arr[ra], RangeNP ra, Cplx res;

function checksum_dom (x, dom -> res)
  Vector3Dcplx x[dom], Range3D dom, Cplx res;

function checksum_set (sp, res, iter, cksum -> cksum*)
  Setup sp, Cplx res, Int iter, Checksums cksum, cksum*;


function init_checksums (sp -> cks)
  Setup sp, Checksums cks;

function compute_init_cond (sp, dom -> u0)
  Setup sp, Vector3Dcplx u0[dom], Range3D dom;

function econd (max, idx -> cond)
  Int max, Int idx, cond;

function verify (sp, cks -> success)
  Setup sp, Checksums cks, Bool success;