#include <iostream>
#include <cstdlib>
#include <mtl/mtl.h>
#include <mtl/lu.h>
using namespace mtl;
// #define DEBUG
#include "op-mg.h"
#include "mtl-op.h"
#include "sparsemult.h"



// #define checkpt() (std::cout << "Checkpoint: file \"" << __FILE__        \
//                   << "\", line " << __LINE__ << std::endl)

// 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;
typedef myTypes::Matrix   Matrix;
typedef myTypes::Perm     Perm;

// Now that we have defined the Vector type, we can use this
#include "vec-util.h"

// SpMatrix is the sparse matrix type used for the operators
typedef matrix<value_type, rectangle<>, compressed<>, row_major>::type SpMatrix;

typedef MatrixOp<Vector, SpMatrix> MatOp;
typedef LUSolveOp<Vector, Matrix, Perm> LUSOp;
typedef GaussSeidelOp<Vector, SpMatrix> GSOp;
typedef GaussSeidelRevOp<Vector, SpMatrix> GSROp;
typedef Op<Vector> *PtrOp;
typedef SpMatrix   *PtrSp;


int my_idx(int i, int j, int n)
{
  if ( i < 0 || j < 0 || j >= 2*n-1 )
    return -1;
  if ( j < n )
    // 1st block
    return (i < n-1) ? (n-1)*j+i : -1;
  else
    // 2nd block
    return (i < 2*n-1) ? n*(n-1) + (2*n-1)*(j-n) + i : -1;
}

void rev_my_idx(int k, int *i, int *j, int n)
{
  if ( k < 0 )
    *i = *j = -1;
  else if ( k < n*(n-1) )
    {
      *j = k / (n-1);
      *i = k - (n-1)*(*j);
    }
  else if ( k < (n-1)*(3*n-1) )
    {
      *j = n+(k-n*(n-1))/(2*n-1);
      *i = k-n*(n-1)-(2*n-1)*(*j-n);
    }
  else
    *i = *j = -1;
}

void set_A0(SpMatrix *A0, int n)
{
  int dim = (n-1)*(3*n-1);
  for ( int k = 0; k < dim; ++k )
    {
      int i, j;
      rev_my_idx(k,&i,&j,n);
      int idx0 = my_idx(i,j,n), idx1;
      if ( (idx1 = my_idx(i,j-1,n)) >= 0 )
        (*A0)(idx0,idx1) = -1;
      if ( (idx1 = my_idx(i-1,j,n)) >= 0 )
        (*A0)(idx0,idx1) = -1;
      (*A0)(idx0,idx0) = 4;
      if ( (idx1 = my_idx(i+1,j,n)) >= 0 )
        (*A0)(idx0,idx1) = -1;
      if ( (idx1 = my_idx(i,j+1,n)) >= 0 )
        (*A0)(idx0,idx1) = -1;
    }
}

void set_Restrict(SpMatrix *Restrict, int n, int nc)
{
  int dim = (n-1)*(3*n-1);
  int dimc = (nc-1)*(3*nc-1);
  for ( int kc = 0; kc < dimc; ++kc )
    {
      // Get coarse grid node (i_c,j_c) for index kc
      int i_c, j_c;
      rev_my_idx(kc,&i_c,&j_c,nc);
      int idxc = my_idx(i_c,j_c,nc), idxf;
      // idxc == kc
      // Corresponding fine grid node is (i_f,j_f)
      int i_f = 2*i_c+1, j_f = 2*j_c+1;
      idxf = my_idx(i_f,j_f,n);
      (*Restrict)(idxc,idxf) = 1;
      if ( (idxf = my_idx(i_f-1,j_f  ,n)) >= 0 )
        (*Restrict)(idxc,idxf) = 0.5;
      if ( (idxf = my_idx(i_f+1,j_f  ,n)) >= 0 )
        (*Restrict)(idxc,idxf) = 0.5;
      if ( (idxf = my_idx(i_f  ,j_f-1,n)) >= 0 )
        (*Restrict)(idxc,idxf) = 0.5;
      if ( (idxf = my_idx(i_f  ,j_f+1,n)) >= 0 )
        (*Restrict)(idxc,idxf) = 0.5;
      if ( (idxf = my_idx(i_f-1,j_f-1,n)) >= 0 )
        (*Restrict)(idxc,idxf) = 0.25;
      if ( (idxf = my_idx(i_f+1,j_f-1,n)) >= 0 )
        (*Restrict)(idxc,idxf) = 0.25;
      if ( (idxf = my_idx(i_f-1,j_f+1,n)) >= 0 )
        (*Restrict)(idxc,idxf) = 0.25;
      if ( (idxf = my_idx(i_f+1,j_f+1,n)) >= 0 )
        (*Restrict)(idxc,idxf) = 0.25;
    }
}


int main(int argc, char *argv[])
{
  int_type n_levels;
  if ( argc > 1 )
    n_levels = atoi(argv[1]);
  else
    n_levels = 3;        // default value
  PtrSp *A       = new PtrSp[n_levels];
  PtrSp *Interp  = new PtrSp[n_levels];
  PtrSp *Restrict= new PtrSp[n_levels];

  PtrOp *Alist   = new PtrOp[n_levels];
  PtrOp *GSlist  = new PtrOp[n_levels];
  PtrOp *GSRlist = new PtrOp[n_levels];
  PtrOp *Ilist   = new PtrOp[n_levels];
  PtrOp *Rlist   = new PtrOp[n_levels];
  int_type dims[n_levels];

  std::cout << "Created arrays" << std::endl;
  // construct A[0] matrix
  std::cout << "About to construct A[0] matrix" << std::endl;
  int_type n = 1 << n_levels;
  int_type dim = (n-1)*(3*n-1);
  dims[0] = dim;
  std::cout << "n = " << n << ", dim = " << dim << std::endl;
  A[0] = new SpMatrix(dim,dim);

  set_A0(A[0], n);
  std::cout << "Constructed A[0] matrix" << std::endl;
  // print_all_matrix((*A[0]));
  // scale A[0]
  // scale(A[0], ...);
  // Turn into MatOp object, and create GSOp & GSROp for relaxation
  Alist[0]   = new MatOp(*A[0]);
  GSlist[0]  = new GSOp(*A[0]);
  GSRlist[0] = new GSROp(*A[0]);
  // GSlist[k] = new GSOp(A[k]);
  // GSRlist[k] = new GSROp(A[k]);

  for ( int level = 0; level < n_levels-1; ++level )
    {
      std::cout << "Setting up level " << level << std::endl;
      // Set up n and dim for the fine level ("level")
      n = 1 << (n_levels - level);
      dim = (n-1)*(3*n-1);
      dims[level] = dim;
      // Construct interpolation operator from level to level+1
      int_type nc = n/2;
      int_type dimc = (nc-1)*(3*nc-1);
      dims[level+1] = dimc;
      std::cout << "level = " << level << std::endl;
      std::cout << "n = " << n << ", dim = " << dim << std::endl;
      std::cout << "nc = " << nc << ", dimc = " << dimc << std::endl;
      Interp[level]   = new SpMatrix(dim,dimc);
      Restrict[level] = new SpMatrix(dimc,dim);

      std::cout << "Constructing Restrict[" << level << "]" << std::endl;
      set_Restrict(Restrict[level],n,nc);
      std::cout << "Constructed Restrict[" << level << "]" << std::endl;
      // print_all_matrix(*Restrict[level]);
      std::cout << "Transposing Restrict[" << level << "]" << std::endl;
      transpose(*Restrict[level],*Interp[level]);
      std::cout << "Transposed Restrict[" << level << "]" << std::endl;
      // print_all_matrix(*Interp[level]);

      A[level+1] = new SpMatrix(dimc,dimc);
      SpMatrix *temp = new SpMatrix(dim,dimc);
      std::cout << "Matrix multiplication for A[" << level+1 << "]" << std::endl;
      mult(*A[level],*Interp[level],*temp);
      mult(*Restrict[level],*temp,*A[level+1]);
      std::cout << "Matrix multiplication done" << std::endl;
      // std::cout << "A[" << level+1 << "] = " << std::endl;
      // print_all_matrix(*A[level+1]);
      // std::cout << "Restrict[" << level << "] = " << std::endl;
      // print_all_matrix(*Restrict[level]);
      // std::cout << "%% A[level]*Interp[level] =" << std::endl;
      // print_all_matrix(*temp);
      // std::cout << "%% A[level+1] = Restrict[level]*A[level]*Interp[level] =" << std::endl;
      // print_all_matrix(*A[level+1]);
    }

  std::cout << "Setting up the operators" << std::endl;

  // Now that all the sparse matrices are constructed, 
  // let's set up the operators (except for the ones from A[0])
  for ( int level = 1; level < n_levels; ++level )
    {
      Alist[level]   = new MatOp(*A[level]);
      GSlist[level]  = new GSOp(*A[level]);
      GSRlist[level] = new GSROp(*A[level]);
      Ilist[level-1] = new MatOp(*Interp[level-1]);
      Rlist[level-1] = new MatOp(*Restrict[level-1]);
    }

  std::cout << "Setting up top-level solver" << std::endl;
  // Finally the top level solver...
  checkpt();
  std::cout << "Dimension of top-level system is " << dims[n_levels-1] << std::endl;
  Matrix LU(dims[n_levels-1], dims[n_levels-1]);
  checkpt();
  copy(*A[n_levels-1], LU);
  checkpt();
  Perm pivot(dims[n_levels-1]);
  // LU factor top-level (coarsest) matrix
  checkpt();
  mtl::lu_factor(LU, pivot);
  checkpt();
  PtrOp top_solve = new LUSOp(LU, pivot);
  checkpt();

  std::cout << "Constructing mg" << std::endl;
  OpMG<myTypes> mg(n_levels, dims, Alist, GSlist, GSRlist, 
                   Ilist, Rlist, top_solve);
  std::cout << "Constructed mg" << std::endl;
  std::cout << "max level # = " << mg.get_max_level() << endl;

  Vector rhs(dims[0]), x(dims[0]), resid(dims[0]);
  for ( int i = 0; i < rhs.size(); ++i )
    rhs[i] = 1;
  // std::cout << "rhs = " << std::endl << rhs;

  for ( int iter = 0; iter < 20; ++iter )
    {
      mg.MGSolve( mg.get_max_level(), x, rhs);
      mg.residual(x, rhs, 0, resid);
      std::cout << "Residual norm = " << infinity_norm(resid) << std::endl;
    }

}