#include "Res_Jac.hpp"

PetscErrorCode FormFunction(SNES snes, Vec y, Vec f, void* ctx)
{
  UserPointers *user = (UserPointers *)ctx;

  VecZeroEntries(f);

//   VecCopy(user->sol_vp, y);

  Vec sol_v, sol_p;
  VecGetSubVector(y, user->is_v, &sol_v);
  VecGetSubVector(y, user->is_p, &sol_p);

  // Update velo
  user->dot_v_np1 -> PlusAX( sol_v, -1.0 );
  user->velo_np1 -> PlusAX( sol_v, -1.0 * user->gamma * user->dt );
  user->dot_v_alpha->PlusAX( sol_v, -1.0 * user->alpha_m );
  user->v_alpha->PlusAX( sol_v, -1.0 * user->alpha_f * user->gamma * user->dt );

  // Update pres
  user->dot_p_np1 -> PlusAX( sol_p, -1.0 );
  user->pres_np1 -> PlusAX( sol_p, -1.0 * user->gamma * user->dt );
  user->dot_p_alpha->PlusAX( sol_p, -1.0 * user->alpha_m );
  user->p_alpha->PlusAX( sol_p, -1.0 * user->alpha_f * user->gamma * user->dt );


  // Update disp
  user->dot_d_np1 -> PlusAX( sol_v, -1.0 * user->val_1 );
  user->disp_np1 -> PlusAX( sol_v, -1.0 * user->val_1 * user->gamma * user->dt );
  user->dot_d_alpha->PlusAX( sol_v, -1.0 * user->val_1 * user->alpha_m );
  user->d_alpha->PlusAX( sol_v, -1.0 * user->val_1 * user->alpha_f * user->gamma * user->dt );

  VecRestoreSubVector(y, user->is_v, &sol_v);
  VecRestoreSubVector(y, user->is_p, &sol_p);

  user->gassem_ptr->Clear_G();

  user->gassem_ptr->Assem_residual( user->curr_time, user->dt,
        user->dot_d_alpha, user->dot_p_alpha, user->dot_v_alpha,
        user->d_alpha, user->p_alpha, user->v_alpha,
        user->lassem_ptr, user->elem_v, user->elem_p, user->ien_v_ptr, user->ien_p_ptr, user->mSize,
        user->phy_bs, user->phy_bt, user->phy_bu, user->exa_phy, user->fnode_phy,
        user->geo_bs, user->geo_bt, user->geo_bu, user->exa_geo, user->fnode_geo,
        user->weight, user->nbc_part_v, user->nbc_part_p, user->ebc_part );

  VecCopy(user->gassem_ptr->G, f);

  VecAssemblyBegin(f);
  VecAssemblyEnd(f);

  double residual_norm = 0.0;
  VecNorm( f, NORM_2, &residual_norm );
  SYS_T::commPrint("  --- res norm: %e \n", residual_norm);
  VecNorm( user->gassem_ptr->G, NORM_2, &residual_norm );
  SYS_T::commPrint("  --- res norm: %e \n", residual_norm);
  

  return 0;

}

PetscErrorCode FormJacobian(SNES snes, Vec y, Mat jac, Mat B, void* ctx)
{
  UserPointers *user = (UserPointers *)ctx;

  MatZeroEntries(B);

  MatZeroEntries(user->gassem_ptr->K);

  Vec sol_v, sol_p;
  VecGetSubVector(y, user->is_v, &sol_v);
  VecGetSubVector(y, user->is_p, &sol_p);

  // Update velo
  user->dot_v_np1 -> PlusAX( sol_v, -1.0 );
  user->velo_np1 -> PlusAX( sol_v, -1.0 * user->gamma * user->dt );
  user->dot_v_alpha->PlusAX( sol_v, -1.0 * user->alpha_m );
  user->v_alpha->PlusAX( sol_v, -1.0 * user->alpha_f * user->gamma * user->dt );

  // Update pres
  user->dot_p_np1 -> PlusAX( sol_p, -1.0 );
  user->pres_np1 -> PlusAX( sol_p, -1.0 * user->gamma * user->dt );
  user->dot_p_alpha->PlusAX( sol_p, -1.0 * user->alpha_m );
  user->p_alpha->PlusAX( sol_p, -1.0 * user->alpha_f * user->gamma * user->dt );


  // Update disp
  user->dot_d_np1 -> PlusAX( sol_v, -1.0 * user->val_1 );
  user->disp_np1 -> PlusAX( sol_v, -1.0 * user->val_1 * user->gamma * user->dt );
  user->dot_d_alpha->PlusAX( sol_v, -1.0 * user->val_1 * user->alpha_m );
  user->d_alpha->PlusAX( sol_v, -1.0 * user->val_1 * user->alpha_f * user->gamma * user->dt );

  VecRestoreSubVector(y, user->is_v, &sol_v);
  VecRestoreSubVector(y, user->is_p, &sol_p);

  user->gassem_ptr->Assem_tangent( user->curr_time, user->dt,
        user->dot_d_alpha, user->dot_p_alpha, user->dot_v_alpha,
        user->d_alpha, user->p_alpha, user->v_alpha,
        user->lassem_ptr, user->elem_v, user->elem_p, user->ien_v_ptr, user->ien_p_ptr, user->mSize,
        user->phy_bs, user->phy_bt, user->phy_bu, user->exa_phy, user->fnode_phy,
        user->geo_bs, user->geo_bt, user->geo_bu, user->exa_geo, user->fnode_geo,
        user->weight, user->nbc_part_v, user->nbc_part_p, user->ebc_part );

  MatDuplicate(user->gassem_ptr->K, MAT_COPY_VALUES, &B);

  PetscReal norm;
  MatNorm(user->gassem_ptr->K, NORM_FROBENIUS, &norm);
  SYS_T::commPrint("  --- K norm: %e \n", norm);
  /*
      Assemble matrix
  */
  MatAssemblyBegin(B,MAT_FINAL_ASSEMBLY);
  MatAssemblyEnd(B,MAT_FINAL_ASSEMBLY);
//   MatView(B,PETSC_VIEWER_STDOUT_WORLD);
  if (jac != B) {
      MatAssemblyBegin(jac,MAT_FINAL_ASSEMBLY);
      MatAssemblyEnd(jac,MAT_FINAL_ASSEMBLY);
  }

  MatNorm(B, NORM_FROBENIUS, &norm);
  SYS_T::commPrint("  --- B norm: %e \n", norm);
  return 0;

}