/*
 *  Lattice Boltzmann grid refinement sample, written in C++,
 *  using the OpenLB library
 *
 *  Copyright (C) 2019 Adrian Kummerländer
 *  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.
 */

#include "olb2D.h"
#ifndef OLB_PRECOMPILED
#include "olb2D.hh"
#endif

#include <vector>

using namespace olb;

typedef double T;

#define DESCRIPTOR descriptors::D2Q9Descriptor

/// Setup geometry relative to cylinder diameter as defined by [SchaeferTurek96]
const T cylinderD = 1.0;

const int N = 5;               // resolution of the model
const T lx  = 22  * cylinderD; // length of the channel
const T ly  = 4.1 * cylinderD; // height of the channel
const T Re  = 100.;            // Reynolds number
const T uLat = 0.05;           // lattice velocity
const T maxPhysT = 60.;        // max. simulation time in s, SI unit


void prepareGeometry(Grid2D<T,DESCRIPTOR>& grid)
{
  OstreamManager clout(std::cout,"prepareGeometry");
  clout << "Prepare Geometry ..." << std::endl;

  auto& converter = grid.getConverter();
  auto& sGeometry = grid.getSuperGeometry();

  sGeometry.rename(0,1);

  const T physSpacing = converter.getPhysDeltaX();

  // Set material number for bounce back boundaries
  {
    const Vector<T,2> wallExtend {lx+physSpacing,  physSpacing};
    const Vector<T,2> wallOrigin {-physSpacing/2, -physSpacing/2};

    IndicatorCuboid2D<T> lowerWall(wallExtend, wallOrigin);
    sGeometry.rename(1,2,lowerWall);

    IndicatorCuboid2D<T> upperWall(wallExtend, wallOrigin + Vector<T,2> {0,ly-physSpacing/2});
    sGeometry.rename(1,2,upperWall);
  }

  // Set material number for inflow and outflow
  {
    const Vector<T,2> extend { physSpacing/2,  ly-physSpacing/2};
    const Vector<T,2> origin {-physSpacing/4, -physSpacing/4};

    IndicatorCuboid2D<T> inflow(extend, origin);
    sGeometry.rename(1,3,inflow);

    IndicatorCuboid2D<T> outflow(extend, origin + Vector<T,2> {lx,0});
    sGeometry.rename(1,4,outflow);
  }

  // Set material number for vertically centered cylinder
  {
    const Vector<T,2> origin {2*cylinderD, 2*cylinderD};
    IndicatorCircle2D<T> obstacle(origin, cylinderD/2);
    sGeometry.rename(1,5,obstacle);
  }

  sGeometry.clean();
  sGeometry.innerClean();
  sGeometry.checkForErrors();
  sGeometry.print();

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

void disableRefinedArea(Grid2D<T,DESCRIPTOR>& coarseGrid,
                        RefiningGrid2D<T,DESCRIPTOR>& fineGrid)
{
  auto& sGeometry = coarseGrid.getSuperGeometry();
  auto refinedOverlap = fineGrid.getRefinedOverlap();
  sGeometry.rename(1,0,*refinedOverlap);
  sGeometry.rename(2,0,*refinedOverlap);
  sGeometry.rename(5,0,*refinedOverlap);
}

void prepareLattice(Grid2D<T,DESCRIPTOR>& grid)
{
  OstreamManager clout(std::cout,"prepareLattice");
  clout << "Prepare lattice ..." << std::endl;

  auto& converter = grid.getConverter();
  auto& sGeometry = grid.getSuperGeometry();
  auto& sLattice  = grid.getSuperLattice();

  Dynamics<T,DESCRIPTOR>& bulkDynamics = grid.addDynamics(
      std::unique_ptr<Dynamics<T,DESCRIPTOR>>(
        new BGKdynamics<T,DESCRIPTOR>(
          grid.getConverter().getLatticeRelaxationFrequency(),
          instances::getBulkMomenta<T,DESCRIPTOR>())));

  sOnLatticeBoundaryCondition2D<T,DESCRIPTOR>& sBoundaryCondition = grid.getOnLatticeBoundaryCondition();
  createLocalBoundaryCondition2D<T,DESCRIPTOR>(sBoundaryCondition);

  const T omega = converter.getLatticeRelaxationFrequency();

  sLattice.defineDynamics(sGeometry, 0, &instances::getNoDynamics<T,DESCRIPTOR>());
  sLattice.defineDynamics(sGeometry, 1, &bulkDynamics); // bulk
  sLattice.defineDynamics(sGeometry, 2, &bulkDynamics); // walls
  sLattice.defineDynamics(sGeometry, 3, &bulkDynamics); // inflow
  sLattice.defineDynamics(sGeometry, 4, &bulkDynamics); // outflow
  sLattice.defineDynamics(sGeometry, 5, &instances::getBounceBack<T,DESCRIPTOR>()); // cylinder

  sBoundaryCondition.addVelocityBoundary(sGeometry, 2, omega);
  sBoundaryCondition.addVelocityBoundary(sGeometry, 3, omega);
  sBoundaryCondition.addPressureBoundary(sGeometry, 4, omega);

  AnalyticalConst2D<T,T> rho0(1.0);
  AnalyticalConst2D<T,T> u0(0.0, 0.0);

  auto materials = sGeometry.getMaterialIndicator({1, 2, 3, 4, 5});
  sLattice.defineRhoU(materials, rho0, u0);
  sLattice.iniEquilibrium(materials, rho0, u0);

  sLattice.initialize();

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

void setBoundaryValues(Grid2D<T,DESCRIPTOR>& grid, int iT)
{
  auto& converter = grid.getConverter();
  auto& sGeometry = grid.getSuperGeometry();
  auto& sLattice  = grid.getSuperLattice();

  const int iTmaxStart = converter.getLatticeTime(0.2*maxPhysT);
  const int iTupdate = 10;

  if ( iT % iTupdate == 0 && iT <= iTmaxStart ) {
    PolynomialStartScale<T,T> StartScale(iTmaxStart, 1);

    T iTvec[1] { T(iT) };
    T frac[1] { };
    StartScale(frac, iTvec);

    const T maxVelocity = converter.getCharLatticeVelocity() * frac[0];
    const T radius = ly/2;
    std::vector<T> axisPoint{0, ly/2};
    std::vector<T> axisDirection{1, 0};
    Poiseuille2D<T> u(axisPoint, axisDirection, maxVelocity, radius);

    sLattice.defineU(sGeometry, 3, u);
  }
}

void getResults(const std::string& prefix,
                Grid2D<T,DESCRIPTOR>& grid,
                int iT)
{
  OstreamManager clout(std::cout,"getResults");

  auto& converter = grid.getConverter();
  auto& sLattice  = grid.getSuperLattice();
  auto& sGeometry = grid.getSuperGeometry();

  SuperVTMwriter2D<T> vtmWriter(prefix + "cylinder2d");
  SuperLatticePhysVelocity2D<T,DESCRIPTOR> velocity(sLattice, converter);
  SuperLatticePhysPressure2D<T,DESCRIPTOR> pressure(sLattice, converter);
  SuperLatticeGeometry2D<T,DESCRIPTOR> geometry(sLattice, sGeometry);
  SuperLatticeKnudsen2D<T,DESCRIPTOR> knudsen(sLattice);
  vtmWriter.addFunctor(geometry);
  vtmWriter.addFunctor(velocity);
  vtmWriter.addFunctor(pressure);
  vtmWriter.addFunctor(knudsen);

  if (iT==0) {
    vtmWriter.createMasterFile();
  }

  if (iT%100==0) {
    vtmWriter.write(iT);
  }
}

int main(int argc, char* argv[])
{
  olbInit(&argc, &argv);
  singleton::directories().setOutputDir("./tmp/");
  OstreamManager clout(std::cout,"main");

  const Vector<T,2> coarseOrigin {0.0, 0.0};
  const Vector<T,2> coarseExtend {lx, ly};
  IndicatorCuboid2D<T> coarseCuboid(coarseExtend, coarseOrigin);

  Grid2D<T,DESCRIPTOR> coarseGrid(coarseCuboid, N, LatticeVelocity<T>(uLat), Re);
  prepareGeometry(coarseGrid);

  const auto coarseDeltaX = coarseGrid.getConverter().getPhysDeltaX();

  const Vector<T,2> fineExtend {8*cylinderD, ly-2*coarseDeltaX};
  const Vector<T,2> fineOrigin {0.75*cylinderD, coarseDeltaX};

  auto& fineGrid = coarseGrid.refine(fineOrigin, fineExtend);
  prepareGeometry(fineGrid);

  const Vector<T,2> fineExtendB {10*coarseDeltaX, ly};
  const Vector<T,2> fineOriginB {lx-10*coarseDeltaX, 0};

  auto& fineOutflowGrid = coarseGrid.refine(fineOriginB, fineExtendB, false);
  prepareGeometry(fineOutflowGrid);

  {
    const Vector<T,2> origin = fineOutflowGrid.getOrigin();
    const Vector<T,2> extend = fineOutflowGrid.getExtend();

    const Vector<T,2> extendY = {0,extend[1]};

    coarseGrid.addFineCoupling(fineOutflowGrid, origin, extendY);
    coarseGrid.addCoarseCoupling(fineOutflowGrid, origin + Vector<T,2> {coarseDeltaX,0}, extendY);

    IndicatorCuboid2D<T> refined(extend, origin + Vector<T,2> {2*coarseDeltaX,0});

    coarseGrid.getSuperGeometry().rename(1,0,refined);
    coarseGrid.getSuperGeometry().rename(2,0,refined);
    coarseGrid.getSuperGeometry().rename(4,0,refined);
  }

  disableRefinedArea(coarseGrid, fineGrid);

  const Vector<T,2> fineExtend2 {4*cylinderD, 2*cylinderD};
  const Vector<T,2> fineOrigin2 {1*cylinderD, 2*cylinderD-fineExtend2[1]/2};

  auto& fineGrid2 = fineGrid.refine(fineOrigin2, fineExtend2);
  prepareGeometry(fineGrid2);

  disableRefinedArea(fineGrid, fineGrid2);

  prepareLattice(coarseGrid);
  prepareLattice(fineGrid);
  prepareLattice(fineOutflowGrid);
  prepareLattice(fineGrid2);

  clout << "Starting simulation..." << endl;
  Timer<T> timer(
    coarseGrid.getConverter().getLatticeTime(maxPhysT),
    coarseGrid.getSuperGeometry().getStatistics().getNvoxel());
  timer.start();
  const int statIter = coarseGrid.getConverter().getLatticeTime(0.5);

  for (int iT = 0; iT < coarseGrid.getConverter().getLatticeTime(maxPhysT); ++iT) {
    setBoundaryValues(coarseGrid, iT);

    coarseGrid.collideAndStream();

    getResults("level0_", coarseGrid, iT);
    getResults("level1_", fineGrid,   iT);
    getResults("level1_outflow_", fineOutflowGrid, iT);
    getResults("level2_", fineGrid2,  iT);

    if (iT%statIter == 0) {
      timer.update(iT);
      timer.printStep();

      coarseGrid.getSuperLattice().getStatistics().print(iT, coarseGrid.getConverter().getPhysTime(iT));
    }
  }

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