# Example translated from C to Julia based on this original source (MIT license):
# https://github.com/hypre-space/hypre/blob/ac9d7d0d7b43cd3d0c7f24ec5d64b58fbf900097/src/examples/ex5.c

# Example 5
#
# Interface:    Linear-Algebraic (IJ)
#
# Sample run:   mpirun -np 4 julia ex5.jl
#
# Description:  This example solves the 2-D Laplacian problem with zero boundary
#               conditions on an n x n grid.  The number of unknowns is N=n^2.
#               The standard 5-point stencil is used, and we solve for the
#               interior nodes only.
#
#               This example solves the same problem as Example 3.  Available
#               solvers are AMG, PCG, and PCG with AMG or Parasails
#               preconditioners.

using MPI
using HYPRE.LibHYPRE # LibHYPRE submodule contains the raw C-interface generated by Clang.jl

function main()

    # Initialize MPI
    # MPI_Init(Ref{Cint}(length(ARGS)), ARGS)
    MPI.Init()
    comm = MPI.COMM_WORLD
    myid = MPI.Comm_rank(comm)
    num_procs = MPI.Comm_size(comm)

    # Initialize HYPRE
    HYPRE_Init()

    # Default problem parameters
    n = 33
    solver_id = 0
    print_system = false

    # TODO: Parse command line

    # Preliminaries: want at least one processor per row
    if n * n < num_procs; n = trunc(Int, sqrt(n)) + 1; end
    N = n * n # global number of rows
    h = 1.0 / (n + 1) # mesh size
    h2 = h * h

    # Each processor knows only of its own rows - the range is denoted by ilower
    # and upper.  Here we partition the rows. We account for the fact that
    # N may not divide evenly by the number of processors.
    local_size = N ÷ num_procs
    extra = N - local_size * num_procs

    ilower = local_size * myid
    ilower += min(myid, extra)

    iupper = local_size * (myid + 1)
    iupper += min(myid + 1, extra)
    iupper = iupper - 1

    # How many rows do I have?
    local_size = iupper - ilower + 1

    # Create the matrix.
    # Note that this is a square matrix, so we indicate the row partition
    # size twice (since number of rows = number of cols)
    Aref = Ref{HYPRE_IJMatrix}(C_NULL)
    HYPRE_IJMatrixCreate(comm, ilower, iupper, ilower, iupper, Aref)
    A = Aref[]

    # Choose a parallel csr format storage (see the User's Manual)
    HYPRE_IJMatrixSetObjectType(A, HYPRE_PARCSR)

    # Initialize before setting coefficients
    HYPRE_IJMatrixInitialize(A)

    # Now go through my local rows and set the matrix entries.
    # Each row has at most 5 entries. For example, if n=3:

    # A = [M -I 0; -I M -I; 0 -I M]
    # M = [4 -1 0; -1 4 -1; 0 -1 4]

    # Note that here we are setting one row at a time, though
    # one could set all the rows together (see the User's Manual).
    let
        values = Vector{Float64}(undef, 5)
        cols = Vector{HYPRE_Int}(undef, 5)
        tmp = Vector{HYPRE_Int}(undef, 2)

        for i in ilower:iupper
            nnz = 1

            # The left identity block:position i-n
            if (i - n) >= 0
                cols[nnz] = i - n
                values[nnz] = -1.0
                nnz += 1
            end

            # The left -1: position i-1
            if i % n != 0
                cols[nnz] = i - 1
                values[nnz] = -1.0
                nnz += 1
            end

            # Set the diagonal: position i
            cols[nnz] = i
            values[nnz] = 4.0
            nnz += 1

            # The right -1: position i+1
            if (i + 1) % n != 0
                cols[nnz] = i + 1
                values[nnz] = -1.0
                nnz += 1
            end

            # The right identity block:position i+n
            if (i + n) < N
                cols[nnz] = i + n
                values[nnz] = -1.0
                nnz += 1
            end

            # Set the values for row i
            tmp[1] = nnz - 1
            tmp[2] = i

            HYPRE_IJMatrixSetValues(A, 1, Ref(tmp[1]), Ref(tmp[2]), cols, values)
        end
    end

    # Assemble after setting the coefficients
    HYPRE_IJMatrixAssemble(A)

    # Note: for the testing of small problems, one may wish to read
    # in a matrix in IJ format (for the format, see the output files
    # from the -print_system option).
    # In this case, one would use the following routine:
    # HYPRE_IJMatrixRead( <filename>, MPI_COMM_WORLD,
    #                     HYPRE_PARCSR, &A );
    # <filename>  = IJ.A.out to read in what has been printed out
    # by -print_system (processor numbers are omitted).
    # A call to HYPRE_IJMatrixRead is an *alternative* to the
    # following sequence of HYPRE_IJMatrix calls:
    # Create, SetObjectType, Initialize, SetValues, and Assemble

    # Get the parcsr matrix object to use
    parcsr_A_ref = Ref{Ptr{Cvoid}}(C_NULL)
    HYPRE_IJMatrixGetObject(A, parcsr_A_ref)
    parcsr_A = convert(Ptr{HYPRE_ParCSRMatrix}, parcsr_A_ref[])

    # Create the rhs and solution
    b_ref = Ref{HYPRE_IJVector}(C_NULL)
    HYPRE_IJVectorCreate(comm, ilower, iupper, b_ref)
    b = b_ref[]
    HYPRE_IJVectorSetObjectType(b, HYPRE_PARCSR)
    HYPRE_IJVectorInitialize(b)

    x_ref = Ref{HYPRE_IJVector}(C_NULL)
    HYPRE_IJVectorCreate(comm, ilower, iupper, x_ref)
    x = x_ref[]
    HYPRE_IJVectorSetObjectType(x, HYPRE_PARCSR)
    HYPRE_IJVectorInitialize(x)

    # Set the rhs values to h^2 and the solution to zero
    let
        rhs_values = Vector{Float64}(undef, local_size)
        x_values = Vector{Float64}(undef, local_size)
        rows = Vector{HYPRE_Int}(undef, local_size)

        for i in 1:local_size
            rhs_values[i] = h2
            x_values[i] = 0.0
            rows[i] = ilower + i - 1
        end

        HYPRE_IJVectorSetValues(b, local_size, rows, rhs_values)
        HYPRE_IJVectorSetValues(x, local_size, rows, x_values)
    end

    HYPRE_IJVectorAssemble(b)
    par_b_ref = Ref{Ptr{Cvoid}}(C_NULL)
    HYPRE_IJVectorGetObject(b, par_b_ref)
    par_b = convert(Ptr{HYPRE_ParVector}, par_b_ref[])

    HYPRE_IJVectorAssemble(x)
    par_x_ref = Ref{Ptr{Cvoid}}(C_NULL)
    HYPRE_IJVectorGetObject(x, par_x_ref)
    par_x = convert(Ptr{HYPRE_ParVector}, par_x_ref[])

    # Print out the system - files names will be IJ.out.A.XXXXX
    # and IJ.out.b.XXXXX, where XXXXX = processor id
    if print_system
        HYPRE_IJMatrixPrint(A, "IJ.out.A")
        HYPRE_IJVectorPrint(b, "IJ.out.b")
    end

    # Choose a solver and solve the system
    solver_ref = Ref{HYPRE_Solver}(C_NULL)

    # AMG
    if solver_id == 0
        # Create solver
        HYPRE_BoomerAMGCreate(solver_ref)
        solver = solver_ref[]

        # Set some parameters (See Reference Manual for more parameters)
        HYPRE_BoomerAMGSetPrintLevel(solver, 3) # print solve info + parameters
        HYPRE_BoomerAMGSetOldDefault(solver)    # Falgout coarsening with modified classical interpolaiton
        HYPRE_BoomerAMGSetRelaxType(solver, 3)  # G-S/Jacobi hybrid relaxation
        HYPRE_BoomerAMGSetRelaxOrder(solver, 1) # uses C/F relaxation
        HYPRE_BoomerAMGSetNumSweeps(solver, 1)  # Sweeeps on each level
        HYPRE_BoomerAMGSetMaxLevels(solver, 20) # maximum number of levels
        HYPRE_BoomerAMGSetTol(solver, 1e-7)     # conv. tolerance

        # Now setup and solve!
        HYPRE_BoomerAMGSetup(solver, parcsr_A, par_b, par_x)
        HYPRE_BoomerAMGSolve(solver, parcsr_A, par_b, par_x)

        # Run info - needed logging turned on
        num_iterations = Ref{HYPRE_Int}()
        final_res_norm = Ref{Float64}()
        HYPRE_BoomerAMGGetNumIterations(solver, num_iterations)
        HYPRE_BoomerAMGGetFinalRelativeResidualNorm(solver, final_res_norm)
        if myid == 0
            println()
            println("Iterations = $(num_iterations[])")
            println("Final Relative Residual Norm = $(final_res_norm[])")
            println()
        end

        # Destroy solver
        HYPRE_BoomerAMGDestroy(solver)
    else
        if myid == 0
            println("Invalid sovler id specified.")
        end
    end

    # Clean up
    HYPRE_IJMatrixDestroy(A)
    HYPRE_IJVectorDestroy(b)
    HYPRE_IJVectorDestroy(x)

    # Finalize HYPRE
    HYPRE_Finalize()

    # Finalize MPI
    MPI.Finalize()

    return 0
end

# Run it
if abspath(PROGRAM_FILE) == @__FILE__
    main()
end