#include "mesh2d.h"
#include "mesh2drectangle.h"
#include "fe_problems2d.h"
#include "fe_problems2dneohooke.h"

#include <petscsys.h>
#include <petscvec.h>
#include "petscksp.h"
#include "petsctao.h"

PetscErrorCode test_eq(Tao /*tao*/, Vec x, Vec ce, void *ctx);
PetscErrorCode test_jacobian_eq(Tao /*tao*/,Vec /*x*/, Mat /*J*/, Mat /*Jpre*/,void */*ctx*/);
PetscErrorCode ComputeInitialValue(FEProblem2D /* neohook_problem */, Vec X);
PetscErrorCode WrapperComputeObjectiveAndGradientRoutine(Tao, Vec X, PetscReal *objective, Vec Grad, void *ctx);
PetscErrorCode WrapperComputeHessianRoutine(Tao, Vec X, Mat Hessian, Mat MassMat, void *ctx);
PetscErrorCode rhsfunction(PetscReal [2]/* point[2] */, PetscReal rhs_value[2], void* ctx);
PetscErrorCode neumannfunction(PetscReal point[2], PetscReal neumann_value[2], void* ctx);


//--------------------------------------------------------------------------------
int main(int argc, char **args)
{

    PetscErrorCode    ierr;
    int               size;
    PetscInt          degree = 2;
    CoefficientType   coeff_type = NO_INCLUSION;
    SplittingType     splitting  = SYMMETRIC;
    DirichletBoundary dirichlet = LEFT_BOUNDARY;
    PetscBool         flag;
    Mesh2D            mesh;
    PetscInt          sizes[2] = {100, 10};
    PetscInt          sizes_subdomains[2] = {1, 1};
    PetscInt          subdomain_coo[2] = {0, 0};
    PetscInt          layers_for_incl[4] = {0, 0, 0, 0};
    PetscReal         lower_corner[2] = {0., 0.};
    PetscReal         upper_corner[2] = {1., 0.1};
    PetscInt          i, j;
    PetscInt          elementcounter;
    const PetscInt*   dirichlet_boundary_points;
    const PetscInt*   mesh_element_incl_flag;
    const PetscInt**  mesh_elements;
    const PetscReal** mesh_points;
    const PetscInt    dofs_per_node = 2;
    PetscInt          npoints, nelements, ndirichlet_points;
    PetscInt          size_element;
    FEProblem2D       neohooke_problem;
    Vec               X, Grad;
    Mat               Hessian, MassMatrix;
    PetscInt          nnz = 30;
    PetscReal         e_in = 210.0;
    PetscReal         nu_in = 0.3;
    PetscReal         e_ou = 210.0;
    PetscReal         nu_ou = 0.3;
    PetscReal         incomp_bound = 0.499;
    PetscReal         volumeforce[2] = {0., -2.};
    PetscReal         neumannforce[2] = { 0., 0.};
    Tao               tao;
    PetscBool         print_path = PETSC_FALSE;
    PetscInt          print_path_steps = 1;
    PetscInt          nloops, maxit_old;
    PetscInt          maxit_rest;
    TaoConvergedReason tao_reason;
    MPI_Comm         comm;

    ierr = PetscInitialize(&argc, &args, NULL, NULL);
    comm = PETSC_COMM_WORLD;
    ierr = MPI_Comm_size(comm, &size); CHKERRMPI(ierr);
    if(size > 1)
        SETERRQ(comm, PETSC_ERR_WRONG_MPI_SIZE, "Incorrect number of processors");

    /* Get options */
    ierr = PetscOptionsGetInt(NULL, NULL, "-nx", &sizes[0], NULL); CHKERRQ(ierr);
    ierr = PetscOptionsGetInt(NULL, NULL, "-ny", &sizes[1], NULL); CHKERRQ(ierr);
    ierr = PetscOptionsGetInt(NULL, NULL, "-degree", &degree, NULL); CHKERRQ(ierr);
    ierr = PetscOptionsHasName(NULL, NULL, "-left_boundary", &flag); CHKERRQ(ierr);
    if(flag)
        dirichlet = LEFT_BOUNDARY;
    ierr = PetscOptionsHasName(NULL, NULL, "-left_and_right_boundary", &flag); CHKERRQ(ierr);
    if(flag)
        dirichlet = LEFT_AND_RIGHT_BOUNDARY;

    ierr = PetscOptionsGetReal(NULL, NULL, "-lowerx", &lower_corner[0], NULL); CHKERRQ(ierr);
    ierr = PetscOptionsGetReal(NULL, NULL, "-lowery", &lower_corner[1], NULL); CHKERRQ(ierr);
    ierr = PetscOptionsGetReal(NULL, NULL, "-upperx", &upper_corner[0], NULL); CHKERRQ(ierr);
    ierr = PetscOptionsGetReal(NULL, NULL, "-uppery", &upper_corner[1], NULL); CHKERRQ(ierr);

    ierr = PetscOptionsGetReal(NULL, NULL, "-e_in", &e_in, NULL); CHKERRQ(ierr);
    ierr = PetscOptionsGetReal(NULL, NULL, "-nu_in", &nu_in, NULL); CHKERRQ(ierr);
    ierr = PetscOptionsGetReal(NULL, NULL, "-e_ou", &e_ou, NULL); CHKERRQ(ierr);
    ierr = PetscOptionsGetReal(NULL, NULL, "-nu_ou", &nu_ou, NULL); CHKERRQ(ierr);
    ierr = PetscOptionsGetReal(NULL, NULL, "-incomp_bound", &incomp_bound, NULL); CHKERRQ(ierr);
    ierr = PetscOptionsGetReal(NULL, NULL, "-volumeforce_x", &volumeforce[0], NULL); CHKERRQ(ierr);
    ierr = PetscOptionsGetReal(NULL, NULL, "-volumeforce_y", &volumeforce[1], NULL); CHKERRQ(ierr);

    ierr = PetscOptionsHasName(NULL, NULL, "-print_path", &print_path); CHKERRQ(ierr);
    ierr = PetscOptionsGetInt(NULL, NULL, "-print_path_steps", &print_path_steps, NULL); CHKERRQ(ierr);

    if(degree == 1)
        nnz = 30;
    else
        nnz = 50;
    ierr = PetscOptionsGetInt(NULL, NULL, "-nnz", &nnz, NULL); CHKERRQ(ierr);

    /* Create Mesh */
    if(degree == 1)
    {
        ierr = Mesh2DRectangleCreateP1(
            sizes,
            lower_corner,
            upper_corner,
            layers_for_incl,
            sizes_subdomains,
            subdomain_coo,
            coeff_type,
            splitting,
            dirichlet,
            &mesh); CHKERRQ(ierr);

        ierr = FEProblems2DNeoHookeP1Create(mesh,
                                            e_in,
                                            nu_in,
                                            e_ou,
                                            nu_ou,
                                            incomp_bound,
                                            rhsfunction,
                                            &volumeforce,
                                            neumannfunction,
                                            &neumannforce,
                                            &neohooke_problem); CHKERRQ(ierr);
    }else if(degree == 2)
    {
        ierr = Mesh2DRectangleCreateP2(
            sizes,
            lower_corner,
            upper_corner,
            layers_for_incl,
            sizes_subdomains,
            subdomain_coo,
            coeff_type,
            splitting,
            dirichlet,
            &mesh); CHKERRQ(ierr);

        ierr = FEProblems2DNeoHookeP2Create(mesh,
                                            e_in,
                                            nu_in,
                                            e_ou,
                                            nu_ou,
                                            incomp_bound,
                                            rhsfunction,
                                            &volumeforce,
                                            neumannfunction,
                                            &neumannforce,
                                            &neohooke_problem); CHKERRQ(ierr);
    }else
        SETERRQ(comm, PETSC_ERR_ARG_WRONG, "Incorrect polynom degree");

    ierr = Mesh2DGetSizes(mesh, &npoints, &nelements, &size_element); CHKERRQ(ierr);
    ierr = Mesh2DGetElements(mesh, &mesh_elements); CHKERRQ(ierr);
    ierr = Mesh2DGetPoints(mesh, &mesh_points); CHKERRQ(ierr);
    ierr = Mesh2DGetElementsPhysicalFlag(mesh, &mesh_element_incl_flag); CHKERRQ(ierr);
    ierr = Mesh2DGetDirichletBoundaryPoints(mesh, &ndirichlet_points, &dirichlet_boundary_points); CHKERRQ(ierr);

    /* Create Mat and vecs */
    ierr = MatCreateSeqAIJ(PETSC_COMM_SELF,
                           npoints * dofs_per_node,
                           npoints * dofs_per_node,
                           nnz,
                           PETSC_NULL,
                           &Hessian); CHKERRQ(ierr);
    ierr = MatCreateSeqAIJ(PETSC_COMM_SELF,
                           npoints * dofs_per_node,
                           npoints * dofs_per_node,
                           nnz,
                           PETSC_NULL,
                           &MassMatrix); CHKERRQ(ierr);
    ierr = MatCreateVecs(Hessian, &X, &Grad); CHKERRQ(ierr);

    /* Compute initial value */
    ierr = ComputeInitialValue(neohooke_problem, X); CHKERRQ(ierr);
    ierr = TaoCreate(comm, &tao); CHKERRQ(ierr);
    ierr = TaoSetFromOptions(tao); CHKERRQ(ierr);

    ierr = TaoSetObjectiveAndGradientRoutine(tao,
                                             WrapperComputeObjectiveAndGradientRoutine,
                                             neohooke_problem); CHKERRQ(ierr);
    ierr = TaoSetHessianRoutine(tao,
                                Hessian,
                                NULL, // MassMatrix,
                                WrapperComputeHessianRoutine,
                                neohooke_problem); CHKERRQ(ierr);
    ierr = TaoSetInitialVector(tao, X); CHKERRQ(ierr);
    ierr = TaoSetFromOptions(tao);

    {
        /* ALMM Test */
        PetscBool issame = PETSC_FALSE;
        Mat LMVM;
        KSP ksp;
        PC pc;
        Mat J;
        Vec eq;
        Tao subsolver;
        ierr = PetscObjectTypeCompare((PetscObject)tao, TAOALMM, &issame); CHKERRQ(ierr);
        PetscInt *idx;
        IS is;
        PetscInt counter = 0;
        const PetscInt *is_idx;
        PetscReal f;
        if(issame)
        {
            /* Create artificial Dirichlet boundary at the right side of the beam by equality constraints */
            ierr = PetscCalloc1(npoints, &idx); CHKERRQ(ierr);
            for(i = 0; i < npoints; ++i)
            {
                if(PetscAbs(mesh_points[i][0] - 1) < 1e-12)
                {
                    idx[counter++] = dofs_per_node * i;
                    idx[counter++] = dofs_per_node * i + 1;
                }
            }
            ierr = ISCreateGeneral(PETSC_COMM_WORLD, counter, idx, PETSC_COPY_VALUES, &is); CHKERRQ(ierr);
            ierr = PetscFree(idx); CHKERRQ(ierr);

            ierr = MatCreateSeqAIJ(PETSC_COMM_SELF,
                                   counter,
                                   npoints * dofs_per_node,
                                   1,
                                   PETSC_NULL,
                                   &J); CHKERRQ(ierr);
            ierr = MatSetUp(J); CHKERRQ(ierr);

            ierr = ISGetIndices(is, &is_idx); CHKERRQ(ierr);
            for(i = 0; i < counter; ++i)
            {
                ierr = MatSetValue(J, i, is_idx[i], 1., INSERT_VALUES); CHKERRQ(ierr);
            }
            ierr = ISRestoreIndices(is, &is_idx); CHKERRQ(ierr);

            ierr = MatAssemblyBegin(J, MAT_FINAL_ASSEMBLY); CHKERRQ(ierr);
            ierr = MatAssemblyEnd(J, MAT_FINAL_ASSEMBLY); CHKERRQ(ierr);

            ierr = MatCreateVecs(J, NULL, &eq); CHKERRQ(ierr);
            ierr = TaoSetEqualityConstraintsRoutine(tao, eq, test_eq, is); CHKERRQ(ierr);
            ierr = TaoSetJacobianEqualityRoutine(tao, J, J, test_jacobian_eq, NULL); CHKERRQ(ierr);
            ierr = VecDestroy(&eq); CHKERRQ(ierr);
            ierr = MatDestroy(&J); CHKERRQ(ierr);

            ierr = TaoALMMGetSubsolver(tao, &subsolver); CHKERRQ(ierr);

            ierr = TaoGetLMVMMatrix(subsolver, &LMVM); CHKERRQ(ierr);
            ierr = MatLMVMGetJ0KSP(LMVM, &ksp); CHKERRQ(ierr);
            ierr = KSPSetType(ksp, KSPPREONLY); CHKERRQ(ierr);
            ierr = KSPGetPC(ksp, &pc); CHKERRQ(ierr);
            ierr = PCSetType(pc, PCLU); CHKERRQ(ierr);
            ierr = PCFactorSetMatSolverType(pc, "mumps"); CHKERRQ(ierr);

            ierr = WrapperComputeHessianRoutine(NULL, X, Hessian, NULL, neohooke_problem); CHKERRQ(ierr);
            ierr = WrapperComputeObjectiveAndGradientRoutine(NULL, X, &f, Grad, neohooke_problem); CHKERRQ(ierr);
            ierr = MatLMVMSetJ0(LMVM, Hessian); CHKERRQ(ierr);
            // ierr = KSPView(ksp, PETSC_VIEWER_STDOUT_SELF); CHKERRQ(ierr);
        }
    }
    ierr = TaoSetUp(tao); CHKERRQ(ierr);


    if(0)
    {
        /* DEBUG */
        Mat LMVM;
        Tao subsolver;
        Vec direction, tmp;
        PetscReal f, steplength = 2, fstart;
        ierr = VecDuplicate(X, &direction); CHKERRQ(ierr);
        ierr = VecDuplicate(X, &tmp); CHKERRQ(ierr);

        //----------------------------------------
        /* Get subsolver */
        ierr = TaoALMMGetSubsolver(tao, &subsolver); CHKERRQ(ierr);
        //----------------------------------------
        /* compute the initial direction */
        ierr = TaoComputeObjectiveAndGradient(subsolver, X, &fstart, tmp);
        ierr = TaoGetLMVMMatrix(subsolver, &LMVM); CHKERRQ(ierr);
        ierr = MatSolve(LMVM, tmp, direction); CHKERRQ(ierr);
        ierr = VecScale(direction, -1.); CHKERRQ(ierr);
        //----------------------------------------
        /* check if descent direction */
        ierr = VecDot(tmp, direction, &f); CHKERRQ(ierr);
        ierr = PetscPrintf(PETSC_COMM_SELF, "vecdot = %16.16lf\n", f); CHKERRQ(ierr);

        //----------------------------------------
        /* Armijo like backtracking */
        ierr = PetscPrintf(PETSC_COMM_SELF, "Armijo-like backtracking for objective of augmented Lagrange function to show that a reduction is possible if the initial value is not overwritten:\n", f); CHKERRQ(ierr);
        for(PetscInt i = 0; i < 40; ++i)
        {
            steplength *= 0.5;
            /* Set initial value to Zero, since it is overwrite by TaoComputeObjective(subsolver, tmp, &f) */
            ierr = VecZeroEntries(X); CHKERRQ(ierr);

            /* compute new point */
            ierr = VecWAXPY(tmp, steplength, direction, X); CHKERRQ(ierr);
            /* evaluate energy at new point */
            ierr = TaoComputeObjective(subsolver, tmp, &f); CHKERRQ(ierr);
            ierr = PetscPrintf(PETSC_COMM_SELF, "fstart = %16.16lf, step =%16.16lf, f = %16.16lf\n", fstart, steplength, f); CHKERRQ(ierr);

        }

        ierr = VecDestroy(&direction); CHKERRQ(ierr);
        ierr = VecDestroy(&tmp); CHKERRQ(ierr);

        /* solve for the subsolver to show that no reduction is achieved in the ALMM objective is */
        ierr = PetscPrintf(PETSC_COMM_SELF, "TaoSolve for the subsolver to show that the initial value is overwritten within the backtracking of the line search:\n", f); CHKERRQ(ierr);
        ierr = TaoSolve(subsolver); CHKERRQ(ierr);
    }

    ierr = TaoSolve(tao); CHKERRQ(ierr);

    ierr = MatDestroy(&Hessian); CHKERRQ(ierr);
    ierr = MatDestroy(&MassMatrix); CHKERRQ(ierr);
    ierr = FEProblems2DNeoHookeDestroy(&neohooke_problem); CHKERRQ(ierr);
    ierr = Mesh2DDestroy(&mesh); CHKERRQ(ierr);

    PetscFinalize();

    return 0;
}

//--------------------------------------------------------------------------------
PetscErrorCode test_eq(Tao /*tao*/, Vec x, Vec ce, void *ctx)
{
    PetscErrorCode ierr;
    IS idx = (IS)ctx;
    PetscInt i, n;
    PetscScalar *vals;
    Vec tmp;

    PetscFunctionBegin;

    ierr = VecGetSubVector(x, idx, &tmp); CHKERRQ(ierr);
    ierr = VecCopy(tmp, ce); CHKERRQ(ierr);

    ierr = VecGetLocalSize(ce, &n); CHKERRQ(ierr);
    ierr = VecGetArray(ce, &vals); CHKERRQ(ierr);
    for(i = 0; i < n; ++i)
    {
        vals[i] -= 0;
    }
    ierr = VecRestoreArray(ce, &vals); CHKERRQ(ierr);

    PetscFunctionReturn(0);
}

//--------------------------------------------------------------------------------
PetscErrorCode test_jacobian_eq(Tao /*tao*/,Vec /*x*/, Mat /*J*/, Mat /*Jpre*/,void */*ctx*/)
{
    PetscFunctionBegin;
    /* Jacobian is constant */
    PetscFunctionReturn(0);
}

//--------------------------------------------------------------------------------
PetscErrorCode ComputeInitialValue(FEProblem2D /* neohook_problem */, Vec X)
{
    PetscErrorCode    ierr;

    PetscFunctionBegin;
    ierr = VecZeroEntries(X); CHKERRQ(ierr);
    PetscFunctionReturn(0);
}

//--------------------------------------------------------------------------------
PetscErrorCode WrapperComputeObjectiveAndGradientRoutine(Tao, Vec X, PetscReal *objective, Vec Grad, void *ctx)
{
    PetscErrorCode ierr;
    FEProblem2D    neohook_problem = (FEProblem2D)ctx;
    PetscInt       ndirichlet_points;
    PetscInt       i;
    PetscScalar    *values;
    const PetscInt *dirichlet_boundary_points;
    const PetscInt dofs_per_node = 2;
    const PetscReal zero = 0.;

    PetscFunctionBegin;
    ierr = FEProblems2DComputeObjectiveAndGradient(neohook_problem, X, objective, Grad); CHKERRQ(ierr);

    ierr = Mesh2DGetDirichletBoundaryPoints(neohook_problem->mesh,
                                            &ndirichlet_points,
                                            &dirichlet_boundary_points); CHKERRQ(ierr);

    ierr = VecGetArray(Grad, &values); CHKERRQ(ierr);
    for(i = 0; i < ndirichlet_points; ++i)
    {
        values[dofs_per_node * dirichlet_boundary_points[i]]     = zero;
        values[dofs_per_node * dirichlet_boundary_points[i] + 1] = zero;
    }
    ierr = VecRestoreArray(Grad, &values); CHKERRQ(ierr);

    PetscFunctionReturn(0);
}

//--------------------------------------------------------------------------------
PetscErrorCode WrapperComputeHessianRoutine(Tao, Vec X, Mat Hessian, Mat MassMat, void *ctx)
{
    PetscErrorCode ierr;
    FEProblem2D    neohook_problem = (FEProblem2D)ctx;
    PetscInt       ndirichlet_points, ndirichlet_dofs, i, counter;
    const PetscInt dofs_per_node = 2;
    const PetscInt *dirichlet_boundary_points;
    PetscInt       *dirichlet_boundary_dofs;

    PetscFunctionBegin;
    ierr = FEProblems2DComputeHessian(neohook_problem, X, Hessian, MassMat); CHKERRQ(ierr);
    ierr = Mesh2DGetDirichletBoundaryPoints(neohook_problem->mesh,
                                            &ndirichlet_points,
                                            &dirichlet_boundary_points); CHKERRQ(ierr);

    ndirichlet_dofs = ndirichlet_points * dofs_per_node;
    ierr = PetscCalloc1(ndirichlet_dofs, &dirichlet_boundary_dofs); CHKERRQ(ierr);
    counter = 0;
    for(i = 0; i < ndirichlet_points; ++i)
    {
        dirichlet_boundary_dofs[counter++] = dirichlet_boundary_points[i] * dofs_per_node;
        dirichlet_boundary_dofs[counter++] = dirichlet_boundary_points[i] * dofs_per_node + 1;
    }

    ierr = MatZeroRowsColumns(Hessian,
                              ndirichlet_dofs,
                              dirichlet_boundary_dofs,
                              1.0,
                              NULL,
                              NULL); CHKERRQ(ierr);

    ierr = PetscFree(dirichlet_boundary_dofs); CHKERRQ(ierr);

    PetscFunctionReturn(0);
}

//--------------------------------------------------------------------------------
PetscErrorCode rhsfunction(PetscReal [2]/* point[2] */, PetscReal rhs_value[2], void* ctx)
{
    PetscReal *volumeforce = (PetscReal *)ctx;

    PetscFunctionBegin;
    rhs_value[0] = volumeforce[0];
    rhs_value[1] = volumeforce[1];
    PetscFunctionReturn(0);
}

//--------------------------------------------------------------------------------
PetscErrorCode neumannfunction(PetscReal point[2], PetscReal neumann_value[2], void* ctx)
{
    PetscReal *neumannforce = (PetscReal *)ctx;

    PetscFunctionBegin;
    if(PetscAbs(point[0] - 1.) < 1e-12)
    {
        neumann_value[0] = neumannforce[0];
        neumann_value[1] = neumannforce[1];
    }else
    {
        neumann_value[0] = 0.;
        neumann_value[1] = 0.;
    }
    PetscFunctionReturn(0);
}