/** @file */
#include <stdio.h>
#include <petscksp.h>
#include "gauss.h"
#include "rhs.h"
#include "global.h"
#include "eval.h"
#include "indices.h"

PetscErrorCode calculate_rhs(Vec rhs, PetscInt rhs_length, PetscInt elems_x,
                             PetscInt elems_y, PetscScalar* u1_vec,
                 PetscScalar* u2_vec, PetscScalar* p_vec,
                 PetscScalar k, PetscScalar nu,
                 PetscInt p_vec_length, PetscInt vel_y_vec_length)
{
  PetscErrorCode ierr;
  PetscInt elems[6];
  PetscScalar rhs_temp;
  /* --------------- X-COMONENT OF VELOCITY --------------- */
  for (PetscInt node_index=0; node_index<rhs_length-p_vec_length-vel_y_vec_length; node_index++)
  {
    /* initialize rhs entry and find all elems in supp(node_index) */
    rhs_temp = 0.0;
    get_ele_indices(node_index, elems_x, elems_y, elems, GRID_TYPE_FINE);
    /* start loop over all those elems and for each elem loop over all gauss points */
    for (PetscInt i=0; i<6; i++)
    {
      for (PetscInt g=0; g< GAUSS_5_NG ; g++)
      {
    if (elems[i] > -1)
    {
    rhs_temp += GAUSS_5_W[g]*evaluate_rhs1(node_index,
                                     elems[i],
                     GAUSS_5_X[g],
                     GAUSS_5_Y[g],
                     elems_x,
                     k,
                     nu,
                     u1_vec,
                     u2_vec,
                     p_vec);
    }
      }
    }
    ierr = VecSetValues(rhs, 1, &node_index, &rhs_temp, INSERT_VALUES);CHKERRQ(ierr);
    PetscPrintf(PETSC_COMM_WORLD, "\n\n");
    PetscPrintf(PETSC_COMM_WORLD, "index %d: %f\n", node_index, rhs_temp);
  }

  /* --------------- Y-COMONENT OF VELOCITY --------------- */
  for (PetscInt node_index=rhs_length-p_vec_length-vel_y_vec_length; node_index<rhs_length-p_vec_length; node_index++)
  {
    /* initialize rhs entry and find all elems in supp(node_index) */
    rhs_temp = 0.0;
    get_ele_indices(node_index-vel_y_vec_length, elems_x, elems_y, elems, GRID_TYPE_FINE);
    /* start loop over all those elems and for each elem loop over all gauss points */
    for (PetscInt i=0; i<6; i++)
    {
      for (PetscInt g=0; g< GAUSS_5_NG ; g++)
      {
    if (elems[i] > -1)
    {
    rhs_temp += GAUSS_5_W[g]*evaluate_rhs2(node_index-vel_y_vec_length,
                                     elems[i],
                     GAUSS_5_X[g],
                     GAUSS_5_Y[g],
                     elems_x,
                     k,
                 nu,
                     u1_vec,
                     u2_vec,
                     p_vec);
    }
      }
    }
    ierr = VecSetValues(rhs, 1, &node_index, &rhs_temp, INSERT_VALUES);CHKERRQ(ierr);
    PetscPrintf(PETSC_COMM_WORLD, "\nY-Component\n");
    PetscPrintf(PETSC_COMM_WORLD, "index %d: %f\n", node_index, rhs_temp);
  }

  /* --------------- PRESSURE --------------- */
  for (PetscInt node_index=rhs_length-p_vec_length; node_index<rhs_length; node_index++)
  {
    /* all entries for pressure-rhs are zero */
    rhs_temp = 0.0;
    ierr = VecSetValues(rhs, 1, &node_index, &rhs_temp, INSERT_VALUES);CHKERRQ(ierr);
  }
  ierr = VecAssemblyBegin(rhs);CHKERRQ(ierr);
  ierr = VecAssemblyEnd(rhs);CHKERRQ(ierr);
  return 0;
}

PetscScalar evaluate_rhs1(PetscInt glob_index,
                   PetscInt ele_index,
                   PetscScalar x1,
           PetscScalar x2,
           PetscInt elems_x,
           PetscScalar k,
           PetscScalar nu,
           PetscScalar* u1_vec,
           PetscScalar* u2_vec,
           PetscScalar* p_vec)
{
  PetscScalar u1=0, u2=0, d1_u1=0, d2_u1=0, lap_u1=0, d1_p=0;
  PetscScalar u_t1;
  PetscScalar phi=0;
  PetscScalar value;

  /* first evaluate some helper terms */
  u1 = eval_FE_fct(u1_vec, P2, ele_index, elems_x, x1, x2, GRID_TYPE_FINE);
  u2 = eval_FE_fct(u2_vec, P2, ele_index, elems_x, x1, x2, GRID_TYPE_FINE);
  d1_u1 = eval_FE_fct(u1_vec, D1_P2, ele_index, elems_x, x1, x2, GRID_TYPE_FINE);
  d2_u1 = eval_FE_fct(u1_vec, D2_P2, ele_index, elems_x, x1, x2, GRID_TYPE_FINE);
  lap_u1 = eval_FE_fct(u1_vec, LAP_P2, ele_index, elems_x, x1, x2, GRID_TYPE_FINE);
  d1_p = eval_FE_fct(p_vec, D1_P1, ele_index, elems_x, x1, x2, GRID_TYPE_COARSE);

  /* then calculate acceleration; */
  u_t1 = nu*lap_u1 - u1 * d1_u1 - u2 * d2_u1 - d1_p;

  /* evaluate the test function */
  phi = eval_glob_basis(P2, glob_index, ele_index, elems_x, x1, x2, GRID_TYPE_FINE);

  /* calculate return value */
  value = (2.0/k) * u1 * phi + u_t1 * phi;

  return value;
}


PetscScalar evaluate_rhs2(PetscInt glob_index,
                   PetscInt ele_index,
                   PetscScalar x1,
           PetscScalar x2,
           PetscInt elems_x,
           PetscScalar k,
           PetscScalar nu,
           PetscScalar* u1_vec,
           PetscScalar* u2_vec,
           PetscScalar* p_vec)
{
  PetscScalar u1=0, u2=0, d1_u2=0, d2_u2=0, lap_u2=0, d2_p=0;
  PetscScalar u_t2;
  PetscScalar phi=0;
  PetscScalar value;

  /* first evaluate some helper terms */
  u1 = eval_FE_fct(u1_vec, P2, ele_index, elems_x, x1, x2, GRID_TYPE_FINE);
  u2 = eval_FE_fct(u2_vec, P2, ele_index, elems_x, x1, x2, GRID_TYPE_FINE);
  d1_u2 = eval_FE_fct(u2_vec, D1_P2, ele_index, elems_x, x1, x2, GRID_TYPE_FINE);
  d2_u2 = eval_FE_fct(u2_vec, D2_P2, ele_index, elems_x, x1, x2, GRID_TYPE_FINE);
  lap_u2 = eval_FE_fct(u2_vec, LAP_P2, ele_index, elems_x, x1, x2, GRID_TYPE_FINE);
  d2_p = eval_FE_fct(p_vec, D2_P1, ele_index, elems_x, x1, x2, GRID_TYPE_COARSE);

  /* then calculate acceleration; */
  u_t2 = nu*lap_u2 - u1 * d1_u2 - u2 * d2_u2 - d2_p;

  /* evaluate the test function */
  phi = eval_glob_basis(P2, glob_index, ele_index, elems_x, x1, x2, GRID_TYPE_FINE);

  /* calculate return value */
  value = (2.0/k) * u2 * phi + u_t2 * phi;

  return value;
}