/* Distributed under the Apache License, Version 2.0.
See accompanying NOTICE file for details.*/

#include "LBM.h"
#include "utils/TimingProfile.h"

// Boundary condition codes
#define BCinterior  0
#define BCinactive -1
#define WALL        0
#define mnWALLN     1
#define mxWALLN     99
#define PHYSICSBC   100
#define IMPOSED_INFLOW 100
#define IMPOSED_OUTFLOW 200

// Constants
// Buffer size for x-direction streaming
#define mL 512
// Tile size for y,z-directions streaming
#define mB 16
// sqrt(2.)
#define SQRT_2 0.70710678118654746

static int DeviceMB = 0; // Max device memory in MB (16GB), should be replace with appropriate CUDA DeviceInfo call

extern "C"
LBM_DECL bool cuda_device_check(Logger& log, int use_device)
{
  int nDevices;
  size_t nMostMemory = 0;
  int nBiggestDevice = -1;
  std::string sBiggestDevice;
  std::string sUsingDevice;

  cudaError_t err = cudaGetDeviceCount(&nDevices);
  if (err != cudaSuccess || nDevices < 0)
  {
    log.Info(cudaGetErrorString(err));
    log.Info("No CUDA compatible device found.");
    return false;
  }

  for (int i = 0; i < nDevices; i++)
  {
    cudaDeviceProp prop;
    cudaGetDeviceProperties(&prop, i);
    log.Info("CUDA Device Number: " + std::to_string(i));
    log.Info("  Device name: " + std::string(prop.name));
    log.Info("  Total Global Memory (Gb): " + std::to_string(prop.totalGlobalMem / 1e+9));
    log.Info("  Shared Memory Per Block (bytes): " + std::to_string(prop.sharedMemPerBlock));
    log.Info("  Memory Clock Rate (KHz): " + std::to_string(prop.memoryClockRate));
    log.Info("  Memory Bus Width (bits):" + std::to_string(prop.memoryBusWidth));
    log.Info("  Peak Memory Bandwidth (GB/s): " + std::to_string(
      2.0 * prop.memoryClockRate * (prop.memoryBusWidth / 8) / 1.0e6));

    if (i == use_device)
    {
      sUsingDevice = prop.name;
      DeviceMB = int(prop.totalGlobalMem / 1e+6);
    }
    if (prop.totalGlobalMem > nMostMemory)
    {
      nBiggestDevice = i;
      sBiggestDevice = prop.name;
      nMostMemory = prop.totalGlobalMem;
    }
  }

  if (use_device >= nDevices)
  {
    log.Warning("Requested device not found.");
    use_device = -1;
  }
  if (use_device < 0)
  {
    use_device = nBiggestDevice;
    sUsingDevice = sBiggestDevice;
    DeviceMB = int(nMostMemory / 1e+6);
    log.Info("Automatically using the largest device found.");
  }
  cudaSetDevice(use_device);
  log.Info("Running CUDA on device " + std::to_string(use_device) + ", a " + sUsingDevice);

  return true;
}

// This is the magnitude and direction of the velocity sets
__constant__ __device__ float VelocitySets[19][3] = { { 0.,0.,0. },{ 0.,0.,1. },{ 0.,0.,-1. },{ 0.,1.,0. },{ 0.,-1.,0. },{ 1.,0.,0. },{ -1.,0.,0. },
{ 0.,SQRT_2,SQRT_2 },{ 0.,-SQRT_2,-SQRT_2 },{ 0.,SQRT_2,-SQRT_2 },{ 0.,-SQRT_2,SQRT_2 },{ SQRT_2,SQRT_2,0. },{ -SQRT_2,-SQRT_2,0. },
{ SQRT_2,-SQRT_2,0. },{ -SQRT_2,SQRT_2,0. },{ SQRT_2,0.,SQRT_2 },{ -SQRT_2,0.,-SQRT_2 },{ SQRT_2,0.,-SQRT_2 },{ -SQRT_2,0.,SQRT_2 } };

// Used Everywhere
__device__ void Device_ComputeDistributionFunctions(float* pressure_flow_stress, float* distribution_functions)
{
  float p, u, v, w, U32, eu, rho;
  p = pressure_flow_stress[0];
  u = pressure_flow_stress[1];
  v = pressure_flow_stress[2];
  w = pressure_flow_stress[3];
  rho = 3 * p;
  U32 = 1.5 * (u * u + v * v + w * w);
  eu = 0.;     distribution_functions[0] = 0.333333333 * rho * (1 - U32);
  eu = u;      distribution_functions[1] = 0.055555556 * rho * (1 - U32 + (3. + 4.5 * eu) * eu);
  eu = -u;     distribution_functions[2] = 0.055555556 * rho * (1 - U32 + (3. + 4.5 * eu) * eu);
  eu = v;      distribution_functions[3] = 0.055555556 * rho * (1 - U32 + (3. + 4.5 * eu) * eu);
  eu = v + u;  distribution_functions[4] = 0.027777778 * rho * (1 - U32 + (3. + 4.5 * eu) * eu);
  eu = v - u;  distribution_functions[5] = 0.027777778 * rho * (1 - U32 + (3. + 4.5 * eu) * eu);
  eu = -v;     distribution_functions[6] = 0.055555556 * rho * (1 - U32 + (3. + 4.5 * eu) * eu);
  eu = -v + u; distribution_functions[7] = 0.027777778 * rho * (1 - U32 + (3. + 4.5 * eu) * eu);
  eu = -v - u; distribution_functions[8] = 0.027777778 * rho * (1 - U32 + (3. + 4.5 * eu) * eu);
  eu = w;      distribution_functions[9] = 0.055555556 * rho * (1 - U32 + (3. + 4.5 * eu) * eu);
  eu = w + u;  distribution_functions[10] = 0.027777778 * rho * (1 - U32 + (3. + 4.5 * eu) * eu);
  eu = w - u;  distribution_functions[11] = 0.027777778 * rho * (1 - U32 + (3. + 4.5 * eu) * eu);
  eu = w + v;  distribution_functions[12] = 0.027777778 * rho * (1 - U32 + (3. + 4.5 * eu) * eu);
  eu = w - v;  distribution_functions[13] = 0.027777778 * rho * (1 - U32 + (3. + 4.5 * eu) * eu);
  eu = -w;     distribution_functions[14] = 0.055555556 * rho * (1 - U32 + (3. + 4.5 * eu) * eu);
  eu = -w + u; distribution_functions[15] = 0.027777778 * rho * (1 - U32 + (3. + 4.5 * eu) * eu);
  eu = -w - u; distribution_functions[16] = 0.027777778 * rho * (1 - U32 + (3. + 4.5 * eu) * eu);
  eu = -w + v; distribution_functions[17] = 0.027777778 * rho * (1 - U32 + (3. + 4.5 * eu) * eu);
  eu = -w - v; distribution_functions[18] = 0.027777778 * rho * (1 - U32 + (3. + 4.5 * eu) * eu);
//    eu = 0.;     distribution_functions[0] = 0.333333333 * (rho - U32);
//    eu = u;      distribution_functions[1] = 0.055555556 * (rho - U32 + (3. + 4.5 * eu) * eu);
//    eu = -u;     distribution_functions[2] = 0.055555556 * (rho - U32 + (3. + 4.5 * eu) * eu);
//    eu = v;      distribution_functions[3] = 0.055555556 * (rho - U32 + (3. + 4.5 * eu) * eu);
//    eu = v + u;  distribution_functions[4] = 0.027777778 * (rho - U32 + (3. + 4.5 * eu) * eu);
//    eu = v - u;  distribution_functions[5] = 0.027777778 * (rho - U32 + (3. + 4.5 * eu) * eu);
//    eu = -v;     distribution_functions[6] = 0.055555556 * (rho - U32 + (3. + 4.5 * eu) * eu);
//    eu = -v + u; distribution_functions[7] = 0.027777778 * (rho - U32 + (3. + 4.5 * eu) * eu);
//    eu = -v - u; distribution_functions[8] = 0.027777778 * (rho - U32 + (3. + 4.5 * eu) * eu);
//    eu = w;      distribution_functions[9] = 0.055555556 * (rho - U32 + (3. + 4.5 * eu) * eu);
//    eu = w + u;  distribution_functions[10] = 0.027777778 * (rho - U32 + (3. + 4.5 * eu) * eu);
//    eu = w - u;  distribution_functions[11] = 0.027777778 * (rho - U32 + (3. + 4.5 * eu) * eu);
//    eu = w + v;  distribution_functions[12] = 0.027777778 * (rho - U32 + (3. + 4.5 * eu) * eu);
//    eu = w - v;  distribution_functions[13] = 0.027777778 * (rho - U32 + (3. + 4.5 * eu) * eu);
//    eu = -w;     distribution_functions[14] = 0.055555556 * (rho - U32 + (3. + 4.5 * eu) * eu);
//    eu = -w + u; distribution_functions[15] = 0.027777778 * (rho - U32 + (3. + 4.5 * eu) * eu);
//    eu = -w - u; distribution_functions[16] = 0.027777778 * (rho - U32 + (3. + 4.5 * eu) * eu);
//    eu = -w + v; distribution_functions[17] = 0.027777778 * (rho - U32 + (3. + 4.5 * eu) * eu);
//    eu = -w - v; distribution_functions[18] = 0.027777778 * (rho - U32 + (3. + 4.5 * eu) * eu);
  for (int n = 0; n < 19; n++)
  {
    if (distribution_functions[n] < 0.)
      distribution_functions[n] = 0.;
  }
}

__device__ void Device_ComputeTemperatureDistributionFunctions(float* pressure_flow_stress, float* distribution_functions)
{
    float p, u, v, w, U32, eu;
    p = pressure_flow_stress[0];
    u = pressure_flow_stress[1];
    v = pressure_flow_stress[2];
    w = pressure_flow_stress[3];
    U32 = 1.5 * (u * u + v * v + w * w);
    eu = 0.;     distribution_functions[0] = 0.333333333  * (p - U32);
    eu = u;      distribution_functions[1] = 0.055555556  * (p - U32 + (3. + 4.5 * eu) * eu);
    eu = -u;     distribution_functions[2] = 0.055555556  * (p - U32 + (3. + 4.5 * eu) * eu);
    eu = v;      distribution_functions[3] = 0.055555556  * (p - U32 + (3. + 4.5 * eu) * eu);
    eu = v + u;  distribution_functions[4] = 0.027777778  * (p - U32 + (3. + 4.5 * eu) * eu);
    eu = v - u;  distribution_functions[5] = 0.027777778  * (p - U32 + (3. + 4.5 * eu) * eu);
    eu = -v;     distribution_functions[6] = 0.055555556  * (p - U32 + (3. + 4.5 * eu) * eu);
    eu = -v + u; distribution_functions[7] = 0.027777778  * (p - U32 + (3. + 4.5 * eu) * eu);
    eu = -v - u; distribution_functions[8] = 0.027777778  * (p - U32 + (3. + 4.5 * eu) * eu);
    eu = w;      distribution_functions[9] = 0.055555556  * (p - U32 + (3. + 4.5 * eu) * eu);
    eu = w + u;  distribution_functions[10] = 0.027777778 * (p - U32 + (3. + 4.5 * eu) * eu);
    eu = w - u;  distribution_functions[11] = 0.027777778 * (p - U32 + (3. + 4.5 * eu) * eu);
    eu = w + v;  distribution_functions[12] = 0.027777778 * (p - U32 + (3. + 4.5 * eu) * eu);
    eu = w - v;  distribution_functions[13] = 0.027777778 * (p - U32 + (3. + 4.5 * eu) * eu);
    eu = -w;     distribution_functions[14] = 0.055555556 * (p - U32 + (3. + 4.5 * eu) * eu);
    eu = -w + u; distribution_functions[15] = 0.027777778 * (p - U32 + (3. + 4.5 * eu) * eu);
    eu = -w - u; distribution_functions[16] = 0.027777778 * (p - U32 + (3. + 4.5 * eu) * eu);
    eu = -w + v; distribution_functions[17] = 0.027777778 * (p - U32 + (3. + 4.5 * eu) * eu);
    eu = -w - v; distribution_functions[18] = 0.027777778 * (p - U32 + (3. + 4.5 * eu) * eu);
    for (int n = 0; n < 19; n++)
    {
        if (distribution_functions[n] < 0.)
            distribution_functions[n] = 0.;
    }
}

/////////////////////////////////////
// 0. Setup Distribution Functions //
/////////////////////////////////////

__global__ void InitializeDistributionFunctions(int nLBMmodel, int num_distribution_functions, int num_full_continuum_mechanics_components, int x_dimension, int y_dimension, int lattice_size, float* device_pressure_flow_stress, float* device_distribution_functions)
{
  int i = threadIdx.x; // Lattice i-index
  int j = blockIdx.x;  // Lattice j-index
  int k = blockIdx.y;  // Lattice k-index
  float distribution_functions[LBM::nPDFmodels];
  float pressure_flow_stress[LBM::nContmodels[1]];

  long int start_loc = i + x_dimension * (j + y_dimension * k);
  long int current_loc = start_loc;

  //Initialie the continuum mechanics
  for (int n = 0; n < num_full_continuum_mechanics_components; n++)
  {
    pressure_flow_stress[n] = device_pressure_flow_stress[current_loc];
    current_loc += lattice_size;
  }

  //Compute the distribution functions from the continuum mechanics
    switch (nLBMmodel)
    {
        case LBM::ModelType::Fluid:
        {
            Device_ComputeDistributionFunctions(pressure_flow_stress, distribution_functions);
            break;
        }
        case LBM::ModelType::Preconditioner:
        case LBM::ModelType::Thermal:
            Device_ComputeTemperatureDistributionFunctions(pressure_flow_stress, distribution_functions);
            break;
    }

  current_loc = start_loc;
  for (int n = 0; n < num_distribution_functions; n++)
  {
    device_distribution_functions[current_loc] = distribution_functions[n];
    current_loc += lattice_size;
  }
}
void SetupDistributionFunctions(int nLBMmodel, size_t* lattice_dimensions, int lattice_size, float* device_pressure_flow_stress, float* device_distribution_functions)
{
  dim3 dG, dB;
  // A CUDA block processes a line of data along x-axis, one thread per lattice site
  dB.x = lattice_dimensions[0];
  dB.y = 1;
  dB.z = 1;
  // The CUDA grid corresponds to dimensions along y and z
  dG.x = lattice_dimensions[1];
  dG.y = lattice_dimensions[2];

  int num_distribution_functions = LBM::nPDFmodels;
  int num_full_continuum_mechanics_components = LBM::nContmodels[nLBMmodel];
  InitializeDistributionFunctions << <dG, dB >> > (nLBMmodel, num_distribution_functions, num_full_continuum_mechanics_components, lattice_dimensions[0], lattice_dimensions[1], lattice_size, device_pressure_flow_stress, device_distribution_functions);

}
//Compute the advection-diffusion distribution functions
__device__ void ComputeAdvectionDiffusionDistributionFunctions(float* continuum_mech_values, float* distribution_functions)
{
  float p;
  p = continuum_mech_values[0];
  p = p > 0. ? p : 0.;
  distribution_functions[0] = 0.333333333 * p;
  distribution_functions[1] = 0.055555556 * p;
  distribution_functions[2] = 0.055555556 * p;
  distribution_functions[3] = 0.055555556 * p;
  distribution_functions[4] = 0.027777778 * p;
  distribution_functions[5] = 0.027777778 * p;
  distribution_functions[6] = 0.055555556 * p;
  distribution_functions[7] = 0.027777778 * p;
  distribution_functions[8] = 0.027777778 * p;
  distribution_functions[9] = 0.055555556 * p;
  distribution_functions[10] = 0.027777778 * p;
  distribution_functions[11] = 0.027777778 * p;
  distribution_functions[12] = 0.027777778 * p;
  distribution_functions[13] = 0.027777778 * p;
  distribution_functions[14] = 0.055555556 * p;
  distribution_functions[15] = 0.027777778 * p;
  distribution_functions[16] = 0.027777778 * p;
  distribution_functions[17] = 0.027777778 * p;
  distribution_functions[18] = 0.027777778 * p;
}
//Compute pressure from distribution functions
__device__ void ComputePressureFromDistributionFunctions(float* distribution_functions, float* continuum_mech_values)
{
  continuum_mech_values[0] = 0.;
  for (int n = 0; n < 19; n++)
  {
    continuum_mech_values[0] += distribution_functions[n];
  }
}

////////////////////////////////////////
// 1. & 3. Impose Boundary Conditions //
////////////////////////////////////////

__device__ void SetFlowBCDistributionFunctions(int nLBMmodel, float* previous_distribution_functions, float* distribution_functions, int bcnrml, float* bc_continuum_mech_values)
{
  float continuum_mech_values[LBM::nContmodels[1]];
  int nrml = bcnrml % 100;
  int bc_type = bcnrml - nrml;

  switch (nLBMmodel)
  {
  case LBM::ModelType::Fluid:
  {
    //Only have a pressure drop boundary condition

    switch (bc_type)  //inlet or outlet
    {
    case IMPOSED_INFLOW: //inlet
    {
      // 1. Equilibrium distribution function update
      //continuum_mech_values[0] = inlet pressure set;
      //continuum_mech_values[1] = calculated velocity at end;
      //continuum_mech_values[2] = calculated velocity at end;
      //continuum_mech_values[3] = calculated velocity at end;
      Device_ComputeDistributionFunctions(continuum_mech_values, distribution_functions);

      // 2. Non-equilibrium distribution function update

    }
    case IMPOSED_OUTFLOW: //outlet
    {
      // 1. Equilibrium distribution function update
      //continuum_mech_values[0] = outlet pressure set;
      //continuum_mech_values[1] = calculated velocity at location 1;
      //continuum_mech_values[2] = calculated velocity at lcoation 1;
      //continuum_mech_values[3] = calculated velocity at location 1;
      Device_ComputeDistributionFunctions(continuum_mech_values, distribution_functions);

      // 2. Non-equilibrium distribution function update
    }
    }

    int flow_direction = 1;
    switch (bc_type)
    {
    case IMPOSED_OUTFLOW:
    {
      flow_direction = -1;
    }
    case IMPOSED_INFLOW:
    {

      for (int n = 0; n < 4; n++)
      {
        continuum_mech_values[n] = bc_continuum_mech_values[n];
      }
      // q[0] contains desired pressure
      // q[1] contains desired normal velocity. Project on Cartesian components
      // vec[0] - z component, vec[1] - y component, vec[2] - x component
      continuum_mech_values[2] = continuum_mech_values[1] * VelocitySets[nrml][1] * flow_direction;
      continuum_mech_values[3] = continuum_mech_values[1] * VelocitySets[nrml][0] * flow_direction;
      continuum_mech_values[1] = continuum_mech_values[1] * VelocitySets[nrml][2] * flow_direction;
      Device_ComputeDistributionFunctions(continuum_mech_values, distribution_functions);
      break;
    }
    default:
    {
      printf("Warning: undefined flow BC: bcnrml=%d bc=%d nrml=%d\n", bcnrml, bc_type, nrml);
      for (int n = 0; n < 19; n++)
      {
        distribution_functions[n] = previous_distribution_functions[n];
      }
    }
    }
    break;
  }
  case LBM::ModelType::Preconditioner: //  Dirichlet boundary condition for either initial pressure approximation
  case LBM::ModelType::Thermal:        //  or steady-state advection-diffusion (e.g., temperature)
  {
    continuum_mech_values[0] = bc_continuum_mech_values[0];
    continuum_mech_values[1] = continuum_mech_values[2] = continuum_mech_values[3] = 0.; // Dirichlet boundary condition on q[0]
    ComputeAdvectionDiffusionDistributionFunctions(continuum_mech_values, distribution_functions);
    break;
  }
  }
}
// Modify F values from most recent time step (stored in previous_distribution_functions) to impose wall BC. Place result in distribution_funtions
__device__ void SetWallDistributionFunctions(float* previous_distribution_functions, float* distribution_funtions)
{
  distribution_funtions[0] = previous_distribution_functions[0];
  distribution_funtions[1] = previous_distribution_functions[2];
  distribution_funtions[2] = previous_distribution_functions[1];
  distribution_funtions[3] = previous_distribution_functions[6];
  distribution_funtions[4] = previous_distribution_functions[8];
  distribution_funtions[5] = previous_distribution_functions[7];
  distribution_funtions[6] = previous_distribution_functions[3];
  distribution_funtions[7] = previous_distribution_functions[5];
  distribution_funtions[8] = previous_distribution_functions[4];
  distribution_funtions[9] = previous_distribution_functions[14];
  distribution_funtions[10] = previous_distribution_functions[16];
  distribution_funtions[11] = previous_distribution_functions[15];
  distribution_funtions[12] = previous_distribution_functions[18];
  distribution_funtions[13] = previous_distribution_functions[17];
  distribution_funtions[14] = previous_distribution_functions[9];
  distribution_funtions[15] = previous_distribution_functions[11];
  distribution_funtions[16] = previous_distribution_functions[10];
  distribution_funtions[17] = previous_distribution_functions[13];
  distribution_funtions[18] = previous_distribution_functions[12];
}
// Modify F values from most recent time step (stored in previous_distribution_functions) to impose flow BC. Place result in distribution_functions
__global__ void ImposeBoundaryConditionsGPU(int nLBMmodel, int num_distribution_functions, int lattice_size, int num_full_continuum_mechanics_components, int x_dimension, int y_dimension, float* device_previous_distribution_functions, float* device_distribution_functions, int* geometry_labels, float* boundary_conditions)
{
  // 1. Impose boundary conditions on most recent PDFs stored in F1
  // 2. Transfer PDFs modified by boundary conditions into F0 in preparation for next iteration

  int i = threadIdx.x;  // Lattice i-index
  int j = blockIdx.x;  // Lattice j-index
  int k = blockIdx.y;  // Lattice k-index
  static int DirD3Q19[19][3] = { { 0,0,0 },{ 0,0,1 },{ 0,0,-1 },{ 0,1,0 },{ 0,-1,0 },{ 1,0,0 },{ -1,0,0 },{ 0,1,1 },{ 0,-1,-1 },{ 0,1,-1 },{ 0,-1,1 },{ 1,1,0 },{ -1,-1,0 },{ 1,-1,0 },{ -1,1,0 },{ 1,0,1 },{ -1,0,-1 },{ 1,0,-1 },{ -1,0,1 } };
  long int loc, loc1;
  int j1[3];
  // f0 , f1    old,new f values on a BC node
  // f0n1, f1n1 old,new f values on node interior to BC node along lattice normal n by 1 lattice step
  // f0n2, f1n2 old,new f values on node interior to BC node along lattice normal n by 2 lattice steps
  float previous_distribution_functions[LBM::nPDFmodels];
  float distribution_functions[LBM::nPDFmodels];
  float bcValue[LBM::nContmodels[1]];
  // Offset within GPU memory of this lattice node
  long int initial_location = i + x_dimension * (j + y_dimension * k);
  int intial_j[3] = { i,j,k };
  // This node's boundary condition code
  int bcLabel = geometry_labels[initial_location];
  // Process node PDFs according to boundary condition type:
  // (-1) Inactive node: reinitialze node values in f1 to undo stream & collide
  if (bcLabel == BCinactive)
  {
    for (int n = 0; n < num_distribution_functions; n++)
    {
      distribution_functions[n] = 1e-40;
    }
  }
  // ( 1) Interior node: place values from current time step into f1
  if (bcLabel == BCinterior)
  {
    loc = initial_location;
    for (int n = 0; n < num_distribution_functions; n++)
    {
      distribution_functions[n] = device_distribution_functions[loc];
      loc += lattice_size;
    }
  }
  // (1-64) Wall node: - place current time values current time step into f0;
  //                   - call WallBC to obtain f1 from f0
  if ((mnWALLN <= bcLabel) && (bcLabel <= mxWALLN))
  {
    int n = bcLabel % 100;  // Code for normal orientation, find indices
    for (int d = 0; d < 3; d++)
    {
      j1[d] = intial_j[d] + DirD3Q19[n][2 - d];
    }

    //   Place interior F1 values into f0 (1 lattice step in)
    loc1 = j1[0] + x_dimension * (j1[1] + y_dimension * j1[2]);
    loc = loc1;
    for (int n = 0; n < num_distribution_functions; n++)
    {
      previous_distribution_functions[n] = device_distribution_functions[loc];
      loc += lattice_size;
    }
    SetWallDistributionFunctions(previous_distribution_functions, distribution_functions);
  }
  // (>100) Physics condition (inflow/outflow/interface)
  if (bcLabel >= PHYSICSBC) {
    // Provide data that might be needed by various physical boundary conditions:
    //   1. Load node continuum values into registers
    loc = initial_location;
    for (int n = 0; n < num_full_continuum_mechanics_components; n++)
    {
      bcValue[n] = boundary_conditions[loc];
      loc += lattice_size;
    }
    //   2. Load node PDFs into registers
    loc = initial_location;
    for (int n = 0; n < num_distribution_functions; n++)
    {
      previous_distribution_functions[n] = device_previous_distribution_functions[loc];
      distribution_functions[n] = device_distribution_functions[loc];
      loc += lattice_size;
    }
    //   3. Load PDFs along interior-pointing normal direction into registers
    // Process BC
    SetFlowBCDistributionFunctions(nLBMmodel, previous_distribution_functions, distribution_functions, bcLabel, bcValue);
  }
  // Store BC-modified values from registers f1 into F0 in preparation for next time step
  // Also place values into F1 that will be used as work space in 2-step streaming operations
  loc = initial_location;
  for (int n = 0; n < num_distribution_functions; n++)
  {
    device_distribution_functions[loc] = device_previous_distribution_functions[loc] = distribution_functions[n];
    loc += lattice_size;
  }
}
void ImposeBoundaryConditions(int nLBMmodel, int lattice_size, size_t* lattice_dimensions, float* device_previous_distribution_functions, float* device_distribution_functions, int* device_geometry_labels, float* boundary_conditions)
{
  int num_distribution_functions = LBM::nPDFmodels;
  int num_full_continuum_mechanics_components = LBM::nContmodels[nLBMmodel];
  dim3 dG, dB;
  // Go through entire volume
  // A CUDA block processes a line of data along x-axis, one thread per lattice site
  dB.x = lattice_dimensions[0];
  dB.y = 1;
  dB.z = 1;
  // The CUDA grid corresponds to dimensions along y and z
  dG.x = lattice_dimensions[1];
  dG.y = lattice_dimensions[2];

  ImposeBoundaryConditionsGPU << <dG, dB >> > (nLBMmodel, num_distribution_functions, lattice_size, num_full_continuum_mechanics_components, lattice_dimensions[0], lattice_dimensions[1], device_previous_distribution_functions, device_distribution_functions, device_geometry_labels, boundary_conditions);
}

///////////////
// 2. Stream //
///////////////

// Translate entries 1 index position to right
__global__ void StreamRight(float* distribution_functions, float* Updated_distributed_functions, int x_dimension, int y_dimension)
{
  int i = threadIdx.x;  // Lattice i-index
  int j = blockIdx.x;  // Lattice j-index
  int k = blockIdx.y;  // Lattice k-index
  __shared__ float line[mL];
  long int loc = i + x_dimension * (j + y_dimension * k);
  if (i < x_dimension)
    line[i] = distribution_functions[loc];
  __syncthreads();
  if (i > 0)
    Updated_distributed_functions[loc] = line[i - 1];
  else
    Updated_distributed_functions[loc] = line[0];
}
// Translate entries 1 index position to left
__global__ void StreamLeft(float* distribution_functions, float* Updated_distributed_functions, int x_dimension, int y_dimension)
{
  int i = threadIdx.x;  // Lattice i-index
  int j = blockIdx.x;  // Lattice j-index
  int k = blockIdx.y;  // Lattice k-index
  long int loc;
  __shared__ float line[mL];
  loc = i + x_dimension * (j + y_dimension * k);
  if (i < x_dimension)
    line[i] = distribution_functions[loc];
  __syncthreads();
  if (i < x_dimension - 1)
    Updated_distributed_functions[loc] = line[i + 1];
  else
    Updated_distributed_functions[loc] = line[x_dimension - 1];
}
// Translate entries 1 index position up
__global__ void StreamUp(float* distribution_functions, float* updated_distributed_functions, int x_dimension, int y_dimension, int z_dimensions)
{
  int tix = threadIdx.x;
  int tiy = threadIdx.y;
  int i = threadIdx.x + blockIdx.x * blockDim.x; // Lattice i-index
  int j = threadIdx.y + blockIdx.y * blockDim.y; // Lattice j-index

  long int loc0, loc1, locBC;
  __shared__ float tile[mB][mB + 1];
  for (int k = 0; k < z_dimensions; k++)
  {
    locBC = i + x_dimension * y_dimension * k;
    loc0 = locBC + x_dimension * (j - 1);
    loc1 = loc0 + x_dimension;
    if ((i < x_dimension) & (j < y_dimension)) {
      if (j > 0)
        tile[tix][tiy] = distribution_functions[loc0];
      else
        tile[tix][tiy] = distribution_functions[locBC];
    }
    __syncthreads();
    if ((i < x_dimension) & (j < y_dimension))
      updated_distributed_functions[loc1] = tile[tix][tiy];
  }
}
// Translate entries 1 index position down
__global__ void StreamDown(float* distribution_functions, float* updated_distributed_functions, int x_dimension, int y_dimension, int z_dimension)
{
  int tix = threadIdx.x;
  int tiy = threadIdx.y;
  int i = threadIdx.x + blockIdx.x * blockDim.x; // Lattice i-index
  int j = threadIdx.y + blockIdx.y * blockDim.y; // Lattice j-index
  long int loc0, loc1, locBC;
  __shared__ float tile[mB][mB + 1];
  for (int k = 0; k < z_dimension; k++)
  {
    locBC = i + x_dimension * j + x_dimension * y_dimension * k;
    loc1 = locBC;
    loc0 = loc1 + x_dimension;
    if ((i < x_dimension) & (j < y_dimension)) {
      if (j < y_dimension - 1)
        tile[tix][tiy] = distribution_functions[loc0];
      else
        tile[tix][tiy] = distribution_functions[locBC];
    }
    __syncthreads();
    if ((i < x_dimension) & (j < y_dimension))
      updated_distributed_functions[loc1] = tile[tix][tiy];
  }
}
// Translate entries 1 index position up
__global__ void StreamFront(float* distribution_functions, float* updated_distributed_functions, int x_dimension, int y_dimension, int z_dimension)
{
  int tix = threadIdx.x;
  int tiy = threadIdx.y;
  int i = threadIdx.x + blockIdx.x * blockDim.x; // Lattice i-index
  int k = threadIdx.y + blockIdx.y * blockDim.y; // Lattice k-index
  long int loc0, loc1, locBC;
  __shared__ float tile[mB][mB + 1];
  for (int j = 0; j < y_dimension; j++)
  {
    locBC = i + x_dimension * j;
    loc0 = locBC + x_dimension * y_dimension * (k - 1);
    loc1 = loc0 + x_dimension * y_dimension;
    if ((i < x_dimension) & (k < z_dimension))
    {
      if (k > 0)
        tile[tix][tiy] = distribution_functions[loc0];
      else
        tile[tix][tiy] = distribution_functions[locBC];
    }
    __syncthreads();
    if ((i < x_dimension) & (k < z_dimension))
      updated_distributed_functions[loc1] = tile[tix][tiy];
  }
}
// Translate entries 1 index position down
__global__ void StreamBack(float* distribution_functions, float* updated_distributed_functions, int x_dimension, int y_dimension, int z_dimension)
{
  int tix = threadIdx.x;
  int tiy = threadIdx.y;
  int i = threadIdx.x + blockIdx.x * blockDim.x; // Lattice i-index
  int k = threadIdx.y + blockIdx.y * blockDim.y; // Lattice k-index
  long int locBC, loc0, loc1;
  __shared__ float tile[mB][mB + 1];
  for (int j = 0; j < y_dimension; j++) {
    locBC = i + x_dimension * j + x_dimension * y_dimension * k;
    loc1 = locBC;
    loc0 = loc1 + x_dimension * y_dimension;
    if ((i < x_dimension) & (k < z_dimension)) {
      if (k < z_dimension - 1)
        tile[tix][tiy] = distribution_functions[loc0];
      else
        tile[tix][tiy] = distribution_functions[locBC];
    }
    __syncthreads();
    if ((i < x_dimension) & (k < z_dimension))
      updated_distributed_functions[loc1] = tile[tix][tiy];
  }
}
//Stream step of the LBM algorithm on the GPU - stream in each direction
void Stream_GPU(int direction, size_t* lattice_dimensions, float* previous_distribution_functions, float* distribution_functions)
{

  dim3 dG, dB;
  int nBx, nBy, nBz;
  switch (direction)
  {
  case 0:
  {
    dB.x = lattice_dimensions[0]; dB.y = 1; dB.z = 1;
    dG.x = lattice_dimensions[1]; dG.y = lattice_dimensions[2];
    StreamRight << <dG, dB >> > (previous_distribution_functions, distribution_functions, lattice_dimensions[0], lattice_dimensions[1]); // stream in positive x-direction
    break;
  }
  case 1:
  {
    dB.x = lattice_dimensions[0]; dB.y = 1; dB.z = 1;
    dG.x = lattice_dimensions[1]; dG.y = lattice_dimensions[2];
    StreamLeft << <dG, dB >> > (previous_distribution_functions, distribution_functions, lattice_dimensions[0], lattice_dimensions[1]); // stream in negative x-direction
    break;
  }
  case 2:
  {
    dB.x = mB; dB.y = mB; dB.z = 1;
    nBx = lattice_dimensions[0] / mB; if (lattice_dimensions[0] % mB > 0) nBx++;
    nBy = lattice_dimensions[1] / mB; if (lattice_dimensions[1] % mB > 0) nBy++;
    dG.x = nBx; dG.y = nBy;
    StreamUp << <dG, dB >> > (previous_distribution_functions, distribution_functions, lattice_dimensions[0], lattice_dimensions[1], lattice_dimensions[2]);  // stream in positive y-direction
    break;
  }
  case 3:
  {
    dB.x = mB; dB.y = mB; dB.z = 1;
    nBx = lattice_dimensions[0] / mB; if (lattice_dimensions[0] % mB > 0) nBx++;
    nBy = lattice_dimensions[1] / mB; if (lattice_dimensions[1] % mB > 0) nBy++;
    dG.x = nBx; dG.y = nBy;
    StreamDown << <dG, dB >> > (previous_distribution_functions, distribution_functions, lattice_dimensions[0], lattice_dimensions[1], lattice_dimensions[2]);  // stream in negative y-direction
    break;
  }
  case 4:
  {
    dB.x = mB; dB.y = mB; dB.z = 1;
    nBx = lattice_dimensions[0] / mB; if (lattice_dimensions[0] % mB > 0) nBx++;
    nBz = lattice_dimensions[2] / mB; if (lattice_dimensions[2] % mB > 0) nBz++;
    dG.x = nBx; dG.y = nBz;
    StreamFront << <dG, dB >> > (previous_distribution_functions, distribution_functions, lattice_dimensions[0], lattice_dimensions[1], lattice_dimensions[2]);  // stream in positive z-direction
    break;
  }
  case 5:
  {
    dB.x = mB; dB.y = mB; dB.z = 1;
    nBx = lattice_dimensions[0] / mB; if (lattice_dimensions[0] % mB > 0) nBx++;
    nBz = lattice_dimensions[2] / mB; if (lattice_dimensions[2] % mB > 0) nBz++;
    dG.x = nBx; dG.y = nBz;
    StreamBack << <dG, dB >> > (previous_distribution_functions, distribution_functions, lattice_dimensions[0], lattice_dimensions[1], lattice_dimensions[2]);  // stream in negative z-direction
    break;
  }
  }
}
//Stream step of the LBM algorithm
void LBM_Stream(int nLBMmodel, int num_distribution_functions, size_t* lattice_dimensions, long int* mF, float* previous_distribution_functions, float* device_distribution_functions)
{
  int  mask = 1;
  int nS[LBM::nPDFmodels];// Copy of StreamDirections so we can modify it
  for (int n = 1; n < num_distribution_functions; n++)
    nS[n] = LBM::StreamDirections[n];
  for (int direction = 0; direction < 6; direction++)
  { // Loop over directions
    for (int n = 1; n < num_distribution_functions; n++)
    { // Loop over PDFs
      if (LBM::VelocityDirectionMasks[n] & mask)
      {
        if (nS[n] % 2 == 1)
          Stream_GPU(direction, lattice_dimensions, previous_distribution_functions + mF[n], device_distribution_functions + mF[n]);
        else
          Stream_GPU(direction, lattice_dimensions, device_distribution_functions + mF[n], previous_distribution_functions + mF[n]);
        nS[n]--;
      }
    }
    mask = mask << 1;
  }
}

////////////////
// 4. Collide //
////////////////

__device__ void ComputeEquilibriumContinuumMechanics(float* continuum_mech_values, float* distribution_functions, float* noneq_distribution_functions = nullptr, float viscous_stress_tensor_coeff = 0)
{
  float rho = 0;
//  printf("%e\n",rho);
  for (int n = 0; n < 19; n++)
  {
    rho += distribution_functions[n];  // Pressurecontinuum_mech_values[0]
  }
//  printf("%e\n",rho);
  continuum_mech_values[0] = 0.333333333 * rho;
  continuum_mech_values[1] = (distribution_functions[1] - distribution_functions[2] + distribution_functions[4] - distribution_functions[5] + distribution_functions[7] - distribution_functions[8] + distribution_functions[10] - distribution_functions[11] + distribution_functions[15] - distribution_functions[16]) / rho; // Velocity in x-dir
  continuum_mech_values[2] = (distribution_functions[3] + distribution_functions[4] + distribution_functions[5] - distribution_functions[6] - distribution_functions[7] - distribution_functions[8] + distribution_functions[12] - distribution_functions[13] + distribution_functions[17] - distribution_functions[18]) / rho; // Velocity in y-dir
  continuum_mech_values[3] = (distribution_functions[9] + distribution_functions[10] + distribution_functions[11] + distribution_functions[12] + distribution_functions[13] - distribution_functions[14] - distribution_functions[15] - distribution_functions[16] - distribution_functions[17] - distribution_functions[18]) / rho; // Velocity in z-dir
  if (noneq_distribution_functions != nullptr)
  {
    continuum_mech_values[4] = viscous_stress_tensor_coeff * (noneq_distribution_functions[4] - noneq_distribution_functions[5] - noneq_distribution_functions[7] + noneq_distribution_functions[8]);  //Sigma_xy
    continuum_mech_values[5] = viscous_stress_tensor_coeff * (noneq_distribution_functions[10] - noneq_distribution_functions[11] - noneq_distribution_functions[15] - noneq_distribution_functions[16]);  //Sigma_xz
    continuum_mech_values[6] = viscous_stress_tensor_coeff * (noneq_distribution_functions[12] - noneq_distribution_functions[13] - noneq_distribution_functions[17] + noneq_distribution_functions[18]);  //Sigma_yz
    continuum_mech_values[7] = viscous_stress_tensor_coeff * (noneq_distribution_functions[1] + noneq_distribution_functions[2] + noneq_distribution_functions[4] + noneq_distribution_functions[5] + noneq_distribution_functions[7] + noneq_distribution_functions[8] + noneq_distribution_functions[10] + noneq_distribution_functions[11] + noneq_distribution_functions[15] + noneq_distribution_functions[16]); // Sigma_xx
    continuum_mech_values[8] = viscous_stress_tensor_coeff * (noneq_distribution_functions[3] + noneq_distribution_functions[4] + noneq_distribution_functions[5] + noneq_distribution_functions[6] + noneq_distribution_functions[7] + noneq_distribution_functions[8] + noneq_distribution_functions[12] + noneq_distribution_functions[13] + noneq_distribution_functions[17] + noneq_distribution_functions[18]); // Sigma_yy
    continuum_mech_values[9] = viscous_stress_tensor_coeff * (noneq_distribution_functions[9] + noneq_distribution_functions[10] + noneq_distribution_functions[11] + noneq_distribution_functions[12] + noneq_distribution_functions[13] + noneq_distribution_functions[14] + noneq_distribution_functions[15] + noneq_distribution_functions[16] + noneq_distribution_functions[17] + noneq_distribution_functions[18]); // Sigma_zz
  }
}
__device__ void ComputeTemperature(float* continuum_mech_values, float* distribution_functions)
{
    continuum_mech_values[0] = 0;
    for (int n = 0; n < 19; n++)
    {
        continuum_mech_values[0] += distribution_functions[n];  // Pressurecontinuum_mech_values[0]
    }
    continuum_mech_values[1] = distribution_functions[1] - distribution_functions[2] + distribution_functions[4] - distribution_functions[5] + distribution_functions[7] - distribution_functions[8] + distribution_functions[10] - distribution_functions[11] + distribution_functions[15] - distribution_functions[16]; // Velocity in x-dir
    continuum_mech_values[2] = distribution_functions[3] + distribution_functions[4] + distribution_functions[5] - distribution_functions[6] - distribution_functions[7] - distribution_functions[8] + distribution_functions[12] - distribution_functions[13] + distribution_functions[17] - distribution_functions[18]; // Velocity in y-dir
    continuum_mech_values[3] = distribution_functions[9] + distribution_functions[10] + distribution_functions[11] + distribution_functions[12] + distribution_functions[13] - distribution_functions[14] - distribution_functions[15] - distribution_functions[16] - distribution_functions[17] - distribution_functions[18]; // Velocity in z-dir
}


//Collide step of the LBM algorithm for GPU
__global__ void Collide_GPU(int nLBMmodel, int num_distribution_functions, int lattice_size, float tau, int x_dimension, int y_dimension, float* device_distribution_functions, float* device_noneq_distribution_functions)
{
  int i = threadIdx.x; // Lattice i-index
  int j = blockIdx.x;  // Lattice j-index
  int k = blockIdx.y;  // Lattice k-index
  float distribution_functions[19];
  float noneq_distribution_functions[19];
  float updated_distribution_functions[19];
  float equilibrium_distribution_functions[19];

  long int loc0 = i + x_dimension * (j + y_dimension * k);
  long int loc = loc0;
  for (int n = 0; n < num_distribution_functions; n++)
  {
    distribution_functions[n] = device_distribution_functions[loc];
    loc += lattice_size;
  }
    switch (nLBMmodel)
    {
        case LBM::ModelType::Fluid:
        {
            float continuum_mech_values[4];
            ComputeEquilibriumContinuumMechanics(continuum_mech_values, distribution_functions);
            Device_ComputeDistributionFunctions(continuum_mech_values, equilibrium_distribution_functions);
            break;
        }
        case LBM::ModelType::Preconditioner:
        case LBM::ModelType::Thermal:
        {
            float equilibrium_continuum_mech_values[4];
            ComputePressureFromDistributionFunctions(distribution_functions, equilibrium_continuum_mech_values);
            ComputeAdvectionDiffusionDistributionFunctions(equilibrium_continuum_mech_values, equilibrium_distribution_functions);
            break;
        }
    }

  for (int n = 0; n < num_distribution_functions; n++) {
      updated_distribution_functions[n] = distribution_functions[n] + (equilibrium_distribution_functions[n] - distribution_functions[n]) / tau;
      noneq_distribution_functions[n] = updated_distribution_functions[n] - equilibrium_distribution_functions[n];
  }
  loc = loc0;
  for (int n = 0; n < num_distribution_functions; n++)
  {
    device_distribution_functions[loc] = updated_distribution_functions[n];
    device_noneq_distribution_functions[loc] = noneq_distribution_functions[n];
    loc += lattice_size;
  }
}
// Collide step of LBM algorithm
void LBM_Collide(int nLBMmodel, int lattice_size, float tau, size_t* lattice_dimensions, float* distribution_functions, float* noneq_distribution_functions)
{
  dim3 dG, dB;
  dB.x = lattice_dimensions[0]; dB.y = 1; dB.z = 1;
  dG.x = lattice_dimensions[1]; dG.y = lattice_dimensions[2];
  int num_distribution_functions = LBM::nPDFmodels;
  Collide_GPU << <dG, dB >> > (nLBMmodel, num_distribution_functions, lattice_size, tau, lattice_dimensions[0], lattice_dimensions[1], distribution_functions, noneq_distribution_functions);
}

////////////////////////////
// 5. Contimuum Mechanics //
////////////////////////////

//Continuum Mechanics Calculations
__global__ void ContinuumMechanics_GPU(int num_distribution_functions, int num_full_continuum_mechanics_components, int lattice_size, int x_direction, int y_direction, float* device_distribution_functions, float* device_continuum_mech_values, float* device_noneq_distribution_functions = nullptr, float viscous_stress_tensor_coeff = 0)
{
  int i = threadIdx.x;  // Lattice i-index
  int j = blockIdx.x;  // Lattice j-index
  int k = blockIdx.y;  // Lattice k-index
  float distribution_functions[19];
  float noneq_distribution_functions[19];
  float continuum_mech_values[10];

  long int loc0 = i + x_direction * (j + y_direction * k);
  long int loc = loc0;
  for (int n = 0; n < num_distribution_functions; n++)
  {
    distribution_functions[n] = device_distribution_functions[loc];
    if (device_noneq_distribution_functions != nullptr)
      noneq_distribution_functions[n] = device_noneq_distribution_functions[loc];
    loc += lattice_size;
  }
  if (device_noneq_distribution_functions == nullptr)
    ComputeEquilibriumContinuumMechanics(continuum_mech_values, distribution_functions);
  else
    ComputeEquilibriumContinuumMechanics(continuum_mech_values, distribution_functions, noneq_distribution_functions, viscous_stress_tensor_coeff);
  loc = loc0;
  for (int m = 0; m < num_full_continuum_mechanics_components; m++)
  {
    device_continuum_mech_values[loc] = continuum_mech_values[m];
    loc += lattice_size;
  }
}

__global__ void ComputeTemperature_GPU(int num_distribution_functions, int num_full_continuum_mechanics_components, int lattice_size, int x_direction, int y_direction, float* device_distribution_functions, float* device_continuum_mech_values)
{
    int i = threadIdx.x;  // Lattice i-index
    int j = blockIdx.x;  // Lattice j-index
    int k = blockIdx.y;  // Lattice k-index
    float distribution_functions[19];
    float continuum_mech_values[10];

    long int loc0 = i + x_direction * (j + y_direction * k);
    long int loc = loc0;
    for (int n = 0; n < num_distribution_functions; n++)
    {
        distribution_functions[n] = device_distribution_functions[loc];
        loc += lattice_size;
    }
        ComputeTemperature(continuum_mech_values, distribution_functions);

    loc = loc0;
    for (int m = 0; m < num_full_continuum_mechanics_components; m++)
    {
        device_continuum_mech_values[loc] = continuum_mech_values[m];
        loc += lattice_size;
    }
}
//Compute Continuum Mechanics
void LBM_ContinuumMechanics(int nLBMmodel, int lattice_size, size_t* lattice_dimensions, float* device_distribution_functions, float* device_continuum_mech_values, int compute_stress, float* device_noneq_distribution_functions = nullptr, float viscous_stress_tensor_coeff = 0)
{
  dim3 dG, dB;
  dB.x = lattice_dimensions[0]; dB.y = 1; dB.z = 1;
  dG.x = lattice_dimensions[1]; dG.y = lattice_dimensions[2];
  int num_distribution_functions = LBM::nPDFmodels;
  int num_full_continuum_mechanics_components = LBM::nContmodels[nLBMmodel];
  switch (nLBMmodel)
  {
  case LBM::ModelType::Fluid:
  {
      if(compute_stress > 0)
      {
      ContinuumMechanics_GPU << <dG, dB >> > (num_distribution_functions, num_full_continuum_mechanics_components, lattice_size, lattice_dimensions[0], lattice_dimensions[1], device_distribution_functions, device_continuum_mech_values, device_noneq_distribution_functions, viscous_stress_tensor_coeff);
      }
      else
      {
      ContinuumMechanics_GPU << <dG, dB >> > (num_distribution_functions, num_full_continuum_mechanics_components, lattice_size, lattice_dimensions[0], lattice_dimensions[1], device_distribution_functions, device_continuum_mech_values);
      }
    break;
  }
  case LBM::ModelType::Preconditioner:
  case LBM::ModelType::Thermal:
  {
    ComputeTemperature_GPU << <dG, dB >> > (num_distribution_functions, num_full_continuum_mechanics_components, lattice_size, lattice_dimensions[0], lattice_dimensions[1], device_distribution_functions, device_continuum_mech_values);
    break;
  }
  }
}

//////////////////////////////////////
// Put all the above steps together //
//////////////////////////////////////

extern "C" int microLBM(LBM::Run& run)
{
  static long int offF[LBM::nPDFmodels];
  static float* device_previous_distribution_functions;
  static float* device_distribution_functions;
  static float* device_noneq_distribution_functions;
  static float *device_continuum_mech_values, *device_imposed_bc;
  static int *device_geometry_labels;

  //----------------------------------------------------------------------------       //
  /* Load run parameters */
  float tmp_n, tmp_d, error;
  float tau;
  float viscous_stress_tensor_coeff =  -1 + 0.5f / run.in.relaxation_time;
  int lattice_size = run.in.dimensions[0]* run.in.dimensions[1]* run.in.dimensions[2];
  bool compute_stress = false;  // flag for computing the stress tensor components. true: compute, false: do not compute
  //-----------------------------------------------------------------------------------//
  /* Allocate device memory */
  long long nBytesGeom = lattice_size * sizeof(int);
  long long nBytesPDF = lattice_size * LBM::nPDFmodels * sizeof(float);
  long long nBytesCont = lattice_size * LBM::nContmodels[run.in.model_type] * sizeof(float);
  long long nBytesQBC = LBM::nContmodels[run.in.model_type] * run.out.nBound * sizeof(float);
  long long nBytesFinact = LBM::nPDFmodels * sizeof(float);
  // Total space
  long long num_total_bytes = nBytesGeom + 3 * nBytesPDF + nBytesCont + nBytesQBC + nBytesFinact;
  long long total_memory_requirement_MB = num_total_bytes / 1024 / 1024;

  if (total_memory_requirement_MB > DeviceMB)
  {
    run.GetLogger()->Info("Not enough memory on GPU device");
    run.GetLogger()->Info("The lattice dimensions = " + run.GetLogger()->to_string(run.in.dimensions));
    run.GetLogger()->Info("nLatticeSites : " + std::to_string(lattice_size));
    run.GetLogger()->Info("nBytesGeom : " + std::to_string(nBytesGeom / (1024 * 1024)));
    run.GetLogger()->Info("nBytesPDFs : " + std::to_string(2 * nBytesPDF / (1024 * 1024)));
    run.GetLogger()->Info("nBytesCont : " + std::to_string(nBytesCont / (1024 * 1024)));
    return 1;
  }
  // Allocate space on GPU device memory:
  cudaMalloc((void**)&device_geometry_labels, nBytesGeom);                       // geometry flags
  cudaMalloc((void**)&device_previous_distribution_functions, nBytesPDF);        // old PDFs
  cudaMalloc((void**)&device_distribution_functions, nBytesPDF);                 // new PDFs
  cudaMalloc((void**)&device_noneq_distribution_functions, nBytesPDF);           // non-equilibrium PDFs
  cudaMalloc((void**)&device_continuum_mech_values, nBytesCont);                 // continuum variables from LBM PDFs
  cudaMalloc((void**)&device_imposed_bc, nBytesQBC);                             // imposed boundary conditions on continuum variables
  
  // Offsets to each PDF
  for (int i = 0; i < LBM::nPDFmodels; i++)
  {
    offF[i] = i * lattice_size;
  }

  float* continuum_mech_values;
  std::vector<float>* timing_ms;
  std::vector<float>* iteration_errors;
  switch (run.in.model_type)
  {
  case LBM::ModelType::Preconditioner:
    continuum_mech_values = run.out.pressure_flows_and_stress;
    iteration_errors = &run.out.pre_conditioning_iteration_error;
    timing_ms = &run.out.pre_conditioning_iteration_timing_ms;
    tau = run.in.relaxation_time;
    break;
  case LBM::ModelType::Fluid:
    continuum_mech_values = run.out.pressure_flows_and_stress;
    iteration_errors = &run.out.fluid_iteration_error;
    timing_ms = &run.out.fluid_iteration_timing_ms;
    tau = run.in.relaxation_time;
    break;
  case LBM::ModelType::Thermal:
    continuum_mech_values = run.out.temperature;
    iteration_errors = &run.out.temperature_iteration_error;
    timing_ms = &run.out.temperature_iteration_timing_ms;
    tau = run.in.relaxation_time_temperature;
    break;
  }

  // Load geometry data to device memory
  cudaMemcpy(device_geometry_labels, run.out.labels, nBytesGeom, cudaMemcpyHostToDevice);              // geometry flags
  cudaMemcpy(device_imposed_bc, run.out.boundary_conditions, nBytesQBC, cudaMemcpyHostToDevice);       // Load boundary conditions to device memory
  cudaMemcpy(device_continuum_mech_values, continuum_mech_values, nBytesCont, cudaMemcpyHostToDevice); // Load continuum field values to device memory (always needed in order to transmit boundary conditions)
  
  if (run.in.initialize_distribution_functions_from_cont_mech)
  {
    // Set initial PDFs from continuum field variables pressure_flow_stress (on Host), device_pressure_flow_stress (on GPU) - need to set up a previous and current distribution function set
    SetupDistributionFunctions(run.in.model_type, run.in.dimensions, lattice_size, device_continuum_mech_values, device_distribution_functions);
    SetupDistributionFunctions(run.in.model_type, run.in.dimensions, lattice_size, device_continuum_mech_values, device_previous_distribution_functions);
  }
  else
  {
    // Transfer PDF values from Host to Device memory (continue from previous computation state)
    cudaMemcpy(device_distribution_functions, run.out.fPDF, nBytesPDF, cudaMemcpyHostToDevice);
    cudaMemcpy(device_previous_distribution_functions, run.out.fPDF, nBytesPDF, cudaMemcpyHostToDevice);
  }

  cudaMemcpy(device_noneq_distribution_functions, run.out.fnPDF, nBytesPDF, cudaMemcpyHostToDevice);

  TimingProfile profiler;
  profiler.Start("ALL");
  //============================ Main iterative loop ======================================
  for(size_t itr = 0; itr < run.cfg.max_iterations; itr++)
  {
    profiler.Start("ITR");
    //    1.  Impose boundary conditions
    ImposeBoundaryConditions(run.in.model_type, lattice_size, run.in.dimensions, device_previous_distribution_functions, device_distribution_functions, device_geometry_labels, device_continuum_mech_values); // BC(F1)    -> F0,F1
    //    2.  Stream step of the LBM algorithm
    LBM_Stream(run.in.model_type, LBM::nPDFmodels, run.in.dimensions, offF, device_previous_distribution_functions, device_distribution_functions); // stream(F0)-> F1
    //    3.  Reimpose Boundary Conditions
    ImposeBoundaryConditions(run.in.model_type, lattice_size, run.in.dimensions, device_previous_distribution_functions, device_distribution_functions, device_geometry_labels, device_continuum_mech_values); // BC(F1)    -> F0,F1
    //    4.  Collide step of the LBM algorithm
    LBM_Collide(run.in.model_type, lattice_size, tau, run.in.dimensions, device_distribution_functions, device_noneq_distribution_functions);         // relax(F1) -> F1
    //    5.  Continuum Mechanics Calculation
    LBM_ContinuumMechanics(run.in.model_type, lattice_size, run.in.dimensions, device_distribution_functions, device_continuum_mech_values, compute_stress, device_noneq_distribution_functions, viscous_stress_tensor_coeff);
    cudaMemcpy(continuum_mech_values, device_continuum_mech_values, nBytesCont, cudaMemcpyDeviceToHost);   // Copy continuum field values

    tmp_n = 0.f;
    tmp_d = 0.f;

    if(run.in.model_type == LBM::ModelType::Thermal)
    {
      for (int i = 0; i < lattice_size; i++)
      {
          tmp_n += std::pow((run.out.previous_results[i] - continuum_mech_values[i]), 2);
          tmp_d += std::pow(continuum_mech_values[i], 2);
          run.out.previous_results[i] = continuum_mech_values[i];
      }
    }
    else
    {
      for (int j = 0; j < lattice_size * 3; j++)
      {
          tmp_n += std::pow((run.out.previous_results[j] - continuum_mech_values[lattice_size + j]), 2);
          tmp_d += std::pow(continuum_mech_values[lattice_size + j],2);
          run.out.previous_results[j] = continuum_mech_values[lattice_size + j];
      }
    }
    error = std::pow(tmp_n / tmp_d, 0.5f);
    iteration_errors->push_back(error);
    timing_ms->push_back((float)profiler.GetElapsedTime_s("ITR") * 1000);

    if (itr == run.cfg.print_iteration)
    {
      run.cfg.print_iteration += run.cfg.print_stride;
      run.GetLogger()->Info("  Convergence error " + std::to_string(error) + " at iteration " + std::to_string(itr));
    }

    if (error < run.cfg.tolerance)
    {
      switch (run.in.model_type)
      {
      case LBM::ModelType::Preconditioner:
        run.out.num_preconditioner_iterations = itr;
        break;
      case LBM::ModelType::Fluid:
        run.out.num_fluid_iterations = itr;
        break;
      case LBM::ModelType::Thermal:
        run.out.num_thermal_iterations = itr;
        break;
      }
      run.GetLogger()->Info("Solution has converged at iteration "+std::to_string(itr)+" with error "+std::to_string(error));
      break;
    }
    else if (itr >= run.cfg.max_iterations - 1 && error > run.cfg.tolerance)
    {
      switch (run.in.model_type)
      {
      case LBM::ModelType::Preconditioner:
        run.out.num_preconditioner_iterations = run.cfg.max_iterations;
        break;
      case LBM::ModelType::Fluid:
        run.out.num_fluid_iterations = run.cfg.max_iterations;
        break;
      case LBM::ModelType::Thermal:
        run.out.num_thermal_iterations = run.cfg.max_iterations;
        break;
      }
      run.GetLogger()->Info("Solution has reached maximum of " + std::to_string(itr+1) + " iterations");
    }
    error = 0.f;
  }
  //======================================================================================
  int tsec = (int)profiler.GetElapsedTime_s("ALL");
  int mins = tsec/60;
  int secs = tsec%60;
  run.GetLogger()->Info("Solution took "+std::to_string(mins)+":"+std::to_string(secs)+" (min:s) time.");

  //  6.  Compute the viscous stress tensor from the latest run only
  compute_stress = true;
  LBM_ContinuumMechanics(run.in.model_type, lattice_size, run.in.dimensions, device_distribution_functions, device_continuum_mech_values, compute_stress, device_noneq_distribution_functions, viscous_stress_tensor_coeff);
  //  7.  Copy Values
  cudaMemcpy(continuum_mech_values, device_continuum_mech_values, nBytesCont, cudaMemcpyDeviceToHost);   // Copy continuum field values
  if (run.in.return_distribution_functions)
  {
    cudaMemcpy(run.out.fPDF, device_distribution_functions, nBytesPDF, cudaMemcpyDeviceToHost); // Copy PDF values
    cudaMemcpy(run.out.fnPDF, device_noneq_distribution_functions, nBytesPDF, cudaMemcpyDeviceToHost); // Copy non-equilibrium PDF values
  }
  //free memory
  cudaFree(device_geometry_labels);
  cudaFree(device_previous_distribution_functions);
  cudaFree(device_distribution_functions);
  cudaFree(device_noneq_distribution_functions);
  cudaFree(device_continuum_mech_values);
  cudaFree(device_imposed_bc);

  return 0;
}
