/*
 *  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

const T lx   = 4.0;           // length of the channel
const T ly   = 1.0;           // height of the channel
const int N  = 10;            // resolution of the model
const T Re   = 10.;           // Reynolds number
const T uMax = 0.01;          // Max lattice speed
const T maxPhysT = 60.;       // max. simulation time in s, SI unit
const T physInterval = 0.25;  // interval for the convergence check in s
const T residuum = 1e-5;      // residuum for the convergence check


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/2};
    const Vector<T,2> wallOrigin {-physSpacing/4, -physSpacing/4};

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

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

  // Set material number for inflow and outflow
  {
    const Vector<T,2> extend { physSpacing/2,  ly};
    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);
  }

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

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

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);

  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

  sBoundaryCondition.addVelocityBoundary(sGeometry, 2, omega); // 0-velocity walls
  sBoundaryCondition.addVelocityBoundary(sGeometry, 3, omega); // velocity inflow
  sBoundaryCondition.addPressureBoundary(sGeometry, 4, omega); // pressure outflow

  const T Lx = converter.getLatticeLength(lx);
  const T Ly = converter.getLatticeLength(ly);
  const T p0 = 8.*converter.getLatticeViscosity()*converter.getCharLatticeVelocity()*Lx/(Ly*Ly);
  AnalyticalLinear2D<T,T> rho(-p0/lx*DESCRIPTOR<T>::invCs2, 0, p0*DESCRIPTOR<T>::invCs2+1);

  const T maxVelocity = converter.getCharLatticeVelocity();
  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);

  auto materials = sGeometry.getMaterialIndicator({1,2,3,4});
  sLattice.defineRhoU(materials, rho, u);
  sLattice.iniEquilibrium(materials, rho, u);

  sLattice.initialize();

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

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

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

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

  const int statIter = converter.getLatticeTime(maxPhysT/10.);

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

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

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

    sLattice.getStatistics().print(iT,converter.getPhysTime(iT));
  }
}

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

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

  const T maxVelocity = converter.getCharPhysVelocity();
  const T radius = ly/2;
  std::vector<T> axisPoint{lx/2, ly/2};
  std::vector<T> axisDirection{1, 0};
  Poiseuille2D<T> uSol(axisPoint, axisDirection, maxVelocity, radius);

  const T Lx = converter.getLatticeLength(lx);
  const T Ly = converter.getLatticeLength(ly);
  T p0 = 8.*converter.getLatticeViscosity()*converter.getCharLatticeVelocity()*Lx/T(Ly*Ly);
  AnalyticalLinear2D<T,T> pSol(-converter.getPhysPressure(p0)/lx, 0, converter.getPhysPressure(p0));

  SuperLatticePhysVelocity2D<T,DESCRIPTOR> u(sLattice, converter);
  SuperLatticePhysPressure2D<T,DESCRIPTOR> p(sLattice, converter);
  auto fluid = sGeometry.getMaterialIndicator(1);

  int tmp[]= { };
  T result[2]= { };

  SuperRelativeErrorL2Norm2D<T> velocityError(u, uSol, fluid);
  velocityError(result, tmp);
  clout << "velocity-L2-error(rel)=" << result[0] << std::endl;

  SuperRelativeErrorL2Norm2D<T> pressureError(p, pSol, fluid);
  pressureError(result, tmp);
  clout << "pressure-L2-error(rel)=" << result[0] << std::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/2, ly};
  IndicatorCuboid2D<T> coarseCuboid(coarseExtend, coarseOrigin);

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

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

  const Vector<T,2> fineExtend {lx/2+coarseDeltaX, ly};
  const Vector<T,2> fineOrigin {lx/2-coarseDeltaX, 0.0};

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

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

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

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

  prepareGeometry(fineGrid);

  prepareLattice(coarseGrid);

  prepareLattice(fineGrid);

  clout << "Total number of active cells: " << coarseGrid.getActiveVoxelN() << endl;
  clout << "starting simulation..." << endl;
  Timer<T> timer(
    coarseGrid.getConverter().getLatticeTime(maxPhysT),
    coarseGrid.getSuperGeometry().getStatistics().getNvoxel());
  util::ValueTracer<T> converge(
    fineGrid.getConverter().getLatticeTime(physInterval),
    residuum);
  timer.start();

  for (int iT = 0; iT < coarseGrid.getConverter().getLatticeTime(maxPhysT); ++iT) {
    if (converge.hasConverged()) {
      clout << "Simulation converged." << endl;
      break;
    }

    coarseGrid.collideAndStream();

    getResults(
      "coarse_",
      coarseGrid,
      iT,
      timer,
      converge.hasConverged());
    getResults(
      "fine_",
      fineGrid,
      iT,
      timer,
      converge.hasConverged());

    converge.takeValue(fineGrid.getSuperLattice().getStatistics().getAverageEnergy(), true);
  }

  getError("coarse", coarseGrid);
  getError("fine", fineGrid);

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