#include <iostream>

#include "matrix.h"
#include "utils.h"

using namespace mtl;

typedef matrix<double, rectangle<>, compressed<>, row_major >::type SpMat;
typedef compressed1D<double> SparseVec;

// mult -- sparse-sparse matrix multiply
//      -- sets out <- out + A*B
//      -- assumes matrices row-major for efficient implementation
void mult(const SpMat &A, const SpMat &B, SpMat &out)
{
  rows_type<SpMat>::type Ar, Br;
  Ar = rows(A);
  Br = rows(B);

  // row for output -- need to re-use
  SparseVec temp_row;

  temp_row.reserve(10);

  SpMat::iterator i;
  for ( i = Ar.begin(); i != Ar.end(); ++i )
    {
      int row_num = (*i).begin().row();
      temp_row.clear();
      SpMat::OneD::iterator j;
      for ( j = (*i).begin(); j != (*i).end(); ++j )
        {
          int col_num = j.column();
          double val = *j;
          SpMat::Row Brow = Br[col_num];
          SpMat::Row::iterator k;
          for ( k = Brow.begin(); k != Brow.end(); ++k )
            temp_row[k.column()] = temp_row[k.column()] + val*(*k);
        }
      SparseVec::iterator p;
      for ( p = temp_row.begin(); p != temp_row.end(); ++p )
        out(row_num,p.index()) += *p;
    }
}

 
int main(int argc, char *argv[])
{
  SpMat A(3,3), B(3,3), C(3,3);

  A(0,0) = 1;  A(0,2) = 2;
  A(1,0) = 3;  A(1,2) = 1;
  A(2,1) = 4;  A(2,2) = 5;

  B(0,1) = 1;  B(0,2) = 2;
  B(1,0) = 3;
  B(2,1) = 4;  B(2,2) = 5;

  C(0,0) = 1;  C(0,1) = 2;

  std::cout << "C" << std::endl;
  print_all_matrix(C);

  mult(A, B, C);

  std::cout << "A" << std::endl;
  print_all_matrix(A);
  std::cout << "B" << std::endl;
  print_all_matrix(B);
  std::cout << "C = A*B" << std::endl;
  print_all_matrix(C);
}