#include <petscmat.h>
#include <petscvec.h>
#include <petsctime.h>
#include <petscsnes.h>
#include <petscdmda.h>
#include "petscsys.h"
#include <stdio.h>


void create_matrix(MPI_Comm comm, Mat* mat, PetscInt Nrow, PetscInt Ncol, MatType type){
  MatCreate(comm, mat);
  MatSetSizes(*mat,PETSC_DECIDE,PETSC_DECIDE,Nrow,Ncol);
  MatSetType(*mat,type);
  MatSetUp(*mat); 
}

void create_vector(MPI_Comm comm, Vec* mat, PetscInt N, VecType type){
  VecCreate(comm, mat);
  VecSetSizes(*mat, PETSC_DECIDE, N);
  VecSetType(*mat,type);
  VecSetUp(*mat);
}

void destroy_matrix(Mat* mat){
  MatDestroy(mat);
}

void destroy_vector(Vec* mat){
  VecDestroy(mat);
}

Mat create_identity_matrix(MPI_Comm comm,
        PetscMPIInt size,
        PetscMPIInt rank,
        PetscInt Nrow,
        PetscInt Ncol,
        PetscScalar a){

  PetscInt i, j;
  PetscScalar val;
  Mat X;
  MatType matType;

  if(size == 1) matType = MATSEQAIJ;
  else matType = MATMPIAIJ;

  create_matrix(comm, &X, Nrow, Ncol, matType); 

  MatAssemblyBegin(X,MAT_FINAL_ASSEMBLY);
  MatZeroEntries(X);
  MatAssemblyEnd(X,MAT_FINAL_ASSEMBLY);
  MatShift(X, a);
  return X;
}

Mat create_dense_matrix(MPI_Comm comm,
	PetscMPIInt size,
	PetscMPIInt rank,
	PetscInt Nrow,
	PetscInt Ncol, 
	PetscScalar a){

  PetscInt i, j;
  PetscScalar val;
  Mat X;
  MatType matType;

  if(size == 1) matType = MATSEQDENSE;
  else matType = MATMPIDENSE; 

  create_matrix(comm, &X, Nrow, Ncol, matType);

  if(rank == 0)
    for(i = 0; i < Nrow; i++)
      for(j = 0; j < Ncol; j++){
        val = a*(i + 1);
        MatSetValues(X, 1, &i, 1, &j, &val, INSERT_VALUES); 
      }

  MatAssemblyBegin(X,MAT_FINAL_ASSEMBLY);
  MatAssemblyEnd(X,MAT_FINAL_ASSEMBLY);
  return X;
}

PetscReal check_error_matrix(Mat* real_solution, Mat* solution, NormType normType){
  Mat w;
  PetscErrorCode ierr;
  MatType type;
  ierr = MatGetType(*real_solution, &type);CHKERRQ(ierr);
  MPI_Comm comm;
  ierr = PetscObjectGetComm((PetscObject)*real_solution,&comm);CHKERRQ(ierr);
  PetscInt Nrow, Ncol;
  MatGetSize(*real_solution,&Nrow,&Ncol);CHKERRQ(ierr);
  create_matrix(comm, &w, Nrow, Ncol, type);
  PetscReal diff;
  ierr = MatCopy(*real_solution,w,SAME_NONZERO_PATTERN);CHKERRQ(ierr);
  ierr = MatAYPX(w,-1.0,*solution,DIFFERENT_NONZERO_PATTERN);CHKERRQ(ierr);
  ierr = MatNorm(w,normType,&diff);CHKERRQ(ierr);
  destroy_matrix(&w);
  return diff;
}

int main(int argc,char **argv){
  PetscInt       A_size, number_rhs, i;
  PetscErrorCode ierr;
  PetscBool      allowd_MatMatSolve, multiple_systems;
  Mat            A, B, X, real_X;
  Vec		 x,b;
  VecType 	 vecType;
  PetscMPIInt    rank, size;

  PetscInitialize(&argc,&argv,0,NULL);

  A_size = 10;
  number_rhs = 2;

  ierr = MPI_Comm_rank(MPI_COMM_WORLD,&rank);CHKERRQ(ierr);
  ierr = MPI_Comm_size(MPI_COMM_WORLD,&size);CHKERRQ(ierr);

  A      = create_identity_matrix(MPI_COMM_WORLD, size, rank, A_size, A_size, 1);
  real_X = create_dense_matrix(MPI_COMM_WORLD, size, rank, A_size, number_rhs, 2);
  ierr =  MatMatMult(A, real_X, MAT_INITIAL_MATRIX, PETSC_DEFAULT, &B);CHKERRQ(ierr); 

  

  MatType matType;
  if(size == 1) matType = MATSEQDENSE;
  else matType = MATMPIDENSE;  
  create_matrix(MPI_COMM_WORLD, &X, A_size, number_rhs, matType);

  if(size == 1) vecType = VECSEQ;
  else vecType = VECMPI;

  PetscLogDouble time1, time2;
  ierr = PetscTime(&time1);CHKERRQ(ierr);
  KSP ksp;
  PC pc;
  Mat F;
  ierr = KSPCreate(MPI_COMM_WORLD, &ksp);CHKERRQ(ierr);
  ierr = KSPSetOperators(ksp, A, A, SAME_PRECONDITIONER);CHKERRQ(ierr);
  ierr = KSPSetType(ksp, KSPPREONLY);CHKERRQ(ierr);
  ierr = KSPGetPC(ksp, &pc);CHKERRQ(ierr);
  ierr = PCSetType(pc, PCLU);CHKERRQ(ierr);
  ierr = PCFactorSetMatSolverPackage(pc, MATSOLVERSUPERLU_DIST);CHKERRQ(ierr);
  ierr = KSPSetUp(ksp);
  ierr = PCFactorGetMatrix(pc, &F);CHKERRQ(ierr);
  ierr = MatMatSolve(F, B, X);CHKERRQ(ierr);
  KSPDestroy(&ksp);
  ierr = PetscTime(&time2);CHKERRQ(ierr);
  PetscReal error;
  error = check_error_matrix(&real_X, &X, NORM_INFINITY);
  time1 = time2 - time1;
  if(size > 1){
    ierr = MPI_Reduce(&time1, &time2, 1, MPI_DOUBLE, MPI_MAX, 0, MPI_COMM_WORLD);
  } else {
    time2 = time1;
  }
  
  if(rank == 0){
    printf("Test: %s\n", "superlu_dist");
    printf("Error: %f\n", (double)error);
    printf("Time: %f\n", (double)time2);
  }
  destroy_matrix(&A);
  destroy_matrix(&B);
  destroy_matrix(&real_X);
  destroy_matrix(&X);
  ierr = PetscFinalize();
  return 0;
}