/*  Lattice Boltzmann sample, written in C++, using the OpenLB
 *  library
 *
 *  Copyright (C) 2008 Orestis Malaspinas, Andrea Parmigiani
 *  E-mail contact: info@openlb.net
 *  The most recent release of OpenLB can be downloaded at
 *  <http://www.openlb.net/>
 *
 *  This program is free software; you can redistribute it and/or
 *  modify it under the terms of the GNU General Public License
 *  as published by the Free Software Foundation; either version 2
 *  of the License, or (at your option) any later version.
 *
 *  This program 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 General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public
 *  License along with this program; if not, write to the Free
 *  Software Foundation, Inc., 51 Franklin Street, Fifth Floor,
 *  Boston, MA  02110-1301, USA.
 */

/* rayleighTaylor2d.cpp:
 * Rayleigh-Taylor instability in 2D, generated by a heavy
 * fluid penetrating a light one. The multi-component fluid model
 * by X. Shan and H. Chen is used. This example shows the usage
 * of multicomponent flow and periodic boundaries.
 */


#include "olb2D.h"
#include "olb2D.hh"   // use only generic version!
#include <cstdlib>
#include <iostream>

using namespace olb;
using namespace olb::descriptors;
using namespace olb::graphics;
using namespace std;

typedef double T;
#define DESCRIPTOR ShanChenDynOmegaForcedD2Q9Descriptor


// Parameters for the simulation setup
const int nx   = 400;
const int ny   = 200;
const int maxIter  = 20000;


// Stores geometry information in form of material numbers
void prepareGeometry( SuperGeometry2D<T>& superGeometry ) {

  OstreamManager clout( std::cout,"prepareGeometry" );
  clout << "Prepare Geometry ..." << std::endl;

  // Sets material number for fluid and boundary
  superGeometry.rename( 0,1 );

  Vector<T,2> origin1( -.5, -.5 );
  Vector<T,2> origin2( -.5,ny/2. );
  Vector<T,2> origin3( -.5, ny-1.5 );
  Vector<T,2> extend1( nx+2, 1. );
  Vector<T,2> extend2( nx+2., ny/2.+2. );

  IndicatorCuboid2D<T> bottom( extend1, origin1 );
  IndicatorCuboid2D<T> upper( extend2, origin2 );
  IndicatorCuboid2D<T> top( extend1, origin3 );

  superGeometry.rename( 1,2,upper );
  superGeometry.rename( 1,3,bottom );
  superGeometry.rename( 2,4,top );

  // Removes all not needed boundary voxels outside the surface
  //superGeometry.clean();
  // Removes all not needed boundary voxels inside the surface
  superGeometry.innerClean();
  superGeometry.checkForErrors();

  superGeometry.print();

  clout << "Prepare Geometry ... OK" << std::endl;
}

void prepareLattice( SuperLattice2D<T, DESCRIPTOR>& sLatticeOne,
                     SuperLattice2D<T, DESCRIPTOR>& sLatticeTwo,
                     Dynamics<T, DESCRIPTOR>& bulkDynamics1,
                     Dynamics<T, DESCRIPTOR>& bulkDynamics2,
                     Dynamics<T, DESCRIPTOR>& bounceBackRho0,
                     Dynamics<T, DESCRIPTOR>& bounceBackRho1,
                     SuperGeometry2D<T>& superGeometry ) {

  OstreamManager clout( std::cout,"prepareLattice" );
  clout << "Prepare Lattice ..." << std::endl;

  // The setup is: periodicity along horizontal direction, bounce-back on top
  // and bottom. The upper half is initially filled with fluid 1 + random noise,
  // and the lower half with fluid 2. Only fluid 1 experiences a forces,
  // directed downwards.

  // define lattice Dynamics
  sLatticeOne.defineDynamics( superGeometry, 0, &instances::getNoDynamics<T, DESCRIPTOR>() );
  sLatticeTwo.defineDynamics( superGeometry, 0, &instances::getNoDynamics<T, DESCRIPTOR>() );

  sLatticeOne.defineDynamics( superGeometry, 1, &bulkDynamics1 );
  sLatticeOne.defineDynamics( superGeometry, 2, &bulkDynamics1 );
  sLatticeOne.defineDynamics( superGeometry, 3, &bulkDynamics1 );
  sLatticeOne.defineDynamics( superGeometry, 4, &bulkDynamics1 );
  sLatticeTwo.defineDynamics( superGeometry, 1, &bulkDynamics2 );
  sLatticeTwo.defineDynamics( superGeometry, 2, &bulkDynamics2 );
  sLatticeTwo.defineDynamics( superGeometry, 3, &bulkDynamics2 );
  sLatticeTwo.defineDynamics( superGeometry, 4, &bulkDynamics2 );

  sLatticeOne.defineDynamics( superGeometry, 3, &bounceBackRho0 );
  sLatticeTwo.defineDynamics( superGeometry, 3, &bounceBackRho1 );
  sLatticeOne.defineDynamics( superGeometry, 4, &bounceBackRho1 );
  sLatticeTwo.defineDynamics( superGeometry, 4, &bounceBackRho0 );

  clout << "Prepare Lattice ... OK" << std::endl;
}

void setBoundaryValues( SuperLattice2D<T, DESCRIPTOR>& sLatticeOne,
                        SuperLattice2D<T, DESCRIPTOR>& sLatticeTwo,
                        T force, int iT, SuperGeometry2D<T>& superGeometry ) {

  if ( iT==0 ) {

    AnalyticalConst2D<T,T> noise( 4.e-2 );
    std::vector<T> v( 2,T() );
    AnalyticalConst2D<T,T> zeroV( v );
    AnalyticalConst2D<T,T> zero( 0. );
    AnalyticalLinear2D<T,T> one( 0.,-force*invCs2<T,DESCRIPTOR>(),0.98+force*ny*invCs2<T,DESCRIPTOR>() );
    AnalyticalConst2D<T,T> onePlus( 0.98+force*ny/2.*invCs2<T,DESCRIPTOR>() );
    AnalyticalRandom2D<T,T> random;
    AnalyticalIdentity2D<T,T> randomOne( random*noise+one );
    AnalyticalIdentity2D<T,T> randomPlus( random*noise+onePlus );
    std::vector<T> F( 2,T() );
    F[1] = -force;
    AnalyticalConst2D<T,T> f( F );

    // for each material set the defineRhou and the Equilibrium

    sLatticeOne.defineRhoU( superGeometry, 1, zero, zeroV );
    sLatticeOne.iniEquilibrium( superGeometry, 1, zero, zeroV );
    sLatticeOne.defineField<descriptors::EXTERNAL_FORCE>( superGeometry, 1, f );
    sLatticeTwo.defineRhoU( superGeometry, 1, randomPlus, zeroV );
    sLatticeTwo.iniEquilibrium( superGeometry, 1, randomPlus, zeroV );

    sLatticeOne.defineRhoU( superGeometry, 2, randomOne, zeroV );
    sLatticeOne.iniEquilibrium( superGeometry, 2, randomOne, zeroV );
    sLatticeOne.defineField<descriptors::EXTERNAL_FORCE>( superGeometry, 2, f );
    sLatticeTwo.defineRhoU( superGeometry, 2, zero, zeroV );
    sLatticeTwo.iniEquilibrium( superGeometry, 2, zero, zeroV );

    /*sLatticeOne.defineRhoU(superGeometry, 3, zero, zeroV);
    sLatticeOne.iniEquilibrium(superGeometry, 3, zero, zeroV);
    sLatticeOne.defineField<descriptors::EXTERNAL_FORCE>(superGeometry, 3, f);
    sLatticeTwo.defineRhoU(superGeometry, 3, one, zeroV);
    sLatticeTwo.iniEquilibrium(superGeometry, 3, one, zeroV);

    sLatticeOne.defineRhoU(superGeometry, 4, one, zeroV);
    sLatticeOne.iniEquilibrium(superGeometry, 4, one, zeroV);
    sLatticeOne.defineField<descriptors::EXTERNAL_FORCE>(superGeometry, 4, f);
    sLatticeTwo.defineRhoU(superGeometry, 4, zero, zeroV);
    sLatticeTwo.iniEquilibrium(superGeometry, 4, zero, zeroV);*/

    // Make the lattice ready for simulation
    sLatticeOne.initialize();
    sLatticeTwo.initialize();
  }
}

void getResults( SuperLattice2D<T, DESCRIPTOR>&    sLatticeTwo,
                 SuperLattice2D<T, DESCRIPTOR>&    sLatticeOne, int iT,
                 SuperGeometry2D<T>& superGeometry, Timer<T>& timer ) {

  OstreamManager clout( std::cout,"getResults" );
  SuperVTMwriter2D<T> vtmWriter( "rayleighTaylor2dsLatticeOne" );

  const int vtkIter = 100;
  const int statIter = 10;

  if ( iT==0 ) {
    // Writes the geometry, cuboid no. and rank no. as vti file for visualization
    SuperLatticeGeometry2D<T, DESCRIPTOR> geometry( sLatticeOne, superGeometry );
    SuperLatticeCuboid2D<T, DESCRIPTOR> cuboid( sLatticeOne );
    SuperLatticeRank2D<T, DESCRIPTOR> rank( sLatticeOne );
    vtmWriter.write( geometry );
    vtmWriter.write( cuboid );
    vtmWriter.write( rank );
    vtmWriter.createMasterFile();
  }

  // Get statistics
  if ( iT%statIter==0 && iT > 0 ) {
    // Timer console output
    timer.update( iT );
    timer.printStep();

    clout << "averageRhoFluidOne="   << sLatticeOne.getStatistics().getAverageRho();
    clout << "; averageRhoFluidTwo=" << sLatticeTwo.getStatistics().getAverageRho() << std::endl;
  }

  // Writes the VTK files
  if ( iT%vtkIter==0 ) {
    clout << "Writing VTK ..." << std::endl;
    SuperLatticeVelocity2D<T, DESCRIPTOR> velocity( sLatticeOne );
    SuperLatticeDensity2D<T, DESCRIPTOR> density( sLatticeOne );
    vtmWriter.addFunctor( velocity );
    vtmWriter.addFunctor( density );
    vtmWriter.write( iT );

    BlockReduction2D2D<T> planeReduction( density, 600, BlockDataSyncMode::ReduceOnly );
    // write output as JPEG
    heatmap::write(planeReduction, iT);

    clout << "Writing VTK ... OK" << std::endl;
  }
}

int main( int argc, char *argv[] ) {

  // === 1st Step: Initialization ===

  olbInit( &argc, &argv );
  singleton::directories().setOutputDir( "./tmp/" );
  OstreamManager clout( std::cout,"main" );

  const T omega1 = 1.0;
  const T omega2 = 1.0;
  const T G      = 3.;
  T force        = 30./( T )ny/( T )ny;

  // === 2nd Step: Prepare Geometry ===
  // Instantiation of a cuboidGeometry with weights

#ifdef PARALLEL_MODE_MPI
  CuboidGeometry2D<T> cGeometry( 0, 0, 1, nx, ny, singleton::mpi().getSize() );
#else
  CuboidGeometry2D<T> cGeometry( 0, 0, 1, nx, ny, 1 );
#endif

  cGeometry.setPeriodicity( true, false );

  HeuristicLoadBalancer<T> loadBalancer( cGeometry );

  SuperGeometry2D<T> superGeometry( cGeometry, loadBalancer, 2 );

  prepareGeometry( superGeometry );

  // === 3rd Step: Prepare Lattice ===

  SuperLattice2D<T, DESCRIPTOR> sLatticeOne( superGeometry );
  SuperLattice2D<T, DESCRIPTOR> sLatticeTwo( superGeometry );

  ForcedBGKdynamics<T, DESCRIPTOR> bulkDynamics1 (
    omega1, instances::getExternalVelocityMomenta<T,DESCRIPTOR>() );
  ForcedBGKdynamics<T, DESCRIPTOR> bulkDynamics2 (
    omega2, instances::getExternalVelocityMomenta<T,DESCRIPTOR>() );

  // A bounce-back node with fictitious density 1,
  //   which is experienced by the partner fluid
  BounceBack<T, DESCRIPTOR> bounceBackRho1( 1. );
  // A bounce-back node with fictitious density 0,
  //   which is experienced by the partner fluid
  BounceBack<T, DESCRIPTOR> bounceBackRho0( 0. );

  std::vector<T> rho0;
  rho0.push_back( 1 );
  rho0.push_back( 1 );
  PsiEqualsRho<T,T> interactionPotential;
  ShanChenForcedGenerator2D<T,DESCRIPTOR> coupling( G,rho0,interactionPotential );

  sLatticeOne.addLatticeCoupling( coupling, sLatticeTwo );

  prepareLattice( sLatticeOne, sLatticeTwo, bulkDynamics1, bulkDynamics2,
                  bounceBackRho0, bounceBackRho1, superGeometry );

  // === 4th Step: Main Loop with Timer ===
  int iT = 0;
  clout << "starting simulation..." << endl;
  Timer<T> timer( maxIter, superGeometry.getStatistics().getNvoxel() );
  timer.start();

  for ( iT=0; iT<maxIter; ++iT ) {

    // === 5th Step: Definition of Initial and Boundary Conditions ===
    setBoundaryValues( sLatticeOne, sLatticeTwo, force, iT, superGeometry );

    // === 6th Step: Collide and Stream Execution ===
    sLatticeOne.collideAndStream();
    sLatticeTwo.collideAndStream();

    sLatticeOne.communicate();
    sLatticeTwo.communicate();

    sLatticeOne.executeCoupling();
    //sLatticeTwo.executeCoupling();

    // === 7th Step: Computation and Output of the Results ===
    getResults( sLatticeTwo, sLatticeOne, iT, superGeometry, timer );
  }

  timer.stop();
  timer.printSummary();
}