// Copyright (C) 2007 Murtazo Nazarov.
// Licensed under the GNU GPL Version 2.
// 
// Based on existing NSESolver of Johan Hoffman

#include <dolfin.h>
#include "../unicorn/CNSSolver.h"

using namespace dolfin::unicorn;
using namespace dolfin;

real margin = DOLFIN_EPS + 1e-3;

// Initial data in 2D
class InitialData_2D : public Function
{
public:
  InitialData_2D(Mesh& mesh) : Function(mesh) {}
  // Return the rank of the value space
  dolfin::uint rank() const
  {
    return 1;
  }

  // Return the dimension of the value space for axis i
  dolfin::uint dim(unsigned int i) const
  {
    return 4;
  }

  void eval(real* values, const real* x)  const
  {
    // rho
    if (x[0] < 0.5)
      values[0] = 1.0;
    else
      values[0] = 0.125;
	
    // m
    values[1] = 0.0;
    values[2] = 0.0;

    // e
    if (x[0] < 0.5)
      values[3] = 2.5;
    else
      values[3] = 0.25;
  }

};

// Initial data in 3D
class InitialData_3D : public Function
{
public:
  InitialData_3D(Mesh& mesh) : Function(mesh) {}

  // Return the rank of the value space
  dolfin::uint rank() const
  {
    return 1;
  }

  // Return the dimension of the value space for axis i
  dolfin::uint dim(unsigned int i) const
  {
    return 5;
  }

  void eval(real* values, const real* x) const
  {
    // rho
    if (x[0] < 0.5)
      values[0] = 1.0;
    else
      values[0] = 0.125;
	
    // m
    values[1] = 0.0;
    values[2] = 0.0;
    values[3] = 0.0;

    // e
    if (x[0] < 0.5)
      values[4] = 2.5;
    else
      values[4] = 0.25;
  }
};
// Force term
class ForceFunction : public Function
{
public:
  ForceFunction(Mesh& mesh) : Function(mesh) {}
  void eval(real* values, const real* x) const
  {
    int d = cell().dim();

    for(int i = 0; i < d; i++)
      {
	values[i] = 0.0;
      }
  }
};


// Sub domain for Dirichlet boundary condition        
class InOutflow : public SubDomain
{
public:
  bool inside(const real* x, bool on_boundary) const
  {
    return on_boundary && 
      (x[0] > 1.0 - margin) || (x[0] < margin);
  }
};

//--------------------------------------------
class TopBottom : public SubDomain
{
public:
  bool inside(const real* x, bool on_boundary) const
  {
    return on_boundary && 
      (x[1] > 1.0 - margin) || (x[1] < margin);
  }
};

//--------------------------------------------
class LeftRight : public SubDomain
{
public:
  bool inside(const real* x, bool on_boundary) const
  {
    return on_boundary && 
      (x[2] > 1.0 - margin) || (x[2] < margin);
  }
};

// Boundary condition for equation 
//--------------------------------------------
class BC_InOutflow : public Function
{
public:
  BC_InOutflow(Mesh& mesh) : Function(mesh) {}
  void eval(real* values, const real* x) const
  {
    values[0] = 0.0;
  }
};

//--------------------------------------------
class BC_TopBottom : public Function
{
public:
  BC_TopBottom(Mesh& mesh) : Function(mesh) {}
  void eval(real* values, const real* x) const
  {
    values[0] = 0.0;
  }
};

//--------------------------------------------
class BC_LeftRight : public Function
{
public:
  BC_LeftRight(Mesh& mesh) : Function(mesh) {}
  void eval(real* values, const real* x) const
  {
    values[0] = 0.0;
  }
};

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

  UnitSquare  mesh(40,40);   // 2D case
  //UnitCube  mesh(20,10,10);      // 3D case

  ForceFunction f(mesh);

  // Set initial data
  InitialData_2D w0(mesh);   // 2D case
  //InitialData_3D w0(mesh);     // 3D case

  // Define subdomains and boundary
  InOutflow inoutboun;
  TopBottom topbottomboun;
  LeftRight leftrightboun;

  BC_InOutflow inoutflow(mesh);
  BC_TopBottom topbottom(mesh);
  //  BC_LeftRight leftright(mesh);  // 3D case

  // Define subsystems:
  SubSystem density(0);
  SubSystem momentum(1);
  SubSystem energy(2);

  /// Create given sub sub system (two levels)                                
  SubSystem m0(1, 0);
  SubSystem m1(1, 1);
  //  SubSystem m2(1, 2);           // 3D case

  DirichletBC bc0(inoutflow, mesh, inoutboun, m0);
  DirichletBC bc1(topbottom, mesh, topbottomboun, m1);
  //  DirichletBC bc2(leftright, mesh, leftrightboun, m2);  // 3D case

  // Collect boundary conditions
  Array <BoundaryCondition*> bcs(&bc0, &bc1);      // 2D case
  //Array <BoundaryCondition*> bcs(&bc0, &bc1, &bc2);  // 3D case

  
  CNSSolver::solve(mesh, f, w0, bcs); 
  
  return 0;
}

