#include <iostream>
#include <ctime>
#include <cmath>

#include <dune/amdis/AMDiS.hpp>
#include <dune/amdis/ProblemStat.hpp>
#include <dune/amdis/Operators.hpp>

#ifdef DOW
#undef DOW
#endif
#define DOW AMDIS_DOW

using namespace AMDiS;

// 3 components: velocity with polynomial degree 2 and pressure with polynomial degree 1

using Grid = Dune::YaspGrid<AMDIS_DIM>;
using StokesParam   = TaylorHoodBasis<Grid::LeafGridView>;
using StokesProblem = ProblemStat<StokesParam>;

int main(int argc, char** argv)
{
    AMDiS::init(argc, argv);

    StokesProblem prob("stokes");
    prob.initialize(INIT_ALL);

    double viscosity = 1.0;
    Parameters::get("stokes->viscosity", viscosity);

    // tree-paths for components
    auto _v = 0_c;
    auto _p = 1_c;

    // <viscosity*grad(u_i), grad(v_i)>
    for (std::size_t i = 0; i < DOW; ++i) {
      auto opL = makeOperator(tag::gradtest_gradtrial{}, viscosity);
      prob.addMatrixOperator(opL, treepath(_v,i), treepath(_v,i));
    }

    // <d_i(v_i), p>
    auto opP = makeOperator(tag::divtestvec_trial{}, 1.0);
    prob.addMatrixOperator(opP, _v, _p);

    // <q, d_i(u_i)>
    auto opDiv = makeOperator(tag::test_divtrialvec{}, 1.0);
    prob.addMatrixOperator(opDiv, _p, _v);

    auto opZero = makeOperator(tag::test_trial{}, 0.0);
    prob.addMatrixOperator(opZero, _p, _p);


    // define boundary regions
    auto left = [](auto const& x) { return x[0] < 1.e-8; };
    auto not_left = [](auto const& x) { return x[0] > 1.0 - 1.e-8 || x[1] < 1.e-8 || x[1] > 1.0 - 1.e-8; };

    // define boundary values
    auto parabolic_y = [](auto const& x) -> Dune::FieldVector<double,DOW>
    {
      return {0.0, x[1]*(1.0 - x[1])};
    };

    auto zero = [](auto const& x) -> Dune::FieldVector<double,DOW>
    {
      return {0.0, 0.0};
    };

    // set boundary conditions for velocity
    prob.addDirichletBC(left, _v, _v, parabolic_y);
    prob.addDirichletBC(not_left, _v, _v, zero);

    prob.addDirichletBC([](auto const& x) { return x[0] < 1.e-8 && x[1] < 1.e-8; }, _p, _p, 0.0);

    AdaptInfo adaptInfo("adapt", prob.getNumComponents());

    // assemble and solve system
    prob.buildAfterCoarsen(adaptInfo, Flag(0));

    // write matrix to file
    if (Parameters::get<int>("stokesMesh->global refinements").value_or(0) < 2) {
      mtl::io::matrix_market_ostream out("stokes_matrix.mtx");
      out << prob.getSystemMatrix()->getMatrix();

      std::cout << "A = \n" << prob.getSystemMatrix()->getMatrix() << '\n';
      std::cout << "x = \n" << prob.getSolution()->getVector() << '\n';
      std::cout << "b = \n" << prob.getRhs()->getVector() << '\n';
    }

    prob.solve(adaptInfo);

    // output solution
    prob.writeFiles(adaptInfo);

    AMDiS::finalize();
    return 0;
}