// SPDX-FileCopyrightText: Copyright (c) 2017-2024, HONEE contributors.
// SPDX-License-Identifier: Apache-2.0 OR BSD-2-Clause
/// @file
/// Functions for setting up and projecting the velocity gradient

#include "../qfunctions/velocity_gradient_projection.h"

#include <petscdmplex.h>

#include <navierstokes.h>

static PetscErrorCode VelocityGradientProjectionCreateDM(NodalProjectionData grad_velo_proj, Honee honee, PetscInt degree) {
  PetscSection section;

  PetscFunctionBeginUser;
  grad_velo_proj->num_comp = 9;  // 9 velocity gradient

  PetscCall(DMClone(honee->dm, &grad_velo_proj->dm));
  PetscCall(DMSetMatrixPreallocateSkip(grad_velo_proj->dm, PETSC_TRUE));
  PetscCall(PetscObjectSetName((PetscObject)grad_velo_proj->dm, "Velocity Gradient Projection"));

  PetscCall(DMSetupByOrder_FEM(PETSC_TRUE, PETSC_TRUE, honee->app_ctx->degree, 1, honee->app_ctx->q_extra, 1, &grad_velo_proj->num_comp,
                               grad_velo_proj->dm));

  PetscCall(DMGetLocalSection(grad_velo_proj->dm, &section));
  PetscCall(PetscSectionSetFieldName(section, 0, ""));
  PetscCall(PetscSectionSetComponentName(section, 0, 0, "VelocityGradientXX"));
  PetscCall(PetscSectionSetComponentName(section, 0, 1, "VelocityGradientXY"));
  PetscCall(PetscSectionSetComponentName(section, 0, 2, "VelocityGradientXZ"));
  PetscCall(PetscSectionSetComponentName(section, 0, 3, "VelocityGradientYX"));
  PetscCall(PetscSectionSetComponentName(section, 0, 4, "VelocityGradientYY"));
  PetscCall(PetscSectionSetComponentName(section, 0, 5, "VelocityGradientYZ"));
  PetscCall(PetscSectionSetComponentName(section, 0, 6, "VelocityGradientZX"));
  PetscCall(PetscSectionSetComponentName(section, 0, 7, "VelocityGradientZY"));
  PetscCall(PetscSectionSetComponentName(section, 0, 8, "VelocityGradientZZ"));
  PetscFunctionReturn(PETSC_SUCCESS);
};

PetscErrorCode VelocityGradientProjectionSetup(Ceed ceed, Honee honee, ProblemData problem, StateVariable state_var_input,
                                               CeedElemRestriction elem_restr_input, CeedBasis basis_input, NodalProjectionData *pgrad_velo_proj) {
  NodalProjectionData grad_velo_proj;
  CeedBasis           basis_grad_velo;
  CeedElemRestriction elem_restr_grad_velo, elem_restr_qd;
  CeedVector          q_data;
  PetscInt            dim, height = 0, dm_field = 0, num_comp_input;
  CeedInt             q_data_size;

  PetscFunctionBeginUser;
  PetscCall(PetscNew(&grad_velo_proj));
  PetscCall(VelocityGradientProjectionCreateDM(grad_velo_proj, honee, honee->app_ctx->degree));

  // -- Get Pre-requisite things
  PetscCall(DMGetDimension(grad_velo_proj->dm, &dim));
  PetscCall(DMGetFieldNumComps(honee->dm, 0, &num_comp_input));
  PetscCall(DMPlexCeedElemRestrictionCreate(ceed, grad_velo_proj->dm, DMLABEL_DEFAULT, DMLABEL_DEFAULT_VALUE, height, dm_field,
                                            &elem_restr_grad_velo));
  PetscCall(DMPlexCeedBasisCreate(ceed, grad_velo_proj->dm, DMLABEL_DEFAULT, DMLABEL_DEFAULT_VALUE, height, dm_field, &basis_grad_velo));
  PetscCall(QDataGet(ceed, grad_velo_proj->dm, DMLABEL_DEFAULT, DMLABEL_DEFAULT_VALUE, &elem_restr_qd, &q_data, &q_data_size));

  {  // -- Build RHS operator
    CeedOperator        op_rhs_assemble;
    CeedQFunction       qf_rhs_assemble;
    CeedElemRestriction elem_restr_x;
    CeedBasis           basis_x;
    CeedVector          x_coord;
    PetscInt            num_comp_x;

    PetscCall(DMGetCoordinateNumComps(honee->dm, &num_comp_x));
    PetscCall(DMPlexCeedCoordinateCreateField(ceed, honee->dm, DMLABEL_DEFAULT, DMLABEL_DEFAULT_VALUE, height, &elem_restr_x, &basis_x, &x_coord));
    switch (state_var_input) {
      case STATEVAR_PRIMITIVE:
        PetscCallCeed(ceed, CeedQFunctionCreateInterior(ceed, 1, VelocityGradientProjectionRHS_Prim, VelocityGradientProjectionRHS_Prim_loc,
                                                        &qf_rhs_assemble));
        break;
      case STATEVAR_CONSERVATIVE:
        PetscCallCeed(ceed, CeedQFunctionCreateInterior(ceed, 1, VelocityGradientProjectionRHS_Conserv, VelocityGradientProjectionRHS_Conserv_loc,
                                                        &qf_rhs_assemble));
        break;
      case STATEVAR_ENTROPY:
        PetscCallCeed(ceed, CeedQFunctionCreateInterior(ceed, 1, VelocityGradientProjectionRHS_Entropy, VelocityGradientProjectionRHS_Entropy_loc,
                                                        &qf_rhs_assemble));
        break;
    }

    PetscCallCeed(ceed, CeedQFunctionSetContext(qf_rhs_assemble, problem->apply_vol_ifunction.qfctx));
    PetscCallCeed(ceed, CeedQFunctionAddInput(qf_rhs_assemble, "q", num_comp_input, CEED_EVAL_INTERP));
    PetscCallCeed(ceed, CeedQFunctionAddInput(qf_rhs_assemble, "Grad_q", num_comp_input * dim, CEED_EVAL_GRAD));
    PetscCallCeed(ceed, CeedQFunctionAddInput(qf_rhs_assemble, "qdata", q_data_size, CEED_EVAL_NONE));
    PetscCallCeed(ceed, CeedQFunctionAddInput(qf_rhs_assemble, "x", num_comp_x, CEED_EVAL_INTERP));
    PetscCallCeed(ceed, CeedQFunctionAddOutput(qf_rhs_assemble, "velocity gradient", grad_velo_proj->num_comp, CEED_EVAL_INTERP));

    PetscCallCeed(ceed, CeedOperatorCreate(ceed, qf_rhs_assemble, NULL, NULL, &op_rhs_assemble));
    PetscCallCeed(ceed, CeedOperatorSetField(op_rhs_assemble, "q", elem_restr_input, basis_input, CEED_VECTOR_ACTIVE));
    PetscCallCeed(ceed, CeedOperatorSetField(op_rhs_assemble, "Grad_q", elem_restr_input, basis_input, CEED_VECTOR_ACTIVE));
    PetscCallCeed(ceed, CeedOperatorSetField(op_rhs_assemble, "qdata", elem_restr_qd, CEED_BASIS_NONE, q_data));
    PetscCallCeed(ceed, CeedOperatorSetField(op_rhs_assemble, "x", elem_restr_x, basis_x, x_coord));
    PetscCallCeed(ceed, CeedOperatorSetField(op_rhs_assemble, "velocity gradient", elem_restr_grad_velo, basis_grad_velo, CEED_VECTOR_ACTIVE));

    PetscCall(OperatorApplyContextCreate(honee->dm, grad_velo_proj->dm, ceed, op_rhs_assemble, NULL, NULL, NULL, NULL, &grad_velo_proj->l2_rhs_ctx));

    PetscCallCeed(ceed, CeedElemRestrictionDestroy(&elem_restr_x));
    PetscCallCeed(ceed, CeedBasisDestroy(&basis_x));
    PetscCallCeed(ceed, CeedVectorDestroy(&x_coord));
    PetscCallCeed(ceed, CeedOperatorDestroy(&op_rhs_assemble));
    PetscCallCeed(ceed, CeedQFunctionDestroy(&qf_rhs_assemble));
  }

  {  // -- Build Mass operator
    CeedOperator  op_mass;
    CeedQFunction qf_mass;
    Mat           mat_mass;

    PetscCall(HoneeMassQFunctionCreate(ceed, grad_velo_proj->num_comp, q_data_size, &qf_mass));
    PetscCallCeed(ceed, CeedOperatorCreate(ceed, qf_mass, NULL, NULL, &op_mass));
    PetscCallCeed(ceed, CeedOperatorSetField(op_mass, "u", elem_restr_grad_velo, basis_grad_velo, CEED_VECTOR_ACTIVE));
    PetscCallCeed(ceed, CeedOperatorSetField(op_mass, "qdata", elem_restr_qd, CEED_BASIS_NONE, q_data));
    PetscCallCeed(ceed, CeedOperatorSetField(op_mass, "v", elem_restr_grad_velo, basis_grad_velo, CEED_VECTOR_ACTIVE));

    PetscCall(MatCreateCeed(grad_velo_proj->dm, grad_velo_proj->dm, op_mass, NULL, &mat_mass));

    // Setup KSP for L^2 projection with lumped mass operator
    PetscCall(KSPCreate(PetscObjectComm((PetscObject)grad_velo_proj->dm), &grad_velo_proj->ksp));
    PetscCall(KSPSetOptionsPrefix(grad_velo_proj->ksp, "velocity_gradient_projection_"));
    {
      PC pc;
      PetscCall(KSPGetPC(grad_velo_proj->ksp, &pc));
      PetscCall(PCSetType(pc, PCJACOBI));
      PetscCall(PCJacobiSetType(pc, PC_JACOBI_ROWSUM));
      PetscCall(KSPSetType(grad_velo_proj->ksp, KSPPREONLY));
    }
    PetscCall(KSPSetFromOptions_WithMatCeed(grad_velo_proj->ksp, mat_mass));

    PetscCallCeed(ceed, CeedQFunctionDestroy(&qf_mass));
    PetscCallCeed(ceed, CeedOperatorDestroy(&op_mass));
    PetscCall(MatDestroy(&mat_mass));
  }

  *pgrad_velo_proj = grad_velo_proj;

  PetscCallCeed(ceed, CeedBasisDestroy(&basis_grad_velo));
  PetscCallCeed(ceed, CeedVectorDestroy(&q_data));
  PetscCallCeed(ceed, CeedElemRestrictionDestroy(&elem_restr_qd));
  PetscCallCeed(ceed, CeedElemRestrictionDestroy(&elem_restr_grad_velo));
  PetscFunctionReturn(PETSC_SUCCESS);
}

PetscErrorCode VelocityGradientProjectionApply(NodalProjectionData grad_velo_proj, Vec Q_loc, Vec VelocityGradient) {
  OperatorApplyContext l2_rhs_ctx = grad_velo_proj->l2_rhs_ctx;

  PetscFunctionBeginUser;
  PetscCall(PetscLogEventBegin(HONEE_VelocityGradientProjection, Q_loc, VelocityGradient, 0, 0));
  PetscCall(ApplyCeedOperatorLocalToGlobal(Q_loc, VelocityGradient, l2_rhs_ctx));

  PetscCall(KSPSolve(grad_velo_proj->ksp, VelocityGradient, VelocityGradient));
  PetscCall(PetscLogEventEnd(HONEE_VelocityGradientProjection, Q_loc, VelocityGradient, 0, 0));
  PetscFunctionReturn(PETSC_SUCCESS);
}