/*
color lattice boltzmann model
 */
#include "models/FreeLeeModel.h"
#include "analysis/distance.h"
#include "analysis/morphology.h"
#include "common/Communication.h"
#include "common/ReadMicroCT.h"
#include <stdlib.h>
#include <time.h>

ScaLBL_FreeLeeModel::ScaLBL_FreeLeeModel(int RANK, int NP,
                                         const Utilities::MPI &COMM)
    : rank(RANK), nprocs(NP), Restart(0), timestep(0), timestepMax(2),
      tauA(1.0), tauB(1.0), tauM(1.0), rhoA(1.0), rhoB(1.0), W(5.0),
      gamma(0.001), kappa(0.0075), beta(0.0024), Fx(0), Fy(0), Fz(0), flux(0),
      din(0), dout(0), inletA(0), inletB(0), outletA(0), outletB(0), tau(1.0),
      rho0(1.0), Nx(0), Ny(0), Nz(0), N(0), Np(0), nprocx(0), nprocy(0),
      nprocz(0), BoundaryCondition(0), Lx(0), Ly(0), Lz(0), comm(COMM) {}
ScaLBL_FreeLeeModel::~ScaLBL_FreeLeeModel() {}

void ScaLBL_FreeLeeModel::getPhase(DoubleArray &PhaseValues) {

    DoubleArray PhaseWideHalo(Nxh, Nyh, Nzh);
    ScaLBL_CopyToHost(PhaseWideHalo.data(), Phi, sizeof(double) * Nh);

    // use halo width = 1 for analysis data
    for (int k = 1; k < Nzh - 1; k++) {
        for (int j = 1; j < Nyh - 1; j++) {
            for (int i = 1; i < Nxh - 1; i++) {
                PhaseValues(i - 1, j - 1, k - 1) = PhaseWideHalo(i, j, k);
            }
        }
    }
}

void ScaLBL_FreeLeeModel::getPotential(DoubleArray &PressureValues,
                                       DoubleArray &MuValues) {

    ScaLBL_Comm->RegularLayout(Map, Pressure, PressureValues);
    ScaLBL_Comm->Barrier();
    comm.barrier();

    ScaLBL_Comm->RegularLayout(Map, mu_phi, MuValues);
    ScaLBL_Comm->Barrier();
    comm.barrier();
}

void ScaLBL_FreeLeeModel::getVelocity(DoubleArray &Vel_x, DoubleArray &Vel_y,
                                      DoubleArray &Vel_z) {

    ScaLBL_Comm->RegularLayout(Map, &Velocity[0], Vel_x);
    ScaLBL_Comm->Barrier();
    comm.barrier();

    ScaLBL_Comm->RegularLayout(Map, &Velocity[Np], Vel_y);
    ScaLBL_Comm->Barrier();
    comm.barrier();

    ScaLBL_Comm->RegularLayout(Map, &Velocity[2 * Np], Vel_z);
    ScaLBL_Comm->Barrier();
    comm.barrier();
}

void ScaLBL_FreeLeeModel::getData_RegularLayout(const double *data,
                                                DoubleArray &regdata) {
    // Gets data (in optimized layout) from the HOST and stores in regular layout
    // Primarly for debugging
    int i, j, k, idx;

    // initialize the array
    regdata.fill(0.f);

    double value;
    for (k = 0; k < Nz; k++) {
        for (j = 0; j < Ny; j++) {
            for (i = 0; i < Nx; i++) {
                idx = Map(i, j, k);
                if (!(idx < 0)) {
                    value = data[idx];
                    regdata(i, j, k) = value;
                }
            }
        }
    }
}

void ScaLBL_FreeLeeModel::ReadParams(string filename) {
    // read the input database
    db = std::make_shared<Database>(filename);
    domain_db = db->getDatabase("Domain");
    freelee_db = db->getDatabase("FreeLee");
    analysis_db = db->getDatabase("Analysis");
    vis_db = db->getDatabase("Visualization");

    // set defaults
    timestepMax = 100000;
    tauA = tauB = 1.0;
    tauM = 1.0; //relaxation time for phase field
    rhoA = rhoB = 1.0;
    tau = 1.0;  //only for single-fluid Lee model
    rho0 = 1.0; //only for single-fluid Lee model
    Fx = Fy = Fz = 0.0;
    gamma = 1e-3; //surface tension
    W = 5.0;      //interfacial thickness
    //beta = 12.0*gamma/W;
    //kappa = 3.0*gamma*W/2.0;//beta and kappa are related to surface tension \gamma
    beta = 0.75 * gamma / W;
    kappa = 0.375 * gamma *
            W; //beta and kappa are related to surface tension \gamma
    Restart = false;
    din = dout = 1.0;
    flux = 0.0;

    // Color Model parameters
    if (freelee_db->keyExists("timestepMax")) {
        timestepMax = freelee_db->getScalar<int>("timestepMax");
    }
    if (freelee_db->keyExists("tau")) { //only for single-fluid Lee model
        tau = freelee_db->getScalar<double>("tau");
    }
    if (freelee_db->keyExists("tauA")) {
        tauA = freelee_db->getScalar<double>("tauA");
    }
    if (freelee_db->keyExists("tauB")) {
        tauB = freelee_db->getScalar<double>("tauB");
    }
    if (freelee_db->keyExists("tauM")) {
        tauM = freelee_db->getScalar<double>("tauM");
    }
    if (freelee_db->keyExists("rho0")) {
        rho0 = freelee_db->getScalar<double>("rho0");
    }
    if (freelee_db->keyExists("rhoA")) {
        rhoA = freelee_db->getScalar<double>("rhoA");
    }
    if (freelee_db->keyExists("rhoB")) {
        rhoB = freelee_db->getScalar<double>("rhoB");
    }
    if (freelee_db->keyExists("F")) {
        Fx = freelee_db->getVector<double>("F")[0];
        Fy = freelee_db->getVector<double>("F")[1];
        Fz = freelee_db->getVector<double>("F")[2];
    }
    if (freelee_db->keyExists("gamma")) {
        gamma = freelee_db->getScalar<double>("gamma");
    }
    if (freelee_db->keyExists("W")) {
        W = freelee_db->getScalar<double>("W");
    }
    if (freelee_db->keyExists("Restart")) {
        Restart = freelee_db->getScalar<bool>("Restart");
    }
    if (freelee_db->keyExists("din")) {
        din = freelee_db->getScalar<double>("din");
    }
    if (freelee_db->keyExists("dout")) {
        dout = freelee_db->getScalar<double>("dout");
    }
    if (freelee_db->keyExists("flux")) {
        flux = freelee_db->getScalar<double>("flux");
    }
    inletA = 1.f;
    inletB = 0.f;
    outletA = 0.f;
    outletB = 1.f;
    //update secondary parameters
    //beta = 12.0*gamma/W;
    //kappa = 3.0*gamma*W/2.0;//beta and kappa are related to surface tension \gamma
    beta = 0.75 * gamma / W;
    kappa = 0.375 * gamma *
            W; //beta and kappa are related to surface tension \gamma
    //if (BoundaryCondition==4) flux *= rhoA; // mass flux must adjust for density (see formulation for details)

    BoundaryCondition = 0;
    if (domain_db->keyExists("BC")) {
        BoundaryCondition = domain_db->getScalar<int>("BC");
    }
}

void ScaLBL_FreeLeeModel::SetDomain() {
    Dm = std::shared_ptr<Domain>(
        new Domain(domain_db, comm)); // full domain for analysis
    Mask = std::shared_ptr<Domain>(
        new Domain(domain_db, comm)); // mask domain removes immobile phases
    // domain parameters
    Nx = Dm->Nx;
    Ny = Dm->Ny;
    Nz = Dm->Nz;
    Lx = Dm->Lx;
    Ly = Dm->Ly;
    Lz = Dm->Lz;
    N = Nx * Ny * Nz;
    Nxh = Nx + 2;
    Nyh = Ny + 2;
    Nzh = Nz + 2;
    Nh = Nxh * Nyh * Nzh;
    id = new signed char[N];
    for (int i = 0; i < Nx * Ny * Nz; i++)
        Dm->id[i] = 1; // initialize this way

    comm.barrier();
    Dm->CommInit();
    comm.barrier();
    // Read domain parameters
    rank = Dm->rank();
    nprocx = Dm->nprocx();
    nprocy = Dm->nprocy();
    nprocz = Dm->nprocz();
}

void ScaLBL_FreeLeeModel::ReadInput() {

    sprintf(LocalRankString, "%05d", rank);
    sprintf(LocalRankFilename, "%s%s", "ID.", LocalRankString);
    sprintf(LocalRestartFile, "%s%s", "Restart.", LocalRankString);

    if (freelee_db->keyExists("image_sequence")) {
        auto ImageList = freelee_db->getVector<std::string>("image_sequence");
        int IMAGE_INDEX = freelee_db->getWithDefault<int>("image_index", 0);
        std::string first_image = ImageList[IMAGE_INDEX];
        Mask->Decomp(first_image);
        IMAGE_INDEX++;
    } else if (domain_db->keyExists("GridFile")) {
        // Read the local domain data
        auto input_id = readMicroCT(*domain_db, MPI_COMM_WORLD);
        // Fill the halo (assuming GCW of 1)
        array<int, 3> size0 = {(int)input_id.size(0), (int)input_id.size(1),
                               (int)input_id.size(2)};
        ArraySize size1 = {(size_t)Mask->Nx, (size_t)Mask->Ny,
                           (size_t)Mask->Nz};
        ASSERT((int)size1[0] == size0[0] + 2 && (int)size1[1] == size0[1] + 2 &&
               (int)size1[2] == size0[2] + 2);
        fillHalo<signed char> fill(MPI_COMM_WORLD, Mask->rank_info, size0,
                                   {1, 1, 1}, 0, 1);
        Array<signed char> id_view;
        id_view.viewRaw(size1, Mask->id.data());
        fill.copy(input_id, id_view);
        fill.fill(id_view);
    } else if (domain_db->keyExists("Filename")) {
        auto Filename = domain_db->getScalar<std::string>("Filename");
        Mask->Decomp(Filename);
    } else {
        Mask->ReadIDs();
    }
    for (int i = 0; i < Nx * Ny * Nz; i++)
        id[i] = Mask->id[i]; // save what was read

    // Generate the signed distance map
    // Initialize the domain and communication
    Array<char> id_solid(Nx, Ny, Nz);
    // Solve for the position of the solid phase
    for (int k = 0; k < Nz; k++) {
        for (int j = 0; j < Ny; j++) {
            for (int i = 0; i < Nx; i++) {
                int n = k * Nx * Ny + j * Nx + i;
                // Initialize the solid phase
                signed char label = Mask->id[n];
                if (label > 0)
                    id_solid(i, j, k) = 1;
                else
                    id_solid(i, j, k) = 0;
            }
        }
    }
    SignDist.resize(Nx, Ny, Nz);
    // Initialize the signed distance function
    for (int k = 0; k < Nz; k++) {
        for (int j = 0; j < Ny; j++) {
            for (int i = 0; i < Nx; i++) {
                // Initialize distance to +/- 1
                SignDist(i, j, k) = 2.0 * double(id_solid(i, j, k)) - 1.0;
            }
        }
    }
    if (rank == 0)
        printf("Initialized solid phase -- Converting to Signed Distance "
               "function \n");
    CalcDist(SignDist, id_solid, *Mask);

    if (rank == 0)
        cout << "Domain set." << endl;
}

void ScaLBL_FreeLeeModel::Create_TwoFluid() {
    /*
	 *  This function creates the variables needed to run two-fluid Lee model
	 */
    //.........................................................
    // Initialize communication structures in averaging domain
    for (int i = 0; i < Nx * Ny * Nz; i++)
        Dm->id[i] = Mask->id[i];
    Mask->CommInit();
    Np = Mask->PoreCount();
    //...........................................................................
    if (rank == 0)
        printf("Create ScaLBL_Communicator \n");
    // Create a communicator for the device (will use optimized layout)
    // ScaLBL_Communicator ScaLBL_Comm(Mask); // original
    ScaLBL_Comm =
        std::shared_ptr<ScaLBL_Communicator>(new ScaLBL_Communicator(Mask));
    //ScaLBL_Comm_Regular  = std::shared_ptr<ScaLBL_Communicator>(new ScaLBL_Communicator(Mask));
    ScaLBL_Comm_WideHalo = std::shared_ptr<ScaLBLWideHalo_Communicator>(
        new ScaLBLWideHalo_Communicator(Mask, 2));

    // create the layout for the LBM
    int Npad = (Np / 16 + 2) * 16;
    if (rank == 0)
        printf("Set up memory efficient layout, %i | %i | %i \n", Np, Npad, N);
    Map.resize(Nx, Ny, Nz);
    Map.fill(-2);
    auto neighborList = new int[18 * Npad];
    Np = ScaLBL_Comm->MemoryOptimizedLayoutAA(Map, neighborList,
                                              Mask->id.data(), Np, 2);
    comm.barrier();

    //...........................................................................
    //                MAIN  VARIABLES ALLOCATED HERE
    //...........................................................................
    // LBM variables
    if (rank == 0)
        printf("Allocating distributions \n");
    //......................device distributions.................................
    dist_mem_size = Np * sizeof(double);
    neighborSize = 18 * (Np * sizeof(int));
    //...........................................................................
    ScaLBL_AllocateDeviceMemory((void **)&NeighborList, neighborSize);
    ScaLBL_AllocateDeviceMemory((void **)&dvcMap, sizeof(int) * Np);
    ScaLBL_AllocateDeviceMemory((void **)&gqbar, 19 * dist_mem_size);
    ScaLBL_AllocateDeviceMemory((void **)&hq, 7 * dist_mem_size);
    ScaLBL_AllocateDeviceMemory((void **)&mu_phi, dist_mem_size);
    ScaLBL_AllocateDeviceMemory((void **)&Den, dist_mem_size);
    ScaLBL_AllocateDeviceMemory((void **)&Phi, sizeof(double) * Nh);
    ScaLBL_AllocateDeviceMemory((void **)&Pressure, sizeof(double) * Np);
    ScaLBL_AllocateDeviceMemory((void **)&Velocity, 3 * sizeof(double) * Np);
    ScaLBL_AllocateDeviceMemory((void **)&ColorGrad, 3 * sizeof(double) * Np);
    //...........................................................................
    // Update GPU data structures
    if (rank == 0)
        printf("Setting up device map and neighbor list \n");
    fflush(stdout);
    int *TmpMap;
    TmpMap = new int[Np];
    for (int k = 1; k < Nz - 1; k++) {
        for (int j = 1; j < Ny - 1; j++) {
            for (int i = 1; i < Nx - 1; i++) {
                int idx = Map(i, j, k);
                if (!(idx < 0))
                    TmpMap[idx] = ScaLBL_Comm_WideHalo->Map(i, j, k);
            }
        }
    }
    // check that TmpMap is valid
    for (int idx = 0; idx < ScaLBL_Comm->LastExterior(); idx++) {
        auto n = TmpMap[idx];
        if (n > Nxh * Nyh * Nzh) {
            printf("Bad value! idx=%i \n", n);
            TmpMap[idx] = Nxh * Nyh * Nzh - 1;
        }
    }
    for (int idx = ScaLBL_Comm->FirstInterior();
         idx < ScaLBL_Comm->LastInterior(); idx++) {
        auto n = TmpMap[idx];
        if (n > Nxh * Nyh * Nzh) {
            printf("Bad value! idx=%i \n", n);
            TmpMap[idx] = Nxh * Nyh * Nzh - 1;
        }
    }
    // copy the device map
    ScaLBL_CopyToDevice(dvcMap, TmpMap, sizeof(int) * Np);
    // copy the neighbor list
    ScaLBL_CopyToDevice(NeighborList, neighborList, neighborSize);
    comm.barrier();
    delete[] TmpMap;
    delete[] neighborList;
}

void ScaLBL_FreeLeeModel::Create_SingleFluid() {
    /*
	 *  This function creates the variables needed to run single-fluid Lee model
	 */
    //.........................................................
    // Initialize communication structures in averaging domain
    for (int i = 0; i < Nx * Ny * Nz; i++)
        Dm->id[i] = Mask->id[i];
    Mask->CommInit();
    Np = Mask->PoreCount();
    //...........................................................................
    if (rank == 0)
        printf("Create ScaLBL_Communicator \n");
    // Create a communicator for the device (will use optimized layout)
    // ScaLBL_Communicator ScaLBL_Comm(Mask); // original
    ScaLBL_Comm =
        std::shared_ptr<ScaLBL_Communicator>(new ScaLBL_Communicator(Mask));

    // create the layout for the LBM
    int Npad = (Np / 16 + 2) * 16;
    if (rank == 0)
        printf("Set up memory efficient layout, %i | %i | %i \n", Np, Npad, N);
    Map.resize(Nx, Ny, Nz);
    Map.fill(-2);
    auto neighborList = new int[18 * Npad];
    Np = ScaLBL_Comm->MemoryOptimizedLayoutAA(Map, neighborList,
                                              Mask->id.data(), Np, 1);
    comm.barrier();

    //...........................................................................
    //                MAIN  VARIABLES ALLOCATED HERE
    //...........................................................................
    // LBM variables
    if (rank == 0)
        printf("Allocating distributions \n");
    //......................device distributions.................................
    dist_mem_size = Np * sizeof(double);
    neighborSize = 18 * (Np * sizeof(int));
    //...........................................................................
    ScaLBL_AllocateDeviceMemory((void **)&NeighborList, neighborSize);
    ScaLBL_AllocateDeviceMemory((void **)&gqbar, 19 * dist_mem_size);
    ScaLBL_AllocateDeviceMemory((void **)&Pressure, sizeof(double) * Np);
    ScaLBL_AllocateDeviceMemory((void **)&Velocity, 3 * sizeof(double) * Np);
    //...........................................................................
    // Update GPU data structures
    if (rank == 0)
        printf("Setting up device map and neighbor list \n");
    // copy the neighbor list
    ScaLBL_CopyToDevice(NeighborList, neighborList, neighborSize);
    comm.barrier();
    delete[] neighborList;
}

void ScaLBL_FreeLeeModel::AssignComponentLabels_ChemPotential_ColorGrad() {
    double *phase;
    phase = new double[Nh];

    size_t NLABELS = 0;
    signed char VALUE = 0;
    double AFFINITY = 0.f;

    auto LabelList = freelee_db->getVector<int>("ComponentLabels");
    auto AffinityList = freelee_db->getVector<double>("ComponentAffinity");

    NLABELS = LabelList.size();
    if (NLABELS != AffinityList.size()) {
        ERROR("Error: ComponentLabels and ComponentAffinity must be the same "
              "length! \n");
    }

    double *label_count;
    double *label_count_global;
    label_count = new double[NLABELS];
    label_count_global = new double[NLABELS];

    // Assign the labels
    for (size_t idx = 0; idx < NLABELS; idx++)
        label_count[idx] = 0;
    for (int k = 0; k < Nzh; k++) {
        for (int j = 0; j < Nyh; j++) {
            for (int i = 0; i < Nxh; i++) {

                //idx for double-halo array 'phase'
                int nh = k * Nxh * Nyh + j * Nxh + i;

                //idx for single-halo array Mask->id[n]
                int x = i - 1;
                int y = j - 1;
                int z = k - 1;
                if (x < 0)
                    x = 0;
                if (y < 0)
                    y = 0;
                if (z < 0)
                    z = 0;
                if (x >= Nx)
                    x = Nx - 1;
                if (y >= Ny)
                    y = Ny - 1;
                if (z >= Nz)
                    z = Nz - 1;
                int n = z * Nx * Ny + y * Nx + x;
                VALUE = id[n];

                // Assign the affinity from the paired list
                for (unsigned int idx = 0; idx < NLABELS; idx++) {
                    //printf("idx=%i, value=%i, %i, \n",idx, VALUE,LabelList[idx]);
                    if (VALUE == LabelList[idx]) {
                        AFFINITY = AffinityList[idx];
                        label_count[idx] += 1.0;
                        idx = NLABELS;
                        //Mask->id[n] = 0; // set mask to zero since this is an immobile component
                    }
                }
                // fluid labels are reserved
                if (VALUE == 1)
                    AFFINITY = 1.0;
                else if (VALUE == 2)
                    AFFINITY = -1.0;
                phase[nh] = AFFINITY;
            }
        }
    }

    // Set Dm to match Mask
    for (int i = 0; i < Nx * Ny * Nz; i++)
        Dm->id[i] = Mask->id[i];

    for (size_t idx = 0; idx < NLABELS; idx++)
        label_count_global[idx] = Dm->Comm.sumReduce(label_count[idx]);

    if (rank == 0) {
        printf("Number of component labels: %lu \n", NLABELS);
        for (unsigned int idx = 0; idx < NLABELS; idx++) {
            VALUE = LabelList[idx];
            AFFINITY = AffinityList[idx];
            double volume_fraction =
                double(label_count_global[idx]) /
                double((Nx - 2) * (Ny - 2) * (Nz - 2) * nprocs);
            printf("   label=%d, affinity=%f, volume fraction==%f\n", VALUE,
                   AFFINITY, volume_fraction);
        }
    }

    //compute color gradient and laplacian of phase field
    double *ColorGrad_host, *mu_phi_host;
    ColorGrad_host = new double[3 * Np];
    mu_phi_host = new double[Np];

    double *Dst;
    Dst = new double[3 * 3 * 3];
    for (int kk = 0; kk < 3; kk++) {
        for (int jj = 0; jj < 3; jj++) {
            for (int ii = 0; ii < 3; ii++) {
                int index = kk * 9 + jj * 3 + ii;
                Dst[index] = sqrt(double(ii - 1) * double(ii - 1) +
                                  double(jj - 1) * double(jj - 1) +
                                  double(kk - 1) * double(kk - 1));
            }
        }
    }
    double w_face = 1.0 / 18.0;
    double w_edge = 1.0 / 36.0;
    double w_corner = 0.f;
    //local
    Dst[13] = 0.f;
    //faces
    Dst[4] = w_face;
    Dst[10] = w_face;
    Dst[12] = w_face;
    Dst[14] = w_face;
    Dst[16] = w_face;
    Dst[22] = w_face;
    // corners
    Dst[0] = w_corner;
    Dst[2] = w_corner;
    Dst[6] = w_corner;
    Dst[8] = w_corner;
    Dst[18] = w_corner;
    Dst[20] = w_corner;
    Dst[24] = w_corner;
    Dst[26] = w_corner;
    // edges
    Dst[1] = w_edge;
    Dst[3] = w_edge;
    Dst[5] = w_edge;
    Dst[7] = w_edge;
    Dst[9] = w_edge;
    Dst[11] = w_edge;
    Dst[15] = w_edge;
    Dst[17] = w_edge;
    Dst[19] = w_edge;
    Dst[21] = w_edge;
    Dst[23] = w_edge;
    Dst[25] = w_edge;

    double cs2_inv = 3.0; //inverse of c_s^2 for D3Q19 lattice
    int width =
        2; //For better readability: make halo width explicity wherever possible
    for (int k = width; k < Nzh - width; k++) {
        for (int j = width; j < Nyh - width; j++) {
            for (int i = width; i < Nxh - width; i++) {

                //idx for double-halo array 'phase'
                int nh = k * Nxh * Nyh + j * Nxh + i;

                int idx = Map(i - width + 1, j - width + 1, k - width + 1);
                if (!(idx < 0)) {
                    double phi_x = 0.f;
                    double phi_y = 0.f;
                    double phi_z = 0.f;
                    double phi_Lap = 0.f; //Laplacian of the phase field
                    for (int kk = 0; kk < 3; kk++) {
                        for (int jj = 0; jj < 3; jj++) {
                            for (int ii = 0; ii < 3; ii++) {

                                int index = kk * 9 + jj * 3 + ii;
                                double weight = Dst[index];

                                int idi = i + ii - 1;
                                int idj = j + jj - 1;
                                int idk = k + kk - 1;

                                if (idi < 0)
                                    idi = 0;
                                if (idj < 0)
                                    idj = 0;
                                if (idk < 0)
                                    idk = 0;
                                if (!(idi < Nxh))
                                    idi = Nxh - 1;
                                if (!(idj < Nyh))
                                    idj = Nyh - 1;
                                if (!(idk < Nzh))
                                    idk = Nzh - 1;

                                int nn = idk * Nxh * Nyh + idj * Nxh + idi;
                                double vec_x = double(ii - 1);
                                double vec_y = double(jj - 1);
                                double vec_z = double(kk - 1);
                                double GWNS = phase[nn];
                                double GWNS_local = phase[nh];
                                phi_x += GWNS * weight * vec_x;
                                phi_y += GWNS * weight * vec_y;
                                phi_z += GWNS * weight * vec_z;
                                phi_Lap +=
                                    weight *
                                    (GWNS -
                                     GWNS_local); //Laplacian of the phase field
                            }
                        }
                    }
                    //store color gradient
                    ColorGrad_host[idx + 0 * Np] = cs2_inv * phi_x;
                    ColorGrad_host[idx + 1 * Np] = cs2_inv * phi_y;
                    ColorGrad_host[idx + 2 * Np] = cs2_inv * phi_z;
                    //compute chemical potential
                    phi_Lap = 2.0 * cs2_inv * phi_Lap;
                    mu_phi_host[idx] = 4.0 * beta * phase[nh] *
                                           (phase[nh] + 1.0) *
                                           (phase[nh] - 1.0) -
                                       kappa * phi_Lap;
                }
            }
        }
    }

    //copy all data to device
    ScaLBL_CopyToDevice(Phi, phase, Nh * sizeof(double));
    ScaLBL_CopyToDevice(ColorGrad, ColorGrad_host, 3 * Np * sizeof(double));
    ScaLBL_CopyToDevice(mu_phi, mu_phi_host, Np * sizeof(double));
    ScaLBL_Comm->Barrier();
    comm.barrier();

    //debug
    //save the phase field and check it
    //FILE *OUTFILE;
    //sprintf(LocalRankFilename,"Phase_Init.%05i.raw",rank);
    //OUTFILE = fopen(LocalRankFilename,"wb");
    //fwrite(phase,8,Nh,OUTFILE);
    //fclose(OUTFILE);

    DoubleArray PhaseField(Nx, Ny, Nz);
    FILE *OUTFILE;

    getData_RegularLayout(mu_phi_host, PhaseField);
    sprintf(LocalRankFilename, "Chem_Init.%05i.raw", rank);
    OUTFILE = fopen(LocalRankFilename, "wb");
    fwrite(PhaseField.data(), 8, N, OUTFILE);
    fclose(OUTFILE);

    getData_RegularLayout(&ColorGrad_host[0], PhaseField);
    FILE *CGX_FILE;
    sprintf(LocalRankFilename, "Gradient_X_Init.%05i.raw", rank);
    CGX_FILE = fopen(LocalRankFilename, "wb");
    fwrite(PhaseField.data(), 8, N, CGX_FILE);
    fclose(CGX_FILE);

    getData_RegularLayout(&ColorGrad_host[Np], PhaseField);
    FILE *CGY_FILE;
    sprintf(LocalRankFilename, "Gradient_Y_Init.%05i.raw", rank);
    CGY_FILE = fopen(LocalRankFilename, "wb");
    fwrite(PhaseField.data(), 8, N, CGY_FILE);
    fclose(CGY_FILE);

    getData_RegularLayout(&ColorGrad_host[2 * Np], PhaseField);
    FILE *CGZ_FILE;
    sprintf(LocalRankFilename, "Gradient_Z_Init.%05i.raw", rank);
    CGZ_FILE = fopen(LocalRankFilename, "wb");
    fwrite(PhaseField.data(), 8, N, CGZ_FILE);
    fclose(CGZ_FILE);

    delete[] phase;
    delete[] ColorGrad_host;
    delete[] mu_phi_host;
    delete[] Dst;
}

void ScaLBL_FreeLeeModel::Initialize_TwoFluid() {
    /*
	 * This function initializes two-fluid Lee model
	 */
    if (rank == 0)
        printf("Initializing phase field, chemical potential and color "
               "gradient\n");
    AssignComponentLabels_ChemPotential_ColorGrad(); //initialize phase field Phi

    if (rank == 0)
        printf("Initializing distributions for momentum transport\n");
    ScaLBL_D3Q19_FreeLeeModel_TwoFluid_Init(gqbar, mu_phi, ColorGrad, Fx, Fy,
                                            Fz, Np);

    if (rank == 0)
        printf("Initializing density field and distributions for phase-field "
               "transport\n");
    ScaLBL_FreeLeeModel_PhaseField_Init(dvcMap, Phi, Den, hq, ColorGrad, rhoA,
                                        rhoB, tauM, W, 0,
                                        ScaLBL_Comm->LastExterior(), Np);
    ScaLBL_FreeLeeModel_PhaseField_Init(
        dvcMap, Phi, Den, hq, ColorGrad, rhoA, rhoB, tauM, W,
        ScaLBL_Comm->FirstInterior(), ScaLBL_Comm->LastInterior(), Np);

    if (Restart == true) {
        //TODO need to revise this function
        if (rank == 0) {
            printf("Reading restart file! \n");
        }

        // Read in the restart file to CPU buffers
        int *TmpMap;
        TmpMap = new int[Np];

        double *cPhi, *cDist, *cDen;
        cPhi = new double[N];
        cDen = new double[2 * Np];
        cDist = new double[19 * Np];
        ScaLBL_CopyToHost(TmpMap, dvcMap, Np * sizeof(int));
        //ScaLBL_CopyToHost(cPhi, Phi, N*sizeof(double));

        ifstream File(LocalRestartFile, ios::binary);
        int idx;
        double value, va, vb;
        for (int n = 0; n < Np; n++) {
            File.read((char *)&va, sizeof(va));
            File.read((char *)&vb, sizeof(vb));
            cDen[n] = va;
            cDen[Np + n] = vb;
        }
        for (int n = 0; n < Np; n++) {
            // Read the distributions
            for (int q = 0; q < 19; q++) {
                File.read((char *)&value, sizeof(value));
                cDist[q * Np + n] = value;
            }
        }
        File.close();

        for (int n = 0; n < ScaLBL_Comm->LastExterior(); n++) {
            va = cDen[n];
            vb = cDen[Np + n];
            value = (va - vb) / (va + vb);
            idx = TmpMap[n];
            if (!(idx < 0) && idx < N)
                cPhi[idx] = value;
        }
        for (int n = ScaLBL_Comm->FirstInterior();
             n < ScaLBL_Comm->LastInterior(); n++) {
            va = cDen[n];
            vb = cDen[Np + n];
            value = (va - vb) / (va + vb);
            idx = TmpMap[n];
            if (!(idx < 0) && idx < N)
                cPhi[idx] = value;
        }

        // Copy the restart data to the GPU
        ScaLBL_CopyToDevice(Den, cDen, 2 * Np * sizeof(double));
        ScaLBL_CopyToDevice(gqbar, cDist, 19 * Np * sizeof(double));
        ScaLBL_CopyToDevice(Phi, cPhi, N * sizeof(double));
        ScaLBL_Comm->Barrier();
        comm.barrier();

        if (rank == 0)
            printf("Initializing phase and density fields on device from "
                   "Restart\n");
        //TODO the following function is to be updated.
        //ScaLBL_FreeLeeModel_PhaseField_InitFromRestart(Den, hq, 0, ScaLBL_Comm->LastExterior(), Np);
        //ScaLBL_FreeLeeModel_PhaseField_InitFromRestart(Den, hq, ScaLBL_Comm->FirstInterior(), ScaLBL_Comm->LastInterior(), Np);
    }

    // establish reservoirs for external bC
    // TODO to be revised
    if (BoundaryCondition == 1 || BoundaryCondition == 2 ||
        BoundaryCondition == 3 || BoundaryCondition == 4) {
        if (Dm->kproc() == 0) {
            ScaLBL_SetSlice_z(Phi, 1.0, Nx, Ny, Nz, 0);
            ScaLBL_SetSlice_z(Phi, 1.0, Nx, Ny, Nz, 1);
            ScaLBL_SetSlice_z(Phi, 1.0, Nx, Ny, Nz, 2);
        }
        if (Dm->kproc() == nprocz - 1) {
            ScaLBL_SetSlice_z(Phi, -1.0, Nx, Ny, Nz, Nz - 1);
            ScaLBL_SetSlice_z(Phi, -1.0, Nx, Ny, Nz, Nz - 2);
            ScaLBL_SetSlice_z(Phi, -1.0, Nx, Ny, Nz, Nz - 3);
        }
    }
    //ScaLBL_CopyToHost(Averages->Phi.data(),Phi,N*sizeof(double));
}

void ScaLBL_FreeLeeModel::Initialize_SingleFluid() {
    /*
	 * This function initializes single-fluid Lee model
	 */
    if (rank == 0)
        printf("Initializing distributions for momentum transport\n");
    ScaLBL_D3Q19_FreeLeeModel_SingleFluid_Init(gqbar, Fx, Fy, Fz, Np);

    if (Restart == true) {
        //TODO need to revise this function
        //remove the phase-related part
    }
}

double ScaLBL_FreeLeeModel::Run_TwoFluid(int returntime) {

    int START_TIME = timestep;
    int EXIT_TIME = min(returntime, timestepMax);
    //************ MAIN ITERATION LOOP ***************************************/
    comm.barrier();
    auto t1 = std::chrono::system_clock::now();
    PROFILE_START("Loop");

    while (timestep < EXIT_TIME) {
        //if ( rank==0 ) { printf("Running timestep %i (%i MB)\n",timestep+1,(int)(Utilities::getMemoryUsage()/1048576)); }
        PROFILE_START("Update");
        // *************ODD TIMESTEP*************
        timestep++;
        //-------------------------------------------------------------------------------------------------------------------
        // Compute the Phase indicator field
        // Read for hq happens in this routine (requires communication)
        ScaLBL_Comm->SendD3Q7AA(hq, 0); //READ FROM NORMAL
        ScaLBL_D3Q7_AAodd_FreeLeeModel_PhaseField(
            NeighborList, dvcMap, hq, Den, Phi, rhoA, rhoB,
            ScaLBL_Comm->FirstInterior(), ScaLBL_Comm->LastInterior(), Np);
        //ScaLBL_D3Q7_AAodd_FreeLee_PhaseField(NeighborList, dvcMap, hq, Den, Phi, ColorGrad, Velocity, rhoA, rhoB, tauM, W, ScaLBL_Comm->FirstInterior(), ScaLBL_Comm->LastInterior(), Np);
        ScaLBL_Comm->RecvD3Q7AA(hq, 0); //WRITE INTO OPPOSITE
        ScaLBL_Comm->Barrier();
        ScaLBL_D3Q7_AAodd_FreeLeeModel_PhaseField(
            NeighborList, dvcMap, hq, Den, Phi, rhoA, rhoB, 0,
            ScaLBL_Comm->LastExterior(), Np);
        //ScaLBL_D3Q7_AAodd_FreeLee_PhaseField(NeighborList, dvcMap, hq, Den, Phi, ColorGrad, Velocity, rhoA, rhoB, tauM, W, 0, ScaLBL_Comm->LastExterior(), Np);

        // Perform the collision operation
        //ScaLBL_D3Q7_ComputePhaseField(dvcMap, hq, Den, Phi, rhoA, rhoB, 0, ScaLBL_Comm->LastInterior(), Np);
        //ScaLBL_Comm_WideHalo->Send(Phi);
        //ScaLBL_Comm_WideHalo->Recv(Phi);
        ScaLBL_Comm->SendD3Q19AA(gqbar); //READ FROM NORMAL
        if (BoundaryCondition > 0 && BoundaryCondition < 5) {
            //TODO to be revised
            // Need to add BC for hq!!!
            ScaLBL_Comm->Color_BC_z(dvcMap, Phi, Den, inletA, inletB);
            ScaLBL_Comm->Color_BC_Z(dvcMap, Phi, Den, outletA, outletB);
        }
        // Halo exchange for phase field
        ScaLBL_Comm_WideHalo->Send(Phi);
        //ScaLBL_D3Q19_AAodd_FreeLeeModel(NeighborList, dvcMap, gqbar, Den, Phi, mu_phi, Velocity, Pressure, ColorGrad, rhoA, rhoB, tauA, tauB,
        //		                        kappa, beta, W, Fx, Fy, Fz, Nxh, Nxh*Nyh, ScaLBL_Comm->FirstInterior(), ScaLBL_Comm->LastInterior(), Np);
        ScaLBL_D3Q19_AAodd_FreeLeeModel_Combined(
            NeighborList, dvcMap, gqbar, hq, Den, Phi, mu_phi, Velocity,
            Pressure, ColorGrad, rhoA, rhoB, tauA, tauB, tauM, kappa, beta, W,
            Fx, Fy, Fz, Nxh, Nxh * Nyh, ScaLBL_Comm->FirstInterior(),
            ScaLBL_Comm->LastInterior(), Np);
        ScaLBL_Comm_WideHalo->Recv(Phi);
        ScaLBL_Comm->RecvD3Q19AA(gqbar); //WRITE INTO OPPOSITE
        ScaLBL_Comm->Barrier();
        // Set BCs
        if (BoundaryCondition == 3) {
            ScaLBL_Comm->D3Q19_Pressure_BC_z(NeighborList, gqbar, din,
                                             timestep);
            ScaLBL_Comm->D3Q19_Pressure_BC_Z(NeighborList, gqbar, dout,
                                             timestep);
        }
        if (BoundaryCondition == 4) {
            din = ScaLBL_Comm->D3Q19_Flux_BC_z(NeighborList, gqbar, flux,
                                               timestep);
            ScaLBL_Comm->D3Q19_Pressure_BC_Z(NeighborList, gqbar, dout,
                                             timestep);
        } else if (BoundaryCondition == 5) {
            ScaLBL_Comm->D3Q19_Reflection_BC_z(gqbar);
            ScaLBL_Comm->D3Q19_Reflection_BC_Z(gqbar);
        }

        //ScaLBL_D3Q19_AAodd_FreeLeeModel(NeighborList, dvcMap, gqbar, Den, Phi, mu_phi, Velocity, Pressure, ColorGrad, rhoA, rhoB, tauA, tauB,
        //		                        kappa, beta, W, Fx, Fy, Fz, Nxh, Nxh*Nyh, 0, ScaLBL_Comm->LastExterior(), Np);
        ScaLBL_D3Q19_AAodd_FreeLeeModel_Combined(
            NeighborList, dvcMap, gqbar, hq, Den, Phi, mu_phi, Velocity,
            Pressure, ColorGrad, rhoA, rhoB, tauA, tauB, tauM, kappa, beta, W,
            Fx, Fy, Fz, Nxh, Nxh * Nyh, 0, ScaLBL_Comm->LastExterior(), Np);
        ScaLBL_Comm->Barrier();

        // *************EVEN TIMESTEP*************
        timestep++;
        // Compute the Phase indicator field
        ScaLBL_Comm->SendD3Q7AA(hq, 0); //READ FROM NORMA
        ScaLBL_D3Q7_AAeven_FreeLeeModel_PhaseField(
            dvcMap, hq, Den, Phi, rhoA, rhoB, ScaLBL_Comm->FirstInterior(),
            ScaLBL_Comm->LastInterior(), Np);
        //ScaLBL_D3Q7_AAeven_FreeLee_PhaseField(dvcMap, hq, Den, Phi, ColorGrad, Velocity, rhoA, rhoB, tauM, W, ScaLBL_Comm->FirstInterior(), ScaLBL_Comm->LastInterior(), Np);
        ScaLBL_Comm->RecvD3Q7AA(hq, 0); //WRITE INTO OPPOSITE
        ScaLBL_Comm->Barrier();
        ScaLBL_D3Q7_AAeven_FreeLeeModel_PhaseField(
            dvcMap, hq, Den, Phi, rhoA, rhoB, 0, ScaLBL_Comm->LastExterior(),
            Np);
        //ScaLBL_D3Q7_AAeven_FreeLee_PhaseField(dvcMap, hq, Den, Phi, ColorGrad, Velocity, rhoA, rhoB, tauM, W, 0, ScaLBL_Comm->LastExterior(), Np);

        // Perform the collision operation
        //ScaLBL_D3Q7_ComputePhaseField(dvcMap, hq, Den, Phi, rhoA, rhoB, ScaLBL_Comm->FirstInterior(), ScaLBL_Comm->LastInterior(), Np);
        //ScaLBL_Comm_WideHalo->Send(Phi);
        //ScaLBL_Comm_WideHalo->Recv(Phi);
        ScaLBL_Comm->SendD3Q19AA(gqbar); //READ FORM NORMAL
        if (BoundaryCondition > 0 && BoundaryCondition < 5) {
            ScaLBL_Comm->Color_BC_z(dvcMap, Phi, Den, inletA, inletB);
            ScaLBL_Comm->Color_BC_Z(dvcMap, Phi, Den, outletA, outletB);
        }
        // Halo exchange for phase field
        ScaLBL_Comm_WideHalo->Send(Phi);
        //ScaLBL_D3Q19_AAeven_FreeLeeModel(dvcMap, gqbar, Den, Phi, mu_phi, Velocity, Pressure, ColorGrad, rhoA, rhoB, tauA, tauB,
        //		                        kappa, beta, W, Fx, Fy, Fz, Nxh, Nxh*Nyh, ScaLBL_Comm->FirstInterior(), ScaLBL_Comm->LastInterior(), Np);
        ScaLBL_D3Q19_AAeven_FreeLeeModel_Combined(
            dvcMap, gqbar, hq, Den, Phi, mu_phi, Velocity, Pressure, ColorGrad,
            rhoA, rhoB, tauA, tauB, tauM, kappa, beta, W, Fx, Fy, Fz, Nxh,
            Nxh * Nyh, ScaLBL_Comm->FirstInterior(),
            ScaLBL_Comm->LastInterior(), Np);
        ScaLBL_Comm_WideHalo->Recv(Phi);
        ScaLBL_Comm->RecvD3Q19AA(gqbar); //WRITE INTO OPPOSITE
        ScaLBL_Comm->Barrier();
        // Set boundary conditions
        if (BoundaryCondition == 3) {
            ScaLBL_Comm->D3Q19_Pressure_BC_z(NeighborList, gqbar, din,
                                             timestep);
            ScaLBL_Comm->D3Q19_Pressure_BC_Z(NeighborList, gqbar, dout,
                                             timestep);
        } else if (BoundaryCondition == 4) {
            din = ScaLBL_Comm->D3Q19_Flux_BC_z(NeighborList, gqbar, flux,
                                               timestep);
            ScaLBL_Comm->D3Q19_Pressure_BC_Z(NeighborList, gqbar, dout,
                                             timestep);
        } else if (BoundaryCondition == 5) {
            ScaLBL_Comm->D3Q19_Reflection_BC_z(gqbar);
            ScaLBL_Comm->D3Q19_Reflection_BC_Z(gqbar);
        }
        //ScaLBL_D3Q19_AAeven_FreeLeeModel(dvcMap, gqbar, Den, Phi, mu_phi, Velocity, Pressure, ColorGrad, rhoA, rhoB, tauA, tauB,
        //		                        kappa, beta, W, Fx, Fy, Fz, Nxh, Nxh*Nyh, 0, ScaLBL_Comm->LastExterior(), Np);
        ScaLBL_D3Q19_AAeven_FreeLeeModel_Combined(
            dvcMap, gqbar, hq, Den, Phi, mu_phi, Velocity, Pressure, ColorGrad,
            rhoA, rhoB, tauA, tauB, tauM, kappa, beta, W, Fx, Fy, Fz, Nxh,
            Nxh * Nyh, 0, ScaLBL_Comm->LastExterior(), Np);
        ScaLBL_Comm->Barrier();
        //************************************************************************
        PROFILE_STOP("Update");
    }
    PROFILE_STOP("Loop");
    PROFILE_SAVE("lbpm_color_simulator", 1);
    //************************************************************************
    if (rank == 0)
        printf("---------------------------------------------------------------"
               "----\n");
    // Compute the walltime per timestep
    auto t2 = std::chrono::system_clock::now();
    double cputime = std::chrono::duration<double>(t2 - t1).count() /
                     (EXIT_TIME - START_TIME);
    // Performance obtained from each node
    double MLUPS = double(Np) / cputime / 1000000;

    return MLUPS;
}

void ScaLBL_FreeLeeModel::Run_SingleFluid() {
    int nprocs = nprocx * nprocy * nprocz;
    const RankInfoStruct rank_info(rank, nprocx, nprocy, nprocz);

    if (rank == 0) {
        printf("********************************************************\n");
        printf("No. of timesteps: %i \n", timestepMax);
        fflush(stdout);
    }

    //.......create and start timer............
    ScaLBL_Comm->Barrier();
    comm.barrier();
    //.........................................

    //************ MAIN ITERATION LOOP ***************************************/
    PROFILE_START("Loop");
    auto t1 = std::chrono::system_clock::now();
    while (timestep < timestepMax) {
        //if ( rank==0 ) { printf("Running timestep %i (%i MB)\n",timestep+1,(int)(Utilities::getMemoryUsage()/1048576)); }
        PROFILE_START("Update");
        // *************ODD TIMESTEP*************
        timestep++;
        //-------------------------------------------------------------------------------------------------------------------
        // Perform the collision operation
        ScaLBL_Comm->SendD3Q19AA(gqbar); //READ FROM NORMAL
        ScaLBL_D3Q19_AAodd_FreeLeeModel_SingleFluid_BGK(
            NeighborList, gqbar, Velocity, Pressure, tau, rho0, Fx, Fy, Fz,
            ScaLBL_Comm->FirstInterior(), ScaLBL_Comm->LastInterior(), Np);
        ScaLBL_Comm->RecvD3Q19AA(gqbar); //WRITE INTO OPPOSITE
        ScaLBL_Comm->Barrier();
        // Set boundary conditions
        // TODO to be revised!
        if (BoundaryCondition == 3) {
            ScaLBL_Comm->D3Q19_Pressure_BC_z(NeighborList, gqbar, din,
                                             timestep);
            ScaLBL_Comm->D3Q19_Pressure_BC_Z(NeighborList, gqbar, dout,
                                             timestep);
        }
        if (BoundaryCondition == 4) {
            din = ScaLBL_Comm->D3Q19_Flux_BC_z(NeighborList, gqbar, flux,
                                               timestep);
            ScaLBL_Comm->D3Q19_Pressure_BC_Z(NeighborList, gqbar, dout,
                                             timestep);
        } else if (BoundaryCondition == 5) {
            ScaLBL_Comm->D3Q19_Reflection_BC_z(gqbar);
            ScaLBL_Comm->D3Q19_Reflection_BC_Z(gqbar);
        }
        ScaLBL_D3Q19_AAodd_FreeLeeModel_SingleFluid_BGK(
            NeighborList, gqbar, Velocity, Pressure, tau, rho0, Fx, Fy, Fz, 0,
            ScaLBL_Comm->LastExterior(), Np);
        ScaLBL_Comm->Barrier();

        // *************EVEN TIMESTEP*************
        timestep++;
        //-------------------------------------------------------------------------------------------------------------------
        // Perform the collision operation
        ScaLBL_Comm->SendD3Q19AA(gqbar); //READ FORM NORMAL
        ScaLBL_D3Q19_AAeven_FreeLeeModel_SingleFluid_BGK(
            gqbar, Velocity, Pressure, tau, rho0, Fx, Fy, Fz,
            ScaLBL_Comm->FirstInterior(), ScaLBL_Comm->LastInterior(), Np);
        ScaLBL_Comm->RecvD3Q19AA(gqbar); //WRITE INTO OPPOSITE
        ScaLBL_Comm->Barrier();
        // Set boundary conditions
        // TODO to be revised!
        if (BoundaryCondition == 3) {
            ScaLBL_Comm->D3Q19_Pressure_BC_z(NeighborList, gqbar, din,
                                             timestep);
            ScaLBL_Comm->D3Q19_Pressure_BC_Z(NeighborList, gqbar, dout,
                                             timestep);
        } else if (BoundaryCondition == 4) {
            din = ScaLBL_Comm->D3Q19_Flux_BC_z(NeighborList, gqbar, flux,
                                               timestep);
            ScaLBL_Comm->D3Q19_Pressure_BC_Z(NeighborList, gqbar, dout,
                                             timestep);
        } else if (BoundaryCondition == 5) {
            ScaLBL_Comm->D3Q19_Reflection_BC_z(gqbar);
            ScaLBL_Comm->D3Q19_Reflection_BC_Z(gqbar);
        }
        ScaLBL_D3Q19_AAeven_FreeLeeModel_SingleFluid_BGK(
            gqbar, Velocity, Pressure, tau, rho0, Fx, Fy, Fz, 0,
            ScaLBL_Comm->LastExterior(), Np);
        ScaLBL_Comm->Barrier();
        //************************************************************************
        PROFILE_STOP("Update");
    }
    PROFILE_STOP("Loop");
    PROFILE_SAVE("lbpm_color_simulator", 1);
    //************************************************************************
    if (rank == 0)
        printf("---------------------------------------------------------------"
               "----\n");
    // Compute the walltime per timestep
    auto t2 = std::chrono::system_clock::now();
    double cputime = std::chrono::duration<double>(t2 - t1).count() / timestep;
    // Performance obtained from each node
    double MLUPS = double(Np) / cputime / 1000000;

    if (rank == 0)
        printf("********************************************************\n");
    if (rank == 0)
        printf("CPU time = %f \n", cputime);
    if (rank == 0)
        printf("Lattice update rate (per core)= %f MLUPS \n", MLUPS);
    MLUPS *= nprocs;
    if (rank == 0)
        printf("Lattice update rate (total)= %f MLUPS \n", MLUPS);
    if (rank == 0)
        printf("********************************************************\n");

    // ************************************************************************
}

void ScaLBL_FreeLeeModel::WriteDebug_TwoFluid() {
    // Copy back final phase indicator field and convert to regular layout
    DoubleArray PhaseData(Nxh, Nyh, Nzh);
    //ScaLBL_Comm->RegularLayout(Map,Phi,PhaseField);
    ScaLBL_CopyToHost(PhaseData.data(), Phi, sizeof(double) * Nh);
    /*
	IntArray MapData(Np);
	ScaLBL_CopyToHost(MapData.data(), dvcMap, sizeof(int)*Np);
	FILE *MAP;
	sprintf(LocalRankFilename,"Map.%05i.raw",rank);
	MAP = fopen(LocalRankFilename,"wb");
	fwrite(MapData.data(),4,Np,MAP);
	fclose(MAP);
	
	FILE *NB;
	//IntArray Neighbors(18,Np);
	//ScaLBL_CopyToHost(Neighbors.data(), NeighborList, sizeof(int)*Np*18);
	sprintf(LocalRankFilename,"neighbors.%05i.raw",rank);
	NB = fopen(LocalRankFilename,"wb");
	fwrite(NeighborList,4,18*Np,NB);
	fclose(NB);

	FILE *DIST;
	DoubleArray DistData(7, Np);
	ScaLBL_CopyToHost(DistData.data(), hq, 7*sizeof(double)*Np);
	sprintf(LocalRankFilename,"h.%05i.raw",rank);
	DIST = fopen(LocalRankFilename,"wb");
	fwrite(DistData.data(),8,7*Np,DIST);
	fclose(DIST);
	
	*/

    FILE *OUTFILE;
    sprintf(LocalRankFilename, "Phase.%05i.raw", rank);
    OUTFILE = fopen(LocalRankFilename, "wb");
    fwrite(PhaseData.data(), 8, Nh, OUTFILE);
    fclose(OUTFILE);

    DoubleArray PhaseField(Nx, Ny, Nz);
    FILE *DIST;
    for (int q = 0; q < 7; q++) {
        ScaLBL_Comm->RegularLayout(Map, &hq[q * Np], PhaseField);

        sprintf(LocalRankFilename, "h%i.%05i.raw", q, rank);
        DIST = fopen(LocalRankFilename, "wb");
        fwrite(PhaseField.data(), 8, Nx * Ny * Nz, DIST);
        fclose(DIST);
    }

    ScaLBL_Comm->RegularLayout(Map, Den, PhaseField);
    FILE *AFILE;
    sprintf(LocalRankFilename, "Density.%05i.raw", rank);
    AFILE = fopen(LocalRankFilename, "wb");
    fwrite(PhaseField.data(), 8, N, AFILE);
    fclose(AFILE);

    ScaLBL_Comm->RegularLayout(Map, Pressure, PhaseField);
    FILE *PFILE;
    sprintf(LocalRankFilename, "Pressure.%05i.raw", rank);
    PFILE = fopen(LocalRankFilename, "wb");
    fwrite(PhaseField.data(), 8, N, PFILE);
    fclose(PFILE);

    ScaLBL_Comm->RegularLayout(Map, mu_phi, PhaseField);
    FILE *CHEMFILE;
    sprintf(LocalRankFilename, "ChemPotential.%05i.raw", rank);
    CHEMFILE = fopen(LocalRankFilename, "wb");
    fwrite(PhaseField.data(), 8, N, CHEMFILE);
    fclose(CHEMFILE);

    ScaLBL_Comm->RegularLayout(Map, &Velocity[0], PhaseField);
    FILE *VELX_FILE;
    sprintf(LocalRankFilename, "Velocity_X.%05i.raw", rank);
    VELX_FILE = fopen(LocalRankFilename, "wb");
    fwrite(PhaseField.data(), 8, N, VELX_FILE);
    fclose(VELX_FILE);

    ScaLBL_Comm->RegularLayout(Map, &Velocity[Np], PhaseField);
    FILE *VELY_FILE;
    sprintf(LocalRankFilename, "Velocity_Y.%05i.raw", rank);
    VELY_FILE = fopen(LocalRankFilename, "wb");
    fwrite(PhaseField.data(), 8, N, VELY_FILE);
    fclose(VELY_FILE);

    ScaLBL_Comm->RegularLayout(Map, &Velocity[2 * Np], PhaseField);
    FILE *VELZ_FILE;
    sprintf(LocalRankFilename, "Velocity_Z.%05i.raw", rank);
    VELZ_FILE = fopen(LocalRankFilename, "wb");
    fwrite(PhaseField.data(), 8, N, VELZ_FILE);
    fclose(VELZ_FILE);

    ScaLBL_Comm->RegularLayout(Map, &ColorGrad[0], PhaseField);
    FILE *CGX_FILE;
    sprintf(LocalRankFilename, "Gradient_X.%05i.raw", rank);
    CGX_FILE = fopen(LocalRankFilename, "wb");
    fwrite(PhaseField.data(), 8, N, CGX_FILE);
    fclose(CGX_FILE);

    ScaLBL_Comm->RegularLayout(Map, &ColorGrad[Np], PhaseField);
    FILE *CGY_FILE;
    sprintf(LocalRankFilename, "Gradient_Y.%05i.raw", rank);
    CGY_FILE = fopen(LocalRankFilename, "wb");
    fwrite(PhaseField.data(), 8, N, CGY_FILE);
    fclose(CGY_FILE);

    ScaLBL_Comm->RegularLayout(Map, &ColorGrad[2 * Np], PhaseField);
    FILE *CGZ_FILE;
    sprintf(LocalRankFilename, "Gradient_Z.%05i.raw", rank);
    CGZ_FILE = fopen(LocalRankFilename, "wb");
    fwrite(PhaseField.data(), 8, N, CGZ_FILE);
    fclose(CGZ_FILE);
}

void ScaLBL_FreeLeeModel::WriteDebug_SingleFluid() {

    DoubleArray PhaseField(Nx, Ny, Nz);

    // Copy back final phase indicator field and convert to regular layout
    ScaLBL_Comm->RegularLayout(Map, Pressure, PhaseField);
    FILE *PFILE;
    sprintf(LocalRankFilename, "Pressure.%05i.raw", rank);
    PFILE = fopen(LocalRankFilename, "wb");
    fwrite(PhaseField.data(), 8, N, PFILE);
    fclose(PFILE);

    ScaLBL_Comm->RegularLayout(Map, &Velocity[0], PhaseField);
    FILE *VELX_FILE;
    sprintf(LocalRankFilename, "Velocity_X.%05i.raw", rank);
    VELX_FILE = fopen(LocalRankFilename, "wb");
    fwrite(PhaseField.data(), 8, N, VELX_FILE);
    fclose(VELX_FILE);

    ScaLBL_Comm->RegularLayout(Map, &Velocity[Np], PhaseField);
    FILE *VELY_FILE;
    sprintf(LocalRankFilename, "Velocity_Y.%05i.raw", rank);
    VELY_FILE = fopen(LocalRankFilename, "wb");
    fwrite(PhaseField.data(), 8, N, VELY_FILE);
    fclose(VELY_FILE);

    ScaLBL_Comm->RegularLayout(Map, &Velocity[2 * Np], PhaseField);
    FILE *VELZ_FILE;
    sprintf(LocalRankFilename, "Velocity_Z.%05i.raw", rank);
    VELZ_FILE = fopen(LocalRankFilename, "wb");
    fwrite(PhaseField.data(), 8, N, VELZ_FILE);
    fclose(VELZ_FILE);
}

void ScaLBL_FreeLeeModel::Create_DummyPhase_MGTest() {
    // Initialize communication structures in averaging domain
    for (int i = 0; i < Nx * Ny * Nz; i++)
        Dm->id[i] = Mask->id[i];
    Mask->CommInit();
    Np = Mask->PoreCount();
    //...........................................................................
    if (rank == 0)
        printf("Create ScaLBL_Communicator \n");
    // Create a communicator for the device (will use optimized layout)
    // ScaLBL_Communicator ScaLBL_Comm(Mask); // original
    ScaLBL_Comm =
        std::shared_ptr<ScaLBL_Communicator>(new ScaLBL_Communicator(Mask));
    //ScaLBL_Comm_Regular  = std::shared_ptr<ScaLBL_Communicator>(new ScaLBL_Communicator(Mask));
    ScaLBL_Comm_WideHalo = std::shared_ptr<ScaLBLWideHalo_Communicator>(
        new ScaLBLWideHalo_Communicator(Mask, 2));

    // create the layout for the LBM
    int Npad = (Np / 16 + 2) * 16;
    if (rank == 0)
        printf("Set up memory efficient layout, %i | %i | %i \n", Np, Npad, N);
    Map.resize(Nx, Ny, Nz);
    Map.fill(-2);
    auto neighborList = new int[18 * Npad];
    Np = ScaLBL_Comm->MemoryOptimizedLayoutAA(Map, neighborList,
                                              Mask->id.data(), Np, 1);
    comm.barrier();

    //...........................................................................
    //                MAIN  VARIABLES ALLOCATED HERE
    //...........................................................................
    // LBM variables
    if (rank == 0)
        printf("Allocating distributions \n");
    //......................device distributions.................................
    dist_mem_size = Np * sizeof(double);
    neighborSize = 18 * (Np * sizeof(int));
    //...........................................................................
    //ScaLBL_AllocateDeviceMemory((void **) &NeighborList, neighborSize);
    ScaLBL_AllocateDeviceMemory((void **)&dvcMap, sizeof(int) * Np);
    //ScaLBL_AllocateDeviceMemory((void **) &gqbar, 19*dist_mem_size);
    //ScaLBL_AllocateDeviceMemory((void **) &hq, 7*dist_mem_size);
    //ScaLBL_AllocateDeviceMemory((void **) &mu_phi, dist_mem_size);
    //ScaLBL_AllocateDeviceMemory((void **) &Den, dist_mem_size);
    ScaLBL_AllocateDeviceMemory((void **)&Phi, sizeof(double) * Nh);
    //ScaLBL_AllocateDeviceMemory((void **) &Pressure, sizeof(double)*Np);
    //ScaLBL_AllocateDeviceMemory((void **) &Velocity, 3*sizeof(double)*Np);
    ScaLBL_AllocateDeviceMemory((void **)&ColorGrad, 3 * sizeof(double) * Np);
    //...........................................................................
    // Update GPU data structures
    if (rank == 0)
        printf("Setting up device map and neighbor list \n");
    fflush(stdout);
    int *TmpMap;
    TmpMap = new int[Np];
    for (int k = 1; k < Nz - 1; k++) {
        for (int j = 1; j < Ny - 1; j++) {
            for (int i = 1; i < Nx - 1; i++) {
                int idx = Map(i, j, k);
                if (!(idx < 0))
                    TmpMap[idx] = ScaLBL_Comm_WideHalo->Map(i, j, k);
            }
        }
    }
    // check that TmpMap is valid
    for (int idx = 0; idx < ScaLBL_Comm->LastExterior(); idx++) {
        auto n = TmpMap[idx];
        if (n > Nxh * Nyh * Nzh) {
            printf("Bad value! idx=%i \n", n);
            TmpMap[idx] = Nxh * Nyh * Nzh - 1;
        }
    }
    for (int idx = ScaLBL_Comm->FirstInterior();
         idx < ScaLBL_Comm->LastInterior(); idx++) {
        auto n = TmpMap[idx];
        if (n > Nxh * Nyh * Nzh) {
            printf("Bad value! idx=%i \n", n);
            TmpMap[idx] = Nxh * Nyh * Nzh - 1;
        }
    }
    // copy the device map
    ScaLBL_CopyToDevice(dvcMap, TmpMap, sizeof(int) * Np);
    // copy the neighbor list
    //ScaLBL_CopyToDevice(NeighborList, neighborList, neighborSize);
    comm.barrier();

    double *phase;
    phase = new double[Nh];

    for (int k = 0; k < Nzh; k++) {
        for (int j = 0; j < Nyh; j++) {
            for (int i = 0; i < Nxh; i++) {

                //idx for double-halo array 'phase'
                int nh = k * Nxh * Nyh + j * Nxh + i;

                //idx for single-halo array Mask->id[n]
                int x = i - 1;
                int y = j - 1;
                int z = k - 1;
                if (x < 0)
                    x = 0;
                if (y < 0)
                    y = 0;
                if (z < 0)
                    z = 0;
                if (x >= Nx)
                    x = Nx - 1;
                if (y >= Ny)
                    y = Ny - 1;
                if (z >= Nz)
                    z = Nz - 1;
                int n = z * Nx * Ny + y * Nx + x;
                phase[nh] = id[n];
            }
        }
    }
    ScaLBL_CopyToDevice(Phi, phase, Nh * sizeof(double));
    ScaLBL_Comm->Barrier();
    comm.barrier();
    delete[] TmpMap;
    delete[] neighborList;
    delete[] phase;
}

void ScaLBL_FreeLeeModel::MGTest() {

    comm.barrier();

    ScaLBL_Comm_WideHalo->Send(Phi);
    ScaLBL_D3Q9_MGTest(dvcMap, Phi, ColorGrad, Nxh, Nxh * Nyh,
                       ScaLBL_Comm->FirstInterior(),
                       ScaLBL_Comm->LastInterior(), Np);
    ScaLBL_Comm_WideHalo->Recv(Phi);
    ScaLBL_D3Q9_MGTest(dvcMap, Phi, ColorGrad, Nxh, Nxh * Nyh, 0,
                       ScaLBL_Comm->LastExterior(), Np);

    //check the sum of ColorGrad
    double cgx_loc = 0.0;
    double cgy_loc = 0.0;
    double cgz_loc = 0.0;
    double cgx, cgy, cgz;
    double *ColorGrad_host;
    ColorGrad_host = new double[3 * Np];
    ScaLBL_CopyToHost(&ColorGrad_host[0], &ColorGrad[0],
                      3 * Np * sizeof(double));
    for (int i = ScaLBL_Comm->FirstInterior(); i < ScaLBL_Comm->LastInterior();
         i++) {
        cgx_loc += ColorGrad_host[0 * Np + i];
        cgy_loc += ColorGrad_host[1 * Np + i];
        cgz_loc += ColorGrad_host[2 * Np + i];
    }
    for (int i = 0; i < ScaLBL_Comm->LastExterior(); i++) {
        cgx_loc += ColorGrad_host[0 * Np + i];
        cgy_loc += ColorGrad_host[1 * Np + i];
        cgz_loc += ColorGrad_host[2 * Np + i];
    }
    cgx = Dm->Comm.sumReduce(cgx_loc);
    cgy = Dm->Comm.sumReduce(cgy_loc);
    cgz = Dm->Comm.sumReduce(cgz_loc);
    if (rank == 0) {
        printf("Sum of all x-component of the mixed gradient = %.2g \n", cgx);
        printf("Sum of all y-component of the mixed gradient = %.2g \n", cgy);
        printf("Sum of all z-component of the mixed gradient = %.2g \n", cgz);
    }

    delete[] ColorGrad_host;
}