#include <slepceps.h>
#include <iostream>
#include <cstdio>
#include <vector>
#include "hdf5.h"
#include "petscviewerhdf5.h"
#include "petscsys.h"
#include <string>
#include <algorithm>

static char help [ ] = "Test shift-and-invert eigensolver using SLEPc";

PetscErrorCode loadSparseMatrixFromHDF5( std::string filename, std::string groupname, Mat & Matrix )
{
    PetscErrorCode slepcerr;

    hid_t file = H5Fopen( filename.c_str(), H5F_ACC_RDONLY, H5P_DEFAULT );
    std::string setname = groupname + "/rows";
    const hid_t rowdataset = H5Dopen( file, setname.c_str(), H5P_DEFAULT );
    setname = groupname + "/cols";
    const hid_t coldataset = H5Dopen( file, setname.c_str(), H5P_DEFAULT );
    setname = groupname + "/vals";
    const hid_t valdataset = H5Dopen( file, setname.c_str(), H5P_DEFAULT );

    const hid_t rowspace = H5Dget_space( rowdataset );
    const std::size_t dimensions_size = 2;
    hsize_t dimensions[dimensions_size];
    hsize_t maxdimensions[dimensions_size];
    dimensions[1] = 1; // in case it's 1D
    H5Sget_simple_extent_dims( rowspace, dimensions, maxdimensions );
    const hsize_t length = dimensions[0];
    int * rowdata = new int[ length ];
    H5Dread( rowdataset, H5T_NATIVE_INT, H5S_ALL, H5S_ALL, H5P_DEFAULT, rowdata );
    int * coldata = new int[ length ];
    H5Dread( coldataset, H5T_NATIVE_INT, H5S_ALL, H5S_ALL, H5P_DEFAULT, coldata );
    double * valdata = new double[ length ];
    H5Dread( valdataset, H5T_NATIVE_DOUBLE, H5S_ALL, H5S_ALL, H5P_DEFAULT, valdata );

    /* Need to figure out how many rows are in the matrix */

    std::vector< int > RowVector( rowdata, rowdata + length );
    std::vector< int > Rows;
    std::unique_copy( RowVector.begin(), RowVector.end(), std::back_inserter(Rows) );
    int rowcount = Rows.size();
    slepcerr = MatSetSizes( Matrix, PETSC_DECIDE, PETSC_DECIDE, (PetscInt)rowcount, (PetscInt)rowcount );
    CHKERRQ( slepcerr );

    slepcerr = MatSetFromOptions( Matrix );
    CHKERRQ( slepcerr );

    slepcerr = MatSetUp( Matrix );
    CHKERRQ( slepcerr );

    PetscInt nodestart, nodeend;
    slepcerr = MatGetOwnershipRange( Matrix, &nodestart, &nodeend );
    CHKERRQ( slepcerr );
    /* Now populate the matrix */
    for( int myrow = nodestart; myrow < nodeend; myrow++ )
    {
        int start       = std::distance( RowVector.begin(), std::find( RowVector.begin(), RowVector.end(), myrow ) );
        hsize_t colsize = std::count( RowVector.begin(), RowVector.end(), myrow );
        PetscReal * myvals = new PetscReal[ colsize ]; 
        PetscInt * colinds = new PetscInt[ colsize ]; 
        for( hsize_t i = 0; i < colsize; i++ )
        {
            myvals[ i ]  = valdata[ start + i ];
            colinds[ i ] = coldata[ start + i ];
        }
        //Set the nth row of the matrix - each node owns one or more rows
        slepcerr = MatSetValues( Matrix, 1, &myrow, colsize, colinds, myvals, INSERT_VALUES );
        CHKERRQ( slepcerr );
        delete [] myvals;
        delete [] colinds;
    }
    delete [] rowdata; 
    delete [] coldata; 
    delete [] valdata; 
    H5Fclose( file );

    return slepcerr;
}

int main( int argc, char** argv )
{
    std::cout.precision( 12 ); 
    Mat H;
    EPS eps;
    PetscReal error;
    PetscErrorCode slepcerr;
    SlepcInitialize( &argc, &argv, (char*)0, help);
    /* Create a matrix stored across all MPI ranks */
    slepcerr = MatCreate( PETSC_COMM_WORLD, &H );
    CHKERRQ( slepcerr );
    std::string ifilename  = std::string(argv[2]);
    std::string ofilename  = std::string(argv[3]);
    std::string Vkey  = std::string(argv[4]);
    double V = 0.0;
    if (Vkey == "005")
    {
        V = 0.05;
    }
    else if (Vkey == "01")
    {
        V = 0.1;
    }
    else if (Vkey == "025")
    {
        V = 0.25;
    }
    else if (Vkey == "05")
    {
        V = 0.5;
    }
    else if (Vkey == "06")
    {
        V = 0.7;
    }
    else if (Vkey == "07")
    {
        V = 0.7;
    }
    else if (Vkey == "08")
    {
        V = 0.8;
    }
    else if (Vkey == "09")
    {
        V = 0.9;
    }
    else if (Vkey == "1")
    {
        V = 1.0;
    }
    else if (Vkey == "125")
    {
        V = 1.25;
    }
    else if (Vkey == "15")
    {
        V = 1.5;
    }
    else if (Vkey == "175")
    {
        V = 1.75;
    }
    else if (Vkey == "2")
    {
        V = 2.0;
    }
    else if (Vkey == "225")
    {
        V = 2.25;
    }
    else if (Vkey == "25")
    {
        V = 2.5;
    }
    else if (Vkey == "275")
    {
        V = 2.75;
    }
    else if (Vkey == "3")
    {
        V = 3.0;
    }
    else if (Vkey == "325")
    {
        V = 3.25;
    }
    else if (Vkey == "35")
    {
        V = 3.5;
    }
    else if (Vkey == "375")
    {
        V = 3.75;
    }
    else if (Vkey == "4")
    {
        V = 4.0;
    }
    MPI_Info info;
    MPI_Info_create(&info);
    MPI_Info_set(info,"romio_ds_write","disable");
    MPI_Info_set(info,"romio_ds_read","disable");
    MPI_Info_set(info,"collective buffering","true");
    MPI_Info_set(info,"romio_lustre_ds_in_coll","disable");
    double L          = std::stod(std::string(argv[5]));
    std::string groupname = "disorder" + std::string(argv[1]);
    slepcerr = loadSparseMatrixFromHDF5( ifilename, groupname, H );
    CHKERRQ( slepcerr );
    std::cout<<"Done loading"<<std::endl;
    MPI_Barrier(PETSC_COMM_WORLD);
    slepcerr = MatAssemblyBegin( H, MAT_FINAL_ASSEMBLY );
    CHKERRQ( slepcerr );
    slepcerr = MatAssemblyEnd( H, MAT_FINAL_ASSEMBLY );
    CHKERRQ( slepcerr );

    /* Need to ensure the vectors have same parallel distribution
    across nodes as the matrix does so we can extract eigenvectors */
    Vec EigReal, EigImag;
    slepcerr = MatCreateVecs( H, NULL, &EigReal );
    CHKERRQ( slepcerr );
    slepcerr = MatCreateVecs( H, NULL, &EigImag );
    CHKERRQ( slepcerr );
   
    /* Create and set up an eigensolver for the Hamiltonian.
    This is not a generalized eigenproblem so we only pass one matrix.
    Right now using Lanczos. For now we use a target value - search for 
    eigenvalues closest in magnitude to the target.*/  
    slepcerr = EPSCreate( PETSC_COMM_WORLD, &eps );
    CHKERRQ( slepcerr );
    
    slepcerr = EPSSetOperators( eps, H, NULL );
    CHKERRQ( slepcerr );
    
    slepcerr = EPSSetType( eps, EPSLANCZOS );
    CHKERRQ( slepcerr );
    slepcerr = EPSSetProblemType( eps, EPS_HEP );
    CHKERRQ( slepcerr );

    EPSSetTarget( eps, V*L/8.0 );
    slepcerr = EPSSetWhichEigenpairs( eps, EPS_TARGET_MAGNITUDE );
    CHKERRQ( slepcerr );

    ST spectral;
    slepcerr = EPSGetST( eps, &spectral );
    CHKERRQ( slepcerr );
    slepcerr = STSetType( spectral, STSINVERT );
    CHKERRQ( slepcerr );
    
    slepcerr = EPSSetFromOptions( eps );
    CHKERRQ( slepcerr );

    std::cout<<"r2solve"<<std::endl;
    MPI_Barrier(PETSC_COMM_WORLD);
    slepcerr = EPSSolve( eps );
    CHKERRQ( slepcerr );
    std::cout<<"Done solve"<<std::endl;
    PetscInt nconverged;
    slepcerr = EPSGetConverged( eps, &nconverged );
    CHKERRQ( slepcerr );
    
    int rank;
    MPI_Comm_rank(PETSC_COMM_WORLD, &rank);
    /* Set up HDF5 file to write eigenvectors to.*/
    PetscViewer hdf5viewer;
    slepcerr = PetscViewerHDF5Open( PETSC_COMM_WORLD, ofilename.c_str(), FILE_MODE_WRITE, &hdf5viewer );
    CHKERRQ( slepcerr );
    /*Create an HDF5 group and move down into it.*/
    slepcerr = PetscViewerHDF5PushGroup( hdf5viewer, groupname.c_str() );
    CHKERRQ( slepcerr );
     
    for( PetscInt i = 0; i < nconverged; i++ )
    {
        PetscScalar real, imag;
        slepcerr = EPSGetEigenpair( eps, i, &real, &imag, EigReal, EigImag );
        CHKERRQ( slepcerr );
        if( !rank )
        {
            std::cout<<real<<" "<<imag<<std::endl;
        }
        char eigname[15];
        sprintf( eigname, "Eig%d", i );
        slepcerr = PetscObjectSetName((PetscObject)EigReal, eigname);
        CHKERRQ( slepcerr );
        slepcerr = VecView( EigReal, hdf5viewer );
        CHKERRQ( slepcerr );
        char attname[15];
        sprintf( attname, "Value" );
        std::string pathname = groupname + "/" + eigname;
        slepcerr = PetscViewerHDF5WriteAttribute( hdf5viewer, pathname.c_str(), attname, PETSC_DOUBLE, &real );
        CHKERRQ( slepcerr );
        slepcerr = PetscViewerFlush( hdf5viewer );
        CHKERRQ( slepcerr );
    }

    slepcerr = EPSDestroy( &eps );
    CHKERRQ( slepcerr );
    
    slepcerr = MatDestroy( &H );
    CHKERRQ( slepcerr );
    
    slepcerr = VecDestroy( &EigReal );
    CHKERRQ( slepcerr );
    slepcerr = VecDestroy( &EigImag );
    CHKERRQ( slepcerr );
    
    slepcerr = PetscViewerDestroy( &hdf5viewer );
    CHKERRQ( slepcerr );
    MPI_Info_free(&info);
    slepcerr = SlepcFinalize();
    CHKERRQ( slepcerr );

    return 0;
}