#include <dolfin.h>
#include "bl.h"

using namespace dolfin;

class FreestreamBoundary : public SubDomain {
public:
    bool inside ( const Array<double> & x, bool on_boundary ) const override {
        return x[1] >= 1.0 - DOLFIN_EPS;
    }
};

class WallBoundary : public SubDomain {
    bool inside ( const Array<double> & x, bool on_boundary ) const override {
        return x[1] <= DOLFIN_EPS;
    }
};

class InflowBoundary : public SubDomain {
    bool inside ( const Array<double> & x, bool on_boundary ) const override {
        return x[0] <= DOLFIN_EPS;
    }
};

class PeriodicBoundary : public SubDomain {
    bool inside ( const Array<double> & x, bool on_boundary ) const override {
        return x[2] < DOLFIN_EPS;
    }

    void map(const Array<double>& x, Array<double>& y) const override {
        y[0] = x[0];
        y[1] = x[1];
        y[2] = x[2] - 1.0;
    }
};

class InputV :  public Expression {
public:
    InputV() : Expression(3) {}
    void eval(Array<double> & values, const Array<double> & x) const override {
        values[0] = 1.0;
        values[1] = 0;
        values[2] = 0;
    }
};

class ExtPressure : public Expression {
public:
    void eval(Array<double> & values, const Array<double> & x) const override {
        values[0] = 0;//x[0];                                        //NOTE: Must not use x[1]
    }
};

Mesh make_mesh(int cells_x,int cells_y,int cells_z,double k = 2.0 )
{
    UnitCubeMesh ans (cells_x, cells_y, cells_z);
    std::vector<double> & V = ans.geometry().x();

    for (size_t i = 1;i < V.size();i+=3)
        V[i] =  1.0 - std::log ( 1  + ( 1.0 - V[i] ) * ( std::exp ( k ) - 1.0 ) ) / k;
    return ans;
}

int main()
{
    Mesh mesh = make_mesh(8,16,8);
    PeriodicBoundary periodic_boundary;

    set_log_level(INFO);

    //Load function space and subspaces
    bl::FunctionSpace V ( mesh , periodic_boundary);

    SubSpace VUx (V,0);
    SubSpace VUy (V,1);
    SubSpace VUz (V,2);

    //Setup boundary conditions
    std::vector<const BoundaryCondition *> bcs;

    //Inflow condition
    InflowBoundary inflow_boundary;
    InputV u0;
    bcs.push_back( new DirichletBC( V, u0, inflow_boundary));

    //Freestream condition
    FreestreamBoundary freestream_boundary;
    Constant u_freestream(1.0);
    Constant w_freestream(0.0);

    bcs.push_back( new DirichletBC( VUx, u_freestream, freestream_boundary) );
    bcs.push_back( new DirichletBC( VUz, w_freestream, freestream_boundary) );

    //Wall condition
    WallBoundary wall_boundary;
    Constant u_noslip(0.0,0.0,0.0);
    bcs.push_back( new DirichletBC( V, u_noslip, wall_boundary )) ;

    //Setups forms
    bl::ResidualForm F ( V );
    bl::JacobianForm J ( V, V );

    //Pressure is a constant parameter
    ExtPressure p_freestream;
    F.p = p_freestream;

    //Setup u
    Function u(V);
    u = u0;
    F.u = u;
    J.u = u;

    NonlinearVariationalProblem bl_problem(F, u, bcs, J);
    NonlinearVariationalSolver solver(bl_problem);
    solver.solve();                                            //< crash

    //Setup solver parameters
    solve( F == 0,u, bcs,J);                                   //< fine

    //Save to file
    File pvd("u.pvd");
    pvd << u;

    //Plot
    plot ( u,"velocity" );
    plot ( p_freestream,mesh,  "pressure");

    interactive();

    return 0;
}