// THIS FILE CONTAINS CODE IN ORDER TO TREAT DIRICHLET CONDITIONS

// --------------------------------------------------
// PUT THE FOLLOWING HEADERS TO A HEADER FILE (.hh) 

/**
 * modify K and b such that solution of K u = b will have u[i] = 0 for Dirichlet nodes i
 * suppose K and b split into blocks
 * K = [ K_DD  K_DR ]     b = [ b_D ]   (D...Dirichlet nodes)
 *     [ K_RD  K_RR ]         [ b_R ]   (R...remaining nodes)
 * then we modify to
 * K = [ I     0    ]     b = [ 0   ]
 *     [ 0     K_RR ]         [ b_R ]
 *
 * @param  mesh  mesh
 * @param  K     Input:  fully assembled stiffness matrix
 *               Output: modified matrix of same size
 * @param  b     Input: fully assembled load vector (sources + Neumann B.C.)
 *               Output: modified vector of same size
 */
void incorporateHomogeneousDirichletBC (const Mesh& mesh, SparseMatrix& K, Vector& b);

/**
 * modify K and b such that solution of K u = b will have u[i] = ug[i] for Dirichlet nodes i
 * suppose K, b, and ug split into blocks
 * K = [ K_DD  K_DR ]     b = [ b_D ]    ug = [ g_D ]    (D...Dirichlet nodes)
 *     [ K_RD  K_RR ]         [ b_R ]         [ ... ]    (R...remaining nodes)
 * then we modify to
 * K = [ I     0    ]     b = [ g_D            ]
 *     [ 0     K_RR ]         [ b_R - K_RD g_D ]
 *
 * @param  mesh  mesh
 * @param  ug    Input: Vector of same size as b carrying prescribed Dirichlet values
 *               (other values are ignored)
 * @param  K     Input:  fully assembled stiffness matrix
 *               Output: modified matrix of same size
 * @param  b     Input: fully assembled load vector (sources + Neumann B.C.)
 *               Output: modified vector of same size
 */
void incorporateInhomogeneousDirichletBC (const Mesh& mesh, const Vector& ug, SparseMatrix& K, Vector& b);

// --------------------------------------------------------
// PUT THE FOLLOWING DEFINITIONS INTO A SOURCE FILE (.cc)

void incorporateHomogeneousDirichletBC (const Mesh& mesh, SparseMatrix& K, Vector& b)
{
  set<int> indices;   // set of column indices
  
  for (int i=0; i<mesh.nVerts; ++i)
    if (mesh.bcVerts[i] == BC_DIRICHLET)
      {
        // get column indices of row i
        K.getColIndices (i, indices);
        // set all entries of column i and row i to zero
        for (set<int>::const_iterator it=indices.begin(); it!=indices.end(); ++it)
          {
            K(i, *it) = 0;
            K(*it, i) = 0;
          }
        K(i, i) = 1;         // diagonal entry
        b[i] = 0;            // homogeneous B.C.
      }
} // incorporateHomogeneousDirichletBC


void incorporateInhomogeneousDirichletBC (const Mesh& mesh, const Vector& ug, SparseMatrix& K, Vector& b)
{
  set<int> indices;  // set of column indices
  Vector ug1(ug);    // copy
  Vector b1;         // helper
  
  // delete non-Dirichlet block from ug1
  for (int i=0; i<mesh.nVerts; ++i)
    {
      if (mesh.bcVerts[i] != BC_DIRICHLET) ug1[i] = 0;
    }
  b1 = K * ug1;
  
  // adapt load vector
  for (int i=0; i<mesh.nVerts; ++i)
    {
      if (mesh.bcVerts[i] != BC_DIRICHLET) b[i] -= b1[i];
      else b[i] = ug[i];
    }
  
  for (int i=0; i<mesh.nVerts; ++i)
    if (mesh.bcVerts[i] == BC_DIRICHLET)
      {
        // get column indices of row i
        K.getColIndices (i, indices);
        // set all entries of column i and row i to zero
        for (set<int>::const_iterator it=indices.begin(); it!=indices.end(); ++it)
          {
            K(i, *it) = 0;
            K(*it, i) = 0;
          }
        K(i, i) = 1;         // diagonal entry
        //b[i] = ug[i]       // already set
      }
} // incorporateInhomogeneousDirichletBC


//------------------------------------------------------
// EXAMPLE OF A MESH WITH DIRICHLET BOUNDARY (PART)
// PUT THIS INTO YOUR MAIN PROGRAM

void initMesh (Mesh& mesh)
{
  mesh.reserve (9, 8, 8);
  
  mesh.verts[0].set (0.0, 0.0);
  mesh.verts[1].set (0.5, 0.0);
  mesh.verts[2].set (1.0, 0.0);
  mesh.verts[3].set (0.0, 0.5);
  mesh.verts[4].set (0.5, 0.5);
  mesh.verts[5].set (1.0, 0.5);
  mesh.verts[6].set (0.0, 1.0);
  mesh.verts[7].set (0.5, 1.0);
  mesh.verts[8].set (1.0, 1.0);
  
  mesh.trigs[0].set (0, 4, 3);
  mesh.trigs[1].set (0, 1, 4);
  mesh.trigs[2].set (1, 5, 4);
  mesh.trigs[3].set (1, 2, 5);
  mesh.trigs[4].set (3, 7, 6);
  mesh.trigs[5].set (3, 4, 7);
  mesh.trigs[6].set (4, 8, 7);
  mesh.trigs[7].set (4, 5, 8);
  
  mesh.segments[0].set (0, 1); mesh.setSegmentBC (0, BC_DIRICHLET);
  // with the second statement, both
  // mesh.bcSegments and mesh.bcVerts are updated accordingly (see mesh.cc)

  mesh.segments[1].set (1, 2); mesh.setSegmentBC (1, BC_DIRICHLET);
  mesh.segments[2].set (2, 5); mesh.setSegmentBC (2, BC_NEUMANN);
  mesh.segments[3].set (5, 8); mesh.setSegmentBC (3, BC_NEUMANN);
  mesh.segments[4].set (8, 7); mesh.setSegmentBC (4, BC_NEUMANN);
  mesh.segments[5].set (7, 6); mesh.setSegmentBC (5, BC_NEUMANN);
  mesh.segments[6].set (6, 3); mesh.setSegmentBC (6, BC_NEUMANN);
  mesh.segments[7].set (3, 0); mesh.setSegmentBC (7, BC_NEUMANN);
}