// Test program & code for my-mg multigrid code

#ifndef TEST_MYMG
#define TEST_MYMG

#include <iostream>
#include "my-mg.h"


namespace MY_MG {
  template<typename Types>
  class oneDPoisson : public multigrid<Types> {
  public:
    // Type definitions for local use
    typedef typename Types::value_type value_type;
    typedef typename Types::int_type int_type;
    typedef typename Types::Vector Vector;
    typedef typename Types::Matrix Matrix;
    typedef typename Types::Perm   Perm;
  private:
    value_type *op_scale;
  public:
    oneDPoisson(int_type n_levels, int_type _dims[], int_type _nu[]) : 
      multigrid<Types>(n_levels, _dims, _nu)
    {
      // need to set operator scale (op_scale)
      op_scale = new double[n_levels];
      op_scale[0] = 1;
      for ( int_type i = 1; i < n_levels; ++i )
        op_scale[i] = op_scale[i-1]/4;
    }


    // Need to over-ride the abstract virtual methods

    // apply_operator -- A[level] = op_scale[level]*tridiag(-1,2,-1)
    Vector &apply_operator(const Vector &v, int_type level, Vector &out)
    {
      value_type alpha = op_scale[level];
      int_type   n = get_dim(level);
      assert(v.size() == n && out.size() == n);
      out[0] = alpha*(2*v[0] - v[1]);
      for ( int_type i = 1; i < n-1; ++i )
        out[i] = alpha*(2*v[i] - v[i-1] - v[i+1]);
      out[n-1] = alpha*(2*v[n-1] - v[n-2]);
      return out;
    }

    // pre_relax -- use Gauss-Seidel (forward)
    Vector &pre_relax(Vector &v, const Vector &b, int_type level)
    {
      assert(v.size() == b.size() && b.size() == get_dim(level));
      return gauss_seidel(v,b,level);
    }

    // post_relax -- Gauss-Seidel (reverse order)
    Vector &post_relax(Vector &v, const Vector &b, int_type level)
    {
      assert(v.size() == b.size() && b.size() == get_dim(level));
      return gauss_seidel_rev(v,b,level);
    }

    // gauss_seidel -- Gauss-Seidel iteration 
    // for op_scale[level]*tridiag(-1,2,-1)*v = b
    Vector &gauss_seidel(Vector &v, const Vector &b, int_type level)
    {
      // Gauss-Seidel for op_scale[level]*tridiag(-1,2,-1)*v = b
      value_type scalar = 1.0/(2*op_scale[level]);
      int_type n = get_dim(level);
      v[0] = b[0]*scalar + v[1]/2;
      for ( int_type i = 1; i < n-1; ++i )
        v[i] = b[i]*scalar + (v[i-1] + v[i+1])/2;
      v[n-1] = b[n-1]*scalar + v[n-2]/2;
      return v;
    }

    // gauss_seidel_rev  -- Gauss-Seidel (reverse order)
    Vector &gauss_seidel_rev(Vector &v, const Vector &b, int_type level)
    {
      // Gauss-Seidel for op_scale[level]*tridiag(-1,2,-1)*v = b
      value_type scalar = 1.0/(2*op_scale[level]);
      int_type n = get_dim(level);
      v[n-1] = b[n-1]*scalar + v[n-2]/2;
      for ( int_type i = n-2; i > 0; --i )
        v[i] = b[i]*scalar + (v[i-1] + v[i+1])/2;
      v[0] = b[0]*scalar + v[1]/2;
      return v;
    }


    // interpolate -- interpolate v (at level+1) to out (at level)
    //   -- stencil is (1/2, 1, 1/2)
    Vector &interpolate(const Vector &v, int_type level, Vector &out)
    {
      // Note: Have to use this->get_max_level() as direct call
      // to get_max_level() causes error as get_max_level()
      // does not depend on any template parameter
      // Alternative: multigrid<Types>::get_max_level()
      assert(level < this->get_max_level() && level >= 0);
      assert(get_dim(level) == 2*get_dim(level+1)+1);
      int_type n = get_dim(level+1);
      out[0]=0.5*v[0];
      assert(out.size() == get_dim(level));
      for ( int_type j = 1; j < n; ++j ) {
        out[2*j-1] = v[j-1];
        out[2*j] = 0.5 * ( v[j-1]+v[j] );
      }
      out[2*n-1] = v[n-1];
      out[2*n] = 0.5*v[n-1];
#ifdef DEBUG
      std::cout << "interpolant of " << std::endl;
      print_vector(v);
      std::cout << "is" << std::endl;
      print_vector(out);
#endif
      return out;
    };

    // restriction -- restrict v (at level-1) to out (at level)
    //   -- stencil is (1/4, 1/2, 1/4)
    Vector &restriction(const Vector &v, int_type level, Vector &out)
    {
      assert(level <= this->get_max_level() && level >= 1);
      int_type n = get_dim(level);
      assert(2*n+1 == get_dim(level-1));
      assert(v.size() == get_dim(level-1) && out.size() == get_dim(level));
      for ( int_type i = 0; i < n; ++i )
        out[i] = 0.25*v[2*i] + 0.5*v[2*i+1] + 0.25*v[2*i+2];
#ifdef DEBUG
      std::cout << "restriction of " << std::endl;
      print_vector(v);
      std::cout << "is" << std::endl;
      print_vector(out);
#endif

      return out;
    }

    // solve -- use direct solver using MTL interface
    Vector &solve(Vector &x, const Vector &b, int_type level)
    {
      assert(x.size() == get_dim(level) && b.size() == get_dim(level));
      Matrix A(get_dim(level),get_dim(level));
      typename mtl::rows_type<Matrix>::type Ar = rows(A);
      Vector temp1(x.size()), temp2(x.size());

      // set up A matrix
      for ( int i = 0; i < x.size(); ++i )
        {
          // set up temp1 = e_i = i'th unit vector
          for ( int j = 0; j < b.size(); ++j )
            temp1[j] = 0;
          temp1[i] = 1;
          // temp2 = A*e_i = i'th column of A
          apply_operator(temp1, level, temp2);
          // set i'th column of A to temp2
          for ( int j = 0; j < b.size(); ++j )
            Ar[i][j] = temp2[j];
        }
      // print_all_matrix(A);
      // Solve A*x = b via lu & lu_solve a la MTL
      Perm pivot(A.ncols());
      lu_factor(A,pivot);
      lu_solve(A,pivot,b,x);
      return x;
    }

  }; // oneDPoisson


}

#endif // TEST_MYMG