/*
   An example of a program solving
     -div(k grad u) + cu = f
   in a 2D domain using linear finite elements.
*/
 
#include <iostream>
#include <fstream>
#include <iomanip>
#include "math.h"


//Include our own previously written code...
//  in general we will be using iterative solvers,
//  but for this particular example (since we're
//  assembling the full global matrix) we'll use
//  the Gaussian Elimination direct solver that
//  we have written for Math609.
#include "cg.cc"

//Include the header file defining the different
//  data structures that we use.
#include "2d_FE.h"

// Include a file that deals with getting the data 
//   for the mesh from a few files.
#include "readmesh.cc"

// Include a file with functions describing the linear
//   finite elements in 2D. 
#include "2d_linear.cc"

// Include a file with various quadratures. Currently only
//   a 3 point rule with nodes at the midpoints is implemented.
// This rule is exact for polynomials of degree 2 and thus
//   is good enough when working with linear elements.  
#include "quadrature_2d.cc"

/************************* The Coefficients ************************/

#define PI M_PI

double k(Point p)
{
   return 1;
}

double c(Point p)
{
   return 5;
}

double f(Point p)
{
   return  (5 + 18*PI*PI)*cos(3*PI*p.x)*cos(3*PI*p.y);
}


// an expected solution (used to check our approximation)
double sol(Point p)
{
   return cos(3*PI*p.x)*cos(3*PI*p.y);
}

/************************* The CG Code ************************/

// I've hardcoded the number 3 here. In general this should be a
//   class, and the dimension of the matrix and the vector should
//   be num_dofs = dofs_per_element();
typedef struct
{
      double LM[3][3]; //local matrix
      int    DM[3]; //DoF mapping
} LocalMatrix;

LocalMatrix *local_matrices;

//Calculate the action of our matrix A (for which
// we're solving Ax=b) on the given vector 'in' and
// saves the result in the given vector 'out'
void myMatrix(double* out, double* in, int n)
{
   for (int i=0; i<n; i++)
      out[i] = 0;

   //the contribution of the i-th local matrix
   for (int ele=0; ele<num_elements; ele++)
   {
      for(int i=0; i<dofs_per_element(); i++)
         for(int j=0; j<dofs_per_element(); j++)
            out[local_matrices[ele].DM[i]] +=
               local_matrices[ele].LM[i][j]*in[local_matrices[ele].DM[j]];
   }
}

/************************* The Main Code ************************/
int main()
{

   std::cout << "Reading the mesh..." << std::endl;
   if(!read_mesh("data.1"))
   {
      std::cout << "Error reading the mesh..." << std::endl;
      exit(1);
   }
   
   int global_dofs = num_nodes;

   local_matrices = new LocalMatrix[num_elements];
      
   double*  b = new double[global_dofs];

   Create_Quadratures();
   Quadrature *quad = &MIDPOINTS_RULE;

   //create the empty matrix and rhs
   for (int ele=0; ele<num_elements; ele++)
   {
      for(int i=0; i<9; i++)
         local_matrices[ele].LM[i/3][i%3] = 0;
   }
   for (int i=0; i<global_dofs; i++)
   {
      b[i] = 0;
   }

   int num_dofs = dofs_per_element();
   // The check for the local degrees of freedom.
   // If num_dofs is not 3, the definition of LocalMatrix should be
   //   modified accordingly.
   if (num_dofs != 3)
   {
      std::cout << "(num_dofs != 3). The program will not work correctly!" << std::endl;
   }

   std::cout << "Assembling the system..." << std::endl;

   // cycle over each of the elements
   //   and fill in the local contributions
   for (int ele=0; ele<num_elements; ele++)
   {
      fill_element_dofs(local_matrices[ele].DM, ele);

      // for all the points in our quadrature
      for (int l = 0; l < quad->size; l++)
      {
         // get the values of the point and the weight
         Point q_point = map_to_FE(quad->points[l], ele);
         double q_weight = quad->weights[l];
         for (int i=0; i<num_dofs; i++)
         {
            for (int j=0; j<num_dofs; j++)
            {
               // and calculate the value at a particualar point that is included in the
               //   quadrature used in computing the integral (k*phi_i',phi_j')+(c*phi_i,phi_j)
               local_matrices[ele].LM[i][j] += q_weight*fabs(jacobian(q_point,ele))*
                  (   k(q_point) * shape_funct_x(q_point, ele, i) * shape_funct_x(q_point, ele, j)
                      +
                      k(q_point) * shape_funct_y(q_point, ele, i) * shape_funct_y(q_point, ele, j)
                      +
                      c(q_point) * shape_funct(q_point, ele, i) * shape_funct(q_point, ele, j)
                  );
            }

            // ... and in a simillar manner - in computing the integral (f,phi_i)
            //  for the right hand side
            b[local_matrices[ele].DM[i]] += q_weight*fabs(jacobian(q_point, ele))*
                                     f(q_point) * shape_funct(q_point, ele, i);
         }
      }
   }

   std::cout << "Imposing the boundary conditions..." << std::endl;

   // impose the boundary conditions
   for (int k=0; k<global_dofs; k++)
   {
      // if this node is not at the boundary we don't care about it
      if (!nodes[k].bdmarker)
         continue;

      int count = 0;
      // if it is - we need to impose the Dirichlet BC
      b[k] = sol(nodes[k].p);

      for(int ele=0; ele<num_elements; ele++)
      {
         // obviously this search will be done in a more clever way
         //  in the general case ;-)
         int i=-1;
         if (local_matrices[ele].DM[0]==k)
            i = 0;
         else if (local_matrices[ele].DM[1]==k)
            i = 1;
         else if (local_matrices[ele].DM[2]==k)
            i = 2;

         if (i>=0)
         {
            for (int j=0; j<num_dofs; j++)
            {
               if (j!=i)
               {
                  local_matrices[ele].LM[i][j] = 0;
                  b[local_matrices[ele].DM[j]] -=
                     local_matrices[ele].LM[j][i]*b[k];
                  local_matrices[ele].LM[j][i] = 0;
               }
            }

            local_matrices[ele].LM[i][i] = 1;
            count++;
         }

      }

      b[k] *= count;
   }

   std::cout << "Solving the system..." << std::endl;
   // solve the system.
   solveCG(myMatrix,b,global_dofs);

   // Print out the numerical solution, the exact solution and the error
   //  at the nodes.
   double max = 0;
   for (int i=0; i<global_dofs; i++)
   {
      double error = fabs(b[i] - sol(nodes[i].p));
      std::cout << "| b[" << std::setw(3) << i << "] - u("<< std::setw(10) << nodes[i].p.x
                << ", " << std::setw(10) << nodes[i].p.y << ")|" << std::endl 
                << "                         = |"
                << std::setw(12) << b[i] << " - " << std::setw(12) << sol(nodes[i].p)
                << " | = " << std::setw(12) << error << std::endl;
      if (max < error)
         max = error;
   }
   std::cout << std::endl << "Discrete l-infiniy norm : " << max << std::endl;

   delete local_matrices;
   delete b;

   Destroy_Quadratures();
}