/*
 *  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 = 0.1;
const int N         = 5;                      // resolution of the cylinder
const T   deltaR    = cylinderD / N;          // coarse lattice spacing
const T   lx        = 22*cylinderD + deltaR;  // length of the channel
const T   ly        = 4.1*cylinderD + deltaR; // height of the channel
const T   cylinderX = 2*cylinderD;
const T   cylinderY = 2*cylinderD + deltaR/2;

const T   Re        = 100.;                   // Reynolds number
const T   tau       = 0.51;                   // relaxation time
const T   maxPhysT  = 16.;                    // max. simulation time in s, SI unit

const Characteristics<T> PhysCharacteristics(
  0.1,         // char. phys. length
  1.0,         // char. phys. velocity
  0.1/Re,      // phsy. kinematic viscosity
  1.0);        // char. phys. density

void prepareGeometry(Grid2D<T,DESCRIPTOR>& grid, Vector<T,2> origin, Vector<T,2> extend)
{
  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 channel walls
  {
    const Vector<T,2> wallExtend { extend[0]+physSpacing/2, physSpacing/2 };
    const Vector<T,2> wallOrigin = origin - physSpacing/4;

    IndicatorCuboid2D<T> lowerWall(wallExtend, wallOrigin);
    sGeometry.rename(1,2,lowerWall);
  }
  {
    const Vector<T,2> wallExtend { extend[0]+physSpacing/2, physSpacing/2 };
    const Vector<T,2> wallOrigin { origin[0]-physSpacing/4, extend[1]-physSpacing/4 };

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

  // Set material number for inflow and outflow
  {
    const Vector<T,2> inflowExtend { physSpacing/2, extend[1]+physSpacing/4 };
    const Vector<T,2> inflowOrigin = origin - physSpacing/4;

    IndicatorCuboid2D<T> inflow(inflowExtend, inflowOrigin);
    sGeometry.rename(1,3,inflow);
  }
  {
    const Vector<T,2> outflowExtend { physSpacing/2, extend[1]+physSpacing/4 };
    const Vector<T,2> outflowOrigin { extend[0]-physSpacing/4, origin[0]-physSpacing/4 };

    IndicatorCuboid2D<T> outflow(outflowExtend, outflowOrigin);
    sGeometry.rename(1,4,outflow);
  }

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

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

  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.reset(*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();
  //createInterpBoundaryCondition2D<T,DESCRIPTOR>(sBoundaryCondition);
  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});
  sLattice.defineRhoU(materials, rho0, u0);
  sLattice.iniEquilibrium(materials, rho0, u0);

  sLattice.initialize();

  clout << "Prepare lattice ... OK" << std::endl;
  sGeometry.print();
}

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.4*16);
  const int iTupdate = 5;

  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() * 3./2. * frac[0];
    Poiseuille2D<T> u(sGeometry, 3, maxVelocity, deltaR/2);

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

void getResults(Grid2D<T,DESCRIPTOR>& grid,
                const std::string& prefix,
                int iT)
{
  auto& converter = grid.getConverter();
  auto& sLattice  = grid.getSuperLattice();
  auto& sGeometry = grid.getSuperGeometry();

  SuperVTMwriter2D<T> vtmWriter(prefix);
  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();
  }

  vtmWriter.write(iT);
}

void takeMeasurements(Grid2D<T,DESCRIPTOR>& grid)
{
  static T maxDrag = 0.0;

  OstreamManager clout(std::cout,"measurement");

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

  SuperLatticePhysPressure2D<T,DESCRIPTOR> pressure(sLattice, converter);
  AnalyticalFfromSuperF2D<T> intpolatePressure(pressure, true);
  SuperLatticePhysDrag2D<T,DESCRIPTOR> dragF(sLattice, sGeometry, 5, converter);

  const T radiusCylinder  = cylinderD/2;

  const T point1[2] { cylinderX - radiusCylinder, cylinderY };
  const T point2[2] { cylinderX + radiusCylinder, cylinderY };

  T pressureInFrontOfCylinder, pressureBehindCylinder;
  intpolatePressure(&pressureInFrontOfCylinder, point1);
  intpolatePressure(&pressureBehindCylinder,    point2);

  T pressureDrop = pressureInFrontOfCylinder - pressureBehindCylinder;
  clout << "pressureDrop=" << pressureDrop;

  const int input[3] {};
  T drag[dragF.getTargetDim()] {};
  dragF(drag, input);
  if (drag[0] > maxDrag) {
    maxDrag = drag[0];
  };
  clout << "; drag=" << drag[0] << "; maxDrag: " << maxDrag << "; lift=" << drag[1] << endl;
}

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,
    RelaxationTime<T>(tau),
    N,
    PhysCharacteristics);
  const Vector<T,2> domainOrigin = coarseGrid.getSuperGeometry().getStatistics().getMinPhysR(0);
  const Vector<T,2> domainExtend = coarseGrid.getSuperGeometry().getStatistics().getPhysExtend(0);

  prepareGeometry(coarseGrid, domainOrigin, domainExtend);

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

  const Vector<T,2> fineOutflowExtend {1*cylinderD, domainExtend[1]};
  const Vector<T,2> fineOutflowOrigin {domainExtend[0]-1*cylinderD, 0};

  auto& fineOutflowGrid = coarseGrid.refine(fineOutflowOrigin, fineOutflowExtend, false);
  prepareGeometry(fineOutflowGrid, domainOrigin, domainExtend);

  {
    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().reset(refined);
  }

  const Vector<T,2> fineOutflowExtend2 {0.5*cylinderD, domainExtend[1]};
  const Vector<T,2> fineOutflowOrigin2 {domainExtend[0]-0.5*cylinderD, 0};

  auto& fineOutflowGrid2 = fineOutflowGrid.refine(fineOutflowOrigin2, fineOutflowExtend2, false);
  prepareGeometry(fineOutflowGrid2, domainOrigin, domainExtend);

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

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

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

    IndicatorCuboid2D<T> refined(extend, origin + Vector<T,2> {coarseDeltaX,0});
    fineOutflowGrid.getSuperGeometry().reset(refined);
  }

  coarseGrid.forEachGrid(prepareLattice);

  clout << "Total number of active cells: " << coarseGrid.getActiveVoxelN() << endl;
  clout << "Starting simulation..." << endl;

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

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

    coarseGrid.collideAndStream();

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

      coarseGrid.forEachGrid("cylinder2d", [&](Grid2D<T,DESCRIPTOR>& grid, const std::string& id) {
        getResults(grid, id, iT);
      });

      takeMeasurements(coarseGrid);
    }
  }

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