#include <iostream>
// #include <vector>
#include <cstdlib>
#include <mtl/mtl.h>
#include <mtl/lu.h>
#include "my-mg.h"
#include "my-vector.h"
#include "test-mymg.h"

using namespace mtl;


// Derived class from multigrid which defines the operators needed
typedef double value_type;
typedef MY_MG::TypeBinder<value_type, 
                          dense1D<value_type>,
                          matrix<value_type>::type, 
                          dense1D<int> > myTypes;
typedef myTypes::int_type int_type;
typedef myTypes::Vector   Vector;


ostream &operator<<(ostream &s, Vector &x)
{
  s << "Vector: dim: " << x.size() << std::endl;
  for ( int_type i = 0; i < x.size(); ++i )
    {
      s << x[i] << "  ";
      if ( i % 6 == 5 )
        s << std::endl;
    }
  s << std::endl;
  return s;
}

typedef MY_MG::oneDPoisson<myTypes> oneDPoissonMG;

int test1()
{
  int_type n_levels = 5;
  // int_type dims[] = { 1023, 511, 255, 127, 63, 31, 15, 7, 3, 1 };
  int_type dims[] = { 31, 15, 7, 3, 1 };
  int_type nu[] = { 1, 1 };
  oneDPoissonMG mg(n_levels, dims, nu);

  Vector b(dims[0]), x(dims[0]);

  for ( int_type i = 0; i < dims[0]; ++i )
    {
      b[i] = rand() / value_type(RAND_MAX);
      x[i] = 0;
    }
  
  std::cout << b;

  mg.solve(x,b,0);

  std::cout << x;

  Vector r(x.size());
  mg.apply_operator(x,0,r); 
  scale(r,value_type(-1));
  add(r,b,r);
  std::cout << "Residual is " << infinity_norm(r) << std::endl;

  for ( int_type i = 0; i < x.size(); ++i )
    x[i] = 0;

  std::cout << "Max level: " << mg.get_max_level() << std::endl;
  mg.MGSolve(mg.get_max_level(), x, b);
  std::cout << "Solution from one step of multigrid:" << std::endl;
  std::cout << x;
  mg.apply_operator(x,0,r); 
  scale(r,value_type(-1));
  add(r,b,r);
  std::cout << "Residual is " << infinity_norm(r) << std::endl;
  std::cout << std::endl << "Repeating..." << std::endl;
  mg.MGSolve(mg.get_max_level(), x, b);
  std::cout << "Solution from two steps of multigrid:" << std::endl;
  std::cout << x;
  mg.apply_operator(x,0,r); 
  scale(r,value_type(-1));
  add(r,b,r);
  std::cout << "Residual is " << infinity_norm(r) << std::endl;
  std::cout << std::endl << "Repeating..." << std::endl;
  mg.MGSolve(mg.get_max_level(), x, b);
  std::cout << "Solution from three steps of multigrid:" << std::endl;
  std::cout << x;
  mg.apply_operator(x,0,r); 
  scale(r,value_type(-1));
  add(r,b,r);
  std::cout << "Residual is " << infinity_norm(r) << std::endl;
  std::cout << std::endl << "Repeating..." << std::endl;
  mg.MGSolve(mg.get_max_level(), x, b);
  std::cout << "Solution from three steps of multigrid:" << std::endl;
  std::cout << x;
  mg.apply_operator(x,0,r); 
  scale(r,value_type(-1));
  add(r,b,r);
  std::cout << "Residual is " << infinity_norm(r) << std::endl;

}

// Performance/convergence tests
int test2(int_type n_levels)
{
  // int_type dims[] = { 1023, 511, 255, 127, 63, 31, 15, 7, 3, 1 };
  int_type *dims = new int_type[n_levels];
  for ( int_type level = 0; level < n_levels; ++level )
    dims[level] = (1 << (n_levels-level)) - 1;
  int_type nu[] = { 1, 1 };
  oneDPoissonMG mg(n_levels, dims, nu);

  std::cout << "Number of levels: " << n_levels << std::endl;

  Vector b(dims[0]), x(dims[0]);

  for ( int_type i = 0; i < dims[0]; ++i )
    {
      b[i] = rand() / value_type(RAND_MAX);
      x[i] = 0;
    }
  
  // std::cout << b;

  // mg.solve(x,b,0);

  // std::cout << x;

  Vector r(x.size());

  // print out original residual norm (l^\infty)
  mg.apply_operator(x,0,r); 
  scale(r,value_type(-1));
  add(r,b,r);
  std::cout << "Residual is " << infinity_norm(r) << std::endl;

  int max_iter = 20;
  for ( int iter = 0; iter < max_iter; ++iter )
    {
      mg.MGSolve(mg.get_max_level(), x, b);
      mg.apply_operator(x,0,r); 
      scale(r,value_type(-1));
      add(r,b,r);
      std::cout << "||r|| = " << infinity_norm(r) << std::endl;
    }
  std::cout << "||r||/(||A||.||x||) = " 
            << infinity_norm(r)/(4*infinity_norm(x)) << std::endl;


}

int main(int argc, char *argv[])
{
  if ( argc == 1 )
    test2(5);
  else
    {
      std::cout << "Using " << argv[1] << " levels" << std::endl;
      test2(atoi(argv[1]));
    }
}