/* This file is part of the Palabos library.
 *
 * Copyright (C) 2011-2015 FlowKit Sarl
 * Route d'Oron 2
 * 1010 Lausanne, Switzerland
 * E-mail contact: contact@flowkit.com
 *
 * The most recent release of Palabos can be downloaded at 
 * <http://www.palabos.org/>
 *
 * The library Palabos is free software: you can redistribute it and/or
 * modify it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * The library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

#ifndef MULTI_FREE_SURFACE_MODEL_3D_HH
#define MULTI_FREE_SURFACE_MODEL_3D_HH

#include "core/globalDefs.h"
#include "core/block3D.h"
#include "latticeBoltzmann/geometricOperationTemplates.h"
#include "atomicBlock/dataProcessor3D.h"
#include "atomicBlock/blockLattice3D.h"
#include "atomicBlock/atomicContainerBlock3D.h"
#include "multiPhysics/freeSurfaceModel3D.h"
#include "multiPhysics/freeSurfaceTemplates.h"
#include "multiPhysics/multiFreeSurfaceModel3D.h"

#include <cmath>

namespace plb {

/* *************** Class MultiFreeSurfaceOneWayCoupling3D ******************************************* */

template<typename T,template<typename U> class Descriptor>
void MultiFreeSurfaceOneWayCoupling3D<T,Descriptor>
        ::processGenericBlocks(Box3D domain,std::vector<AtomicBlock3D*> atomicBlocks)
{
    PLB_ASSERT(atomicBlocks.size() == 20);

    using namespace twoPhaseFlag;

    std::vector<AtomicBlock3D*> atomicBlocks1, atomicBlocks2;

    for (int i = 0; i < 10; i++)
        atomicBlocks1.push_back(atomicBlocks[i]);

    for (int i = 0; i < 10; i++)
        atomicBlocks2.push_back(atomicBlocks[10 + i]);

    FreeSurfaceProcessorParam3D<T,Descriptor> param1(atomicBlocks1);
    FreeSurfaceProcessorParam3D<T,Descriptor> param2(atomicBlocks2);

    for (plint iX=domain.x0; iX<=domain.x1; ++iX) {
        for (plint iY=domain.y0; iY<=domain.y1; ++iY) {
            for (plint iZ=domain.z0; iZ<=domain.z1; ++iZ) {         
                if (param1.flag(iX,iY,iZ)==interface && isWet(param2.flag(iX,iY,iZ))) {
                    // Velocity coupling
                    T rhoBar1; 
                    Array<T,3> j1;
                    momentTemplates<T,Descriptor>::get_rhoBar_j(param1.cell(iX,iY,iZ), rhoBar1, j1);
                    T density1 = Descriptor<T>::fullRho(rhoBar1);

                    T rhoBar2; 
                    Array<T,3> j2;
                    momentTemplates<T,Descriptor>::get_rhoBar_j(param2.cell(iX,iY,iZ), rhoBar2, j2);
                    T density2 = Descriptor<T>::fullRho(rhoBar2);

                    T tau = T(1)/param1.cell(iX,iY,iZ).getDynamics().getOmega();
                    for (plint iD=0; iD<3; ++iD) {
                        //j1[iD] += interactionStrength*density1*(j2[iD]/density2 - j1[iD]/density1)*tau;
                        j1[iD] += interactionStrength*rhoDefault1*(j2[iD]/density2 - j1[iD]/density1)*tau;
                    }

                    param1.setMomentum(iX,iY,iZ, j1);
                }
            }
        }
    }   
}

/* *************** Class MultiFreeSurfaceVelocityContinuityCoupling3D ******************************************* */

template<typename T,template<typename U> class Descriptor>
void MultiFreeSurfaceVelocityContinuityCoupling3D<T,Descriptor>
        ::processGenericBlocks(Box3D domain,std::vector<AtomicBlock3D*> atomicBlocks)
{
    PLB_ASSERT(atomicBlocks.size() == 20);

    using namespace twoPhaseFlag;

    std::vector<AtomicBlock3D*> atomicBlocks1, atomicBlocks2;

    for (int i = 0; i < 10; i++)
        atomicBlocks1.push_back(atomicBlocks[i]);

    for (int i = 0; i < 10; i++)
        atomicBlocks2.push_back(atomicBlocks[10 + i]);

    FreeSurfaceProcessorParam3D<T,Descriptor> param1(atomicBlocks1);
    FreeSurfaceProcessorParam3D<T,Descriptor> param2(atomicBlocks2);

    T convertRho1to2 = rhoDefault2/rhoDefault1;
    T convertRho2to1 = rhoDefault1/rhoDefault2;

    for (plint iX=domain.x0; iX<=domain.x1; ++iX) {
        for (plint iY=domain.y0; iY<=domain.y1; ++iY) {
            for (plint iZ=domain.z0; iZ<=domain.z1; ++iZ) {         
                
                Array<T,3> velocity((T) 0.0, (T) 0.0, (T) 0.0);

                // Enforce continuity of velocity only when both fluid cells belong to the interface.
                if (param1.flag(iX,iY,iZ)==interface && param2.flag(iX,iY,iZ)==interface) {
                //if ((param2.flag(iX,iY,iZ)==interface && isWet(param1.flag(iX,iY,iZ))) ||
                //    (param1.flag(iX,iY,iZ)==interface && isWet(param2.flag(iX,iY,iZ)))) {
                    T rhoBar1; 
                    Array<T,3> j1;
                    momentTemplates<T,Descriptor>::get_rhoBar_j(param1.cell(iX,iY,iZ), rhoBar1, j1);
                    T density1 = Descriptor<T>::fullRho(rhoBar1);

                    T rhoBar2; 
                    Array<T,3> j2;
                    momentTemplates<T,Descriptor>::get_rhoBar_j(param2.cell(iX,iY,iZ), rhoBar2, j2);
                    T density2 = Descriptor<T>::fullRho(rhoBar2);

                    velocity = 0.5 * (j1/density1 + j2/density2);
                }

                if (param1.flag(iX,iY,iZ)==interface && isWet(param2.flag(iX,iY,iZ))) {
                    // Compute the average pressure from fluid cells of both fluids.
                    T averageDensity = T();
                    plint numNeighbors = 0;
                    for (plint i=-1; i<2; i++) {
                        for (plint j=-1; j<2; j++) {
                            for (plint k=-1; k<2; k++) {
                                if (isFullWet(param1.flag(iX+i,iY+j,iZ+k))) {
                                    ++numNeighbors;
                                    averageDensity += param1.getDensity(iX+i,iY+j,iZ+k);
                                }
                                if (isFullWet(param2.flag(iX+i,iY+j,iZ+k))) {
                                    ++numNeighbors;
                                    averageDensity += convertRho2to1 * param2.getDensity(iX+i,iY+j,iZ+k);
                                }
                            }
                        }
                    }
                    if (numNeighbors==0) {
                        averageDensity = rhoDefault1;
                    }
                    else {
                        averageDensity /= (T)numNeighbors;
                    }
                    param1.setDensity(iX,iY,iZ, averageDensity);

                    param1.volumeFraction(iX,iY,iZ) = param1.mass(iX,iY,iZ) / averageDensity;
                    if (param2.flag(iX,iY,iZ) == interface) {
                        param1.setMomentum(iX,iY,iZ, param1.getDensity(iX,iY,iZ) * velocity);
                    }
                }

                if (param2.flag(iX,iY,iZ)==interface && isWet(param1.flag(iX,iY,iZ))) {
                    // Compute the average pressure from fluid cells of both fluids.
                    T averageDensity = T();
                    plint numNeighbors = 0;
                    for (plint i=-1; i<2; i++) {
                        for (plint j=-1; j<2; j++) {
                            for (plint k=-1; k<2; k++) {
                                if (isFullWet(param1.flag(iX+i,iY+j,iZ+k))) {
                                    ++numNeighbors;
                                    averageDensity += convertRho1to2 * param1.getDensity(iX+i,iY+j,iZ+k);
                                }
                                if (isFullWet(param2.flag(iX+i,iY+j,iZ+k))) {
                                    ++numNeighbors;
                                    averageDensity += param2.getDensity(iX+i,iY+j,iZ+k);
                                }
                            }
                        }
                    }
                    if (numNeighbors==0) {
                        averageDensity = rhoDefault2;
                    }
                    else {
                        averageDensity /= (T)numNeighbors;
                    }
                    param2.setDensity(iX,iY,iZ, averageDensity);

                    param2.volumeFraction(iX,iY,iZ) = param2.mass(iX,iY,iZ) / averageDensity;
                    if (param1.flag(iX,iY,iZ) == interface) {
                        param2.setMomentum(iX,iY,iZ, param2.getDensity(iX,iY,iZ) * velocity);
                    }
                }
            }
        }
    }   
}

/* *************** Class MultiFreeSurfaceRepellingForceCoupling3D ******************************************* */

template<typename T,template<typename U> class Descriptor>
T MultiFreeSurfaceRepellingForceCoupling3D<T,Descriptor>::deltaFunction(T r, T h)
{
    Precision precision = floatingPointPrecision<T>();
    T eps = getEpsilon<T>(precision);

    PLB_ASSERT(r > (T) 0 || std::fabs(r) <= eps);
    PLB_ASSERT(h > (T) 0 && std::fabs(h) >  eps);

    static T pi = std::acos((T) -1.0);

    T delta = 0.0;
    if (r < h || std::fabs(r - h) <= eps) {
        delta = 0.5 * (1.0 + std::cos(pi * r / h));
    }

    return delta;
}

template<typename T,template<typename U> class Descriptor>
void MultiFreeSurfaceRepellingForceCoupling3D<T,Descriptor>
        ::processGenericBlocks(Box3D domain,std::vector<AtomicBlock3D*> atomicBlocks)
{
    PLB_ASSERT(atomicBlocks.size() == 20);

    using namespace twoPhaseFlag;

    std::vector<AtomicBlock3D*> atomicBlocks1, atomicBlocks2;

    for (int i = 0; i < 10; i++)
        atomicBlocks1.push_back(atomicBlocks[i]);

    for (int i = 0; i < 10; i++)
        atomicBlocks2.push_back(atomicBlocks[10 + i]);

    FreeSurfaceProcessorParam3D<T,Descriptor> param1(atomicBlocks1);
    FreeSurfaceProcessorParam3D<T,Descriptor> param2(atomicBlocks2);

    T convertRho1to2 = rhoDefault2/rhoDefault1;
    T convertRho2to1 = rhoDefault1/rhoDefault2;

    for (plint iX=domain.x0; iX<=domain.x1; ++iX) {
        for (plint iY=domain.y0; iY<=domain.y1; ++iY) {
            for (plint iZ=domain.z0; iZ<=domain.z1; ++iZ) {         

                // Implement the repelling force in terms of a momentum correction. Remember that fluids can have strong
                // density differences. In order to have the units of a momentum, this correction term must be
                // multiplied by the density.
                if (isWet(param1.flag(iX,iY,iZ)) && isWet(param2.flag(iX,iY,iZ))) {
                    // Add a force to both fluids. Remember the envelope is 3.
                    Array<T,3> force1((T) 0, (T) 0, (T) 0);
                    for (plint i=-2; i<3; i++) {
                        plint I = iX + i;
                        for (plint j=-2; j<3; j++) {
                            plint J = iY + j;
                            for (plint k=-2; k<3; k++) {
                                plint K = iZ + k;
                                if (isWet(param1.flag(I,J,K)) && isWet(param2.flag(I,J,K))) {
                                    T VFdiff = param1.volumeFraction(I,J,K) + param2.volumeFraction(I,J,K) - 1.0;
                                    if (VFdiff>(T)0.0) {
                                        Array<T,3> outwardNormalToFluid1 = param1.getNormal(I,J,K) - param2.getNormal(I,J,K);
                                        outwardNormalToFluid1 /= norm(outwardNormalToFluid1);
                                        T r = std::sqrt((T)i*i + (T)j*j + (T)k*k);
                                        T delta = deltaFunction(r, (T) 3); // h is 3 because envelope is 3.
                                        force1 += -VFdiff*delta*outwardNormalToFluid1;
                                    }
                                }
                            }
                        }
                    }
                    force1 *= interactionStrength;

                    // force1 is an acceleration.
                    T tau1 = T(1)/param1.cell(iX,iY,iZ).getDynamics().getOmega();
                    param1.setMomentum(iX,iY,iZ, param1.getMomentum(iX,iY,iZ) + rhoDefault1*tau1*force1);
                    T tau2 = T(1)/param2.cell(iX,iY,iZ).getDynamics().getOmega();
                    param2.setMomentum(iX,iY,iZ, param2.getMomentum(iX,iY,iZ) + rhoDefault2*tau2*(-force1));

                    /*
                    // force1 is a force.
                    T tau1 = T(1)/param1.cell(iX,iY,iZ).getDynamics().getOmega();
                    param1.setMomentum(iX,iY,iZ, param1.getMomentum(iX,iY,iZ) + tau1*force1);
                    T tau2 = T(1)/param2.cell(iX,iY,iZ).getDynamics().getOmega();
                    param2.setMomentum(iX,iY,iZ, param2.getMomentum(iX,iY,iZ) + tau2*(-force1));
                    */
                }

                Array<T,3> j;

                if (param1.flag(iX,iY,iZ)==interface && isWet(param2.flag(iX,iY,iZ))) {
                    // Compute the average pressure from fluid cells of both fluids.
                    T averageDensity = T();
                    plint numNeighbors = 0;
                    for (plint i=-1; i<2; i++) {
                        for (plint j=-1; j<2; j++) {
                            for (plint k=-1; k<2; k++) {
                                if (isFullWet(param1.flag(iX+i,iY+j,iZ+k))) {
                                    ++numNeighbors;
                                    averageDensity += param1.getDensity(iX+i,iY+j,iZ+k);
                                }
                                if (isFullWet(param2.flag(iX+i,iY+j,iZ+k))) {
                                    ++numNeighbors;
                                    averageDensity += convertRho2to1 * param2.getDensity(iX+i,iY+j,iZ+k);
                                }
                            }
                        }
                    }
                    if (numNeighbors==0) {
                        averageDensity = rhoDefault1;
                    }
                    else {
                        averageDensity /= (T)numNeighbors;
                    }
                    T oldDensity = param1.getDensity(iX,iY,iZ);
                    param1.setDensity(iX,iY,iZ, averageDensity);
                    param1.volumeFraction(iX,iY,iZ) = param1.mass(iX,iY,iZ) / averageDensity;
                    j = param1.getMomentum(iX,iY,iZ);
                    j *= averageDensity / oldDensity;
                    param1.setMomentum(iX,iY,iZ, j);


                    /*
                    // Implement the repelling force in terms of a momentum correction. Remember that fluids can have strong
                    // density differences. In order to have the units of a momentum, this correction term must be
                    // multiplied by the density.
                    T VFdiff = param1.volumeFraction(iX,iY,iZ) + param2.volumeFraction(iX,iY,iZ) - 1.0;
                    Array<T,3> normalToInterface = param1.getNormal(iX, iY, iZ);                    
                    //Array<T,3> normalToInterface = -param2.getNormal(iX, iY, iZ);                    
                    if (VFdiff>(T)0.0) {
                        T tau = T(1)/param2.cell(iX,iY,iZ).getDynamics().getOmega();
                        param2.setMomentum(iX,iY,iZ, param2.getMomentum(iX,iY,iZ) +
                                interactionStrength*rhoDefault2*VFdiff*tau*normalToInterface);
                    }
                    */
                }

                if (param2.flag(iX,iY,iZ)==interface && isWet(param1.flag(iX,iY,iZ))) {
                    // Compute the average pressure from fluid cells of both fluids.
                    T averageDensity = T();
                    plint numNeighbors = 0;
                    for (plint i=-1; i<2; i++) {
                        for (plint j=-1; j<2; j++) {
                            for (plint k=-1; k<2; k++) {
                                if (isFullWet(param1.flag(iX+i,iY+j,iZ+k))) {
                                    ++numNeighbors;
                                    averageDensity += convertRho1to2 * param1.getDensity(iX+i,iY+j,iZ+k);
                                }
                                if (isFullWet(param2.flag(iX+i,iY+j,iZ+k))) {
                                    ++numNeighbors;
                                    averageDensity += param2.getDensity(iX+i,iY+j,iZ+k);
                                }
                            }
                        }
                    }
                    if (numNeighbors==0) {
                        averageDensity = rhoDefault2;
                    }
                    else {
                        averageDensity /= (T)numNeighbors;
                    }

                    T oldDensity = param2.getDensity(iX,iY,iZ);
                    param2.setDensity(iX,iY,iZ, averageDensity);
                    param2.volumeFraction(iX,iY,iZ) = param2.mass(iX,iY,iZ) / averageDensity;
                    j = param2.getMomentum(iX,iY,iZ);
                    j *= averageDensity / oldDensity;
                    param2.setMomentum(iX,iY,iZ, j);

                    /*
                    // Implement the repelling force in terms of a momentum correction. Remember that fluids can have strong
                    // density differences. In order to have the units of a momentum, this correction term must be
                    // multiplied by the density.
                    T VFdiff = param1.volumeFraction(iX,iY,iZ) + param2.volumeFraction(iX,iY,iZ) - 1.0;
                    Array<T,3> normalToInterface = param2.getNormal(iX,iY,iZ);
                    //Array<T,3> normalToInterface = -param1.getNormal(iX,iY,iZ);
                    if (VFdiff>(T)0.0) {
                        T tau = T(1)/param1.cell(iX,iY,iZ).getDynamics().getOmega();
                        param1.setMomentum(iX,iY,iZ, param1.getMomentum(iX,iY,iZ) +
                                interactionStrength*rhoDefault1*VFdiff*tau*normalToInterface);
                    }
                    */
                }
            }
        }
    }   
}

/* *************** Class MultiFreeSurfaceComplexCoupling3D ******************************************* */

template<typename T,template<typename U> class Descriptor>
T MultiFreeSurfaceComplexCoupling3D<T,Descriptor>::deltaFunction(T r, T h)
{
    Precision precision = floatingPointPrecision<T>();
    T eps = getEpsilon<T>(precision);

    PLB_ASSERT(r > (T) 0 || std::fabs(r) <= eps);
    PLB_ASSERT(h > (T) 0 && std::fabs(h) >  eps);

    static T pi = std::acos((T) -1.0);

    T delta = 0.0;
    if (r < h || std::fabs(r - h) <= eps) {
        delta = 0.5 * (1.0 + std::cos(pi * r / h));
    }

    return delta;
}

template<typename T,template<typename U> class Descriptor>
void MultiFreeSurfaceComplexCoupling3D<T,Descriptor>
        ::processGenericBlocks(Box3D domain,std::vector<AtomicBlock3D*> atomicBlocks)
{
    PLB_ASSERT(atomicBlocks.size() == 20);

    using namespace twoPhaseFlag;

    std::vector<AtomicBlock3D*> atomicBlocks1, atomicBlocks2;

    for (int i = 0; i < 10; i++)
        atomicBlocks1.push_back(atomicBlocks[i]);

    for (int i = 0; i < 10; i++)
        atomicBlocks2.push_back(atomicBlocks[10 + i]);

    FreeSurfaceProcessorParam3D<T,Descriptor> param1(atomicBlocks1);
    FreeSurfaceProcessorParam3D<T,Descriptor> param2(atomicBlocks2);

    T convertRho1to2 = rhoDefault2/rhoDefault1;
    T convertRho2to1 = rhoDefault1/rhoDefault2;

    for (plint iX=domain.x0; iX<=domain.x1; ++iX) {
        for (plint iY=domain.y0; iY<=domain.y1; ++iY) {
            for (plint iZ=domain.z0; iZ<=domain.z1; ++iZ) {         

                Array<T,3> velocity((T) 0.0, (T) 0.0, (T) 0.0);

                if (param1.flag(iX,iY,iZ)==interface && param2.flag(iX,iY,iZ)==interface) {
                    T rhoBar1; 
                    Array<T,3> j1;
                    momentTemplates<T,Descriptor>::get_rhoBar_j(param1.cell(iX,iY,iZ), rhoBar1, j1);
                    T density1 = Descriptor<T>::fullRho(rhoBar1);

                    T rhoBar2; 
                    Array<T,3> j2;
                    momentTemplates<T,Descriptor>::get_rhoBar_j(param2.cell(iX,iY,iZ), rhoBar2, j2);
                    T density2 = Descriptor<T>::fullRho(rhoBar2);

                    velocity = 0.5 * (j1/density1 + j2/density2);
                }

                // Compute the repelling force.
                
                Array<T,3> force1((T) 0.0, (T) 0.0, (T) 0.0);

                if (isWet(param1.flag(iX,iY,iZ)) && isWet(param2.flag(iX,iY,iZ))) {
                    // Remember the envelope is 3.
                    for (plint i=-2; i<3; i++) {
                        plint I = iX + i;
                        for (plint j=-2; j<3; j++) {
                            plint J = iY + j;
                            for (plint k=-2; k<3; k++) {
                                plint K = iZ + k;
                                if (isWet(param1.flag(I,J,K)) && isWet(param2.flag(I,J,K))) {
                                    T VFdiff = param1.volumeFraction(I,J,K) + param2.volumeFraction(I,J,K) - 1.0;
                                    if (VFdiff>(T)0.0) {
                                        Array<T,3> outwardNormalToFluid1 = param1.getNormal(I,J,K) - param2.getNormal(I,J,K);
                                        outwardNormalToFluid1 /= norm(outwardNormalToFluid1);
                                        T r = std::sqrt((T)i*i + (T)j*j + (T)k*k);
                                        T delta = deltaFunction(r, (T) 3); // h is 3 because envelope is 3.
                                        force1 += -VFdiff*delta*outwardNormalToFluid1;
                                    }
                                }
                            }
                        }
                    }
                    force1 *= interactionStrength;
                }

                if (param1.flag(iX,iY,iZ)==interface && isWet(param2.flag(iX,iY,iZ))) {
                    // Compute the average pressure from fluid cells of both fluids.
                    T averageDensity = T();
                    plint numNeighbors = 0;
                    for (plint i=-1; i<2; i++) {
                        for (plint j=-1; j<2; j++) {
                            for (plint k=-1; k<2; k++) {
                                if (isFullWet(param1.flag(iX+i,iY+j,iZ+k))) {
                                    ++numNeighbors;
                                    averageDensity += param1.getDensity(iX+i,iY+j,iZ+k);
                                }
                                if (isFullWet(param2.flag(iX+i,iY+j,iZ+k))) {
                                    ++numNeighbors;
                                    averageDensity += convertRho2to1 * param2.getDensity(iX+i,iY+j,iZ+k);
                                }
                            }
                        }
                    }
                    if (numNeighbors==0) {
                        averageDensity = rhoDefault1;
                    }
                    else {
                        averageDensity /= (T)numNeighbors;
                    }
                    param1.setDensity(iX,iY,iZ, averageDensity);
                    param1.volumeFraction(iX,iY,iZ) = param1.mass(iX,iY,iZ) / averageDensity;
                    if (param2.flag(iX,iY,iZ) == interface) {
                        param1.setMomentum(iX,iY,iZ, param1.getDensity(iX,iY,iZ) * velocity);
                    }
                }

                if (param2.flag(iX,iY,iZ)==interface && isWet(param1.flag(iX,iY,iZ))) {
                    // Compute the average pressure from fluid cells of both fluids.
                    T averageDensity = T();
                    plint numNeighbors = 0;
                    for (plint i=-1; i<2; i++) {
                        for (plint j=-1; j<2; j++) {
                            for (plint k=-1; k<2; k++) {
                                if (isFullWet(param1.flag(iX+i,iY+j,iZ+k))) {
                                    ++numNeighbors;
                                    averageDensity += convertRho1to2 * param1.getDensity(iX+i,iY+j,iZ+k);
                                }
                                if (isFullWet(param2.flag(iX+i,iY+j,iZ+k))) {
                                    ++numNeighbors;
                                    averageDensity += param2.getDensity(iX+i,iY+j,iZ+k);
                                }
                            }
                        }
                    }
                    if (numNeighbors==0) {
                        averageDensity = rhoDefault2;
                    }
                    else {
                        averageDensity /= (T)numNeighbors;
                    }
                    param2.setDensity(iX,iY,iZ, averageDensity);
                    param2.volumeFraction(iX,iY,iZ) = param2.mass(iX,iY,iZ) / averageDensity;
                    if (param1.flag(iX,iY,iZ) == interface) {
                        param2.setMomentum(iX,iY,iZ, param2.getDensity(iX,iY,iZ) * velocity);
                    }
                }

                // Implement the repelling force in terms of a momentum correction. Remember that fluids can have strong
                // density differences. In order to have the units of a momentum, this correction term must be
                // multiplied by the density.
                
                if (isWet(param1.flag(iX,iY,iZ)) && isWet(param2.flag(iX,iY,iZ))) {
                    // Add a force to both fluids. Remember the envelope is 3.

                    // force1 is an acceleration.
                    T tau1 = T(1)/param1.cell(iX,iY,iZ).getDynamics().getOmega();
                    param1.setMomentum(iX,iY,iZ, param1.getMomentum(iX,iY,iZ) + rhoDefault1*tau1*force1);
                    T tau2 = T(1)/param2.cell(iX,iY,iZ).getDynamics().getOmega();
                    param2.setMomentum(iX,iY,iZ, param2.getMomentum(iX,iY,iZ) + rhoDefault2*tau2*(-force1));

                    /*
                    // force1 is a force.
                    T tau1 = T(1)/param1.cell(iX,iY,iZ).getDynamics().getOmega();
                    param1.setMomentum(iX,iY,iZ, param1.getMomentum(iX,iY,iZ) + tau1*force1);
                    T tau2 = T(1)/param2.cell(iX,iY,iZ).getDynamics().getOmega();
                    param2.setMomentum(iX,iY,iZ, param2.getMomentum(iX,iY,iZ) + tau2*(-force1));
                    */
                }
            }
        }
    }   
}

}  // namespace plb

#endif  // MULTI_FREE_SURFACE_MODEL_3D_HH

