/* Ergo, version 3.5, a program for linear scaling electronic structure
 * calculations.
 * Copyright (C) 2016 Elias Rudberg, Emanuel H. Rubensson, Pawel Salek,
 * and Anastasia Kruchinina.
 * 
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 * 
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 * 
 * Primary academic reference:
 * Kohn−Sham Density Functional Theory Electronic Structure Calculations 
 * with Linearly Scaling Computational Time and Memory Usage,
 * Elias Rudberg, Emanuel H. Rubensson, and Pawel Salek,
 * J. Chem. Theory Comput. 7, 340 (2011),
 * <http://dx.doi.org/10.1021/ct100611z>
 * 
 * For further information about Ergo, see <http://www.ergoscf.org>.
 */
#include "matrix_typedefs.h" // definitions of matrix types and interval type (source)
#include "realtype.h"   // definitions of types (utilities_basic)
#include "matrix_utilities.h"
#include "integral_matrix_wrappers.h"
#include "SizesAndBlocks.h"
#include "Matrix.h"
#include "Vector.h"
#include "MatrixSymmetric.h"
#include "MatrixTriangular.h"
#include "MatrixGeneral.h"
#include "VectorGeneral.h"
#include "output.h"

#include "files_dense.h"
#include "files_sparse.h"

#include <iostream>
#include <fstream>
#include <string.h>

#include "ErgoMatrix.h"

using namespace std;

typedef ergo_real real;
typedef symmMatrix MatrixType;

typedef ErgoMatrix<MatrixType> MatrixWrapperErgo;


int print_matrix = 0;
real TOL_SCALAR = 1e-12;
real TOL_EQUAL_MAT  = 1e-8;


int get_matrix_from_sparse(char *filename, MatrixType &X)
{
  vector<int> I, J;
  vector<ergo_real> val;
  int N, M;
  if(read_matrix_from_mtx(filename, I, J, val, N, M) == -1) 
    return -1;
    
  assert(N==M);

  /********** Initialization of SizesAndBlocks */
  int size = N;
  int nlevels = 5; //!!!
  std::vector<int> blockSizes(nlevels);
  blockSizes[nlevels - 1] = 1; // should always be one
    for (int ind = nlevels - 2; ind >= 0; ind--)
      blockSizes[ind] = blockSizes[ind + 1] * 10;
  mat::SizesAndBlocks rows(blockSizes, size);
  mat::SizesAndBlocks cols(blockSizes, size);
  /********************************************/  
  X.resetSizesAndBlocks(rows,cols);
  assert(X.get_nrows()*X.get_ncols() == N*N);

  X.assign_from_sparse(I, J, val);

  return 0;
}


void print_elements(MatrixType &X)
{
  int sizeX = X.get_nrows();
  cout << "rows of X = " << sizeX << endl;

  std::vector<int> row(sizeX*sizeX);
  std::vector<int> col(sizeX*sizeX);
  std::vector<real> val(sizeX*sizeX);

  X.get_all_values(row, col, val);
  assert(row.size() == col.size());

  std::vector<double> matrix(sizeX*sizeX);

  for(int i = 0; i < (int)row.size(); ++i )
    {
      matrix[row[i]*sizeX + col[i]] = val[i];
      matrix[col[i]*sizeX + row[i]] = val[i];
    }

  for (int i = 0; i < sizeX*sizeX; i++)
    {
      if(i%sizeX == 0 &&  i != 0) cout << endl;
      cout << matrix[i]<< " ";
    }
  cout << endl;
}




int main()
try{
  // READ THE MATRIX
  char filename[] = "test_files/F_1.mtx";

  MatrixType M;
  if( get_matrix_from_sparse(filename, M) == -1 )
    { 
      printf("Cannot read matrix from the file...\n");
      return EXIT_FAILURE;
    }

  if(print_matrix == 1)
    print_elements(M);

  // get rows and columns
  int n, m;
  m = M.get_ncols();
  n = M.get_nrows();
  assert(n == m);


  cout << "=================================" << endl;  
  cout << "Check constructors..." << endl;

  {
  MatrixWrapperErgo A;
  MatrixWrapperErgo B(M);
  MatrixWrapperErgo C(B);

  mat::SizesAndBlocks rows;
  mat::SizesAndBlocks cols;
  M.getRows(rows); 
  M.getCols(cols); 

  A.read_from_mtx(filename, rows, cols); // read M

  cout << "=================================" << endl;  
  cout << "Check eucledian norm of the difference..." << endl;
  real eucl = ErgoMatrix<MatrixType>::eucl_diff(A, B, 1e-12);
  assert(eucl < TOL_EQUAL_MAT);

  // check transfer
  MatrixWrapperErgo D(A);
  B.transfer_from(D); // here D is empty
  eucl = ErgoMatrix<MatrixType>::eucl_diff(A, B, 1e-12);
  assert(eucl < TOL_EQUAL_MAT);
  }


  MatrixWrapperErgo A(M);


  cout << "=================================" << endl;  
  cout << "Check gershgorin bounds..." << endl;

  real eigmin, eigmax;
  A.gershgorin(eigmin, eigmax);
  assert(abs(eigmin - (-17.120054726093)) < TOL_SCALAR);
  assert(abs(eigmax - 5.181705707709) < TOL_SCALAR);

  cout << "=================================" << endl;  
  cout << "Check square of the matrix..." << endl;
  {
    MatrixWrapperErgo B;
    A.square(B);
    MatrixType Tmp, Tmp2;
    A.get_matrix(Tmp);
    Tmp2 = (real)1.0*Tmp*Tmp;
    MatrixWrapperErgo C(Tmp2);
    real eucl_tmp = MatrixWrapperErgo::eucl_diff(C, B, 1e-12);
    assert(eucl_tmp < TOL_EQUAL_MAT);
  }

  {
    MatrixWrapperErgo B(A);
    cout << "=================================" << endl;  
    cout << "Check frobenius norm of the difference..." << endl;

    real frob = MatrixWrapperErgo::frob_diff(A, B);
    assert(frob < TOL_EQUAL_MAT);

    cout << "=================================" << endl;  
    cout << "Check mixed norm of the difference..." << endl;

    real mixed = MatrixWrapperErgo::mixed_diff(A, B, 1e-12);
    assert(mixed < TOL_EQUAL_MAT);

    cout << "=================================" << endl;  
    cout << "Check non-symm multiplication with forcing symmetry" << endl;

    MatrixWrapperErgo B2(A);
    B2.mult_force_symm(B);

    MatrixWrapperErgo C;
    A.square(C);

    real eucl_tmp = MatrixWrapperErgo::eucl_diff(B2, C, 1e-12);
    assert(eucl_tmp < TOL_EQUAL_MAT);
    
  }


  cout << "=================================" << endl;  
  cout << "Check truncation..." << endl;

  real th = A.thresh(0.0001, mat::euclNorm);
  cout << "Truncation return value th = " << th << endl;


  cout << "OK!" << endl;

  return 0;
}
 catch(char const * e)
   {
     std::cout << e << std::endl;
   }
 catch(std::exception &e)
   {
     std::cout << e.what() << std::endl;
   }