#! /usr/bin/python

# test file based on Murge-Fortran.F90

from mpi4py import MPI
try:
    from pypastix import MURGE, Matrix
except Exception:
    try:
        from pyhips import MURGE, Matrix
    except Exception:
        raise
import numpy as np

execfile('utils.py')

# ... MPI Initialization
size = MPI.COMM_WORLD.Get_size()
rank = MPI.COMM_WORLD.Get_rank()

null = MPI.COMM_NULL
comm = MPI.COMM_WORLD
# ...

murge = MURGE(nmatrices = 1)
A = Matrix(murge, n, id=0)

# ... Graph initialization
#_indptr  = np.int64(csr_mat.indptr.copy())  + 1
#_indices = np.int32(csr_mat.indices.copy()) + 1
#A.GraphGlobal(indptr=_indptr, indices=_indices, format='csr')

_rows = np.int64(coo_mat.row.copy()) + 1
_cols = np.int64(coo_mat.col.copy()) + 1
ic = np.argsort(_cols)

rows = np.int32([_rows[i] for i in ic])
cols = np.int32([_cols[i] for i in ic])

A.GraphGlobal(rows=rows, cols=cols, format='ijv')
# ...

# Get Local nodes
nodelist = A.GetLocalNodeList()
#print "nodelist : ", nodelist
#import sys; sys.exit(0)

# compute the number of non-zeros;
A.localnnz = computeLocalnnz(nodelist)

# ...
# Assembling
# ...
A.AssemblyBegin()

for i in nodelist:
    if (i == 1) or (i == n):
        # Boundaries
        val = GetCoef(i,i,xmin,xmax,n)
        A.AssemblySetValue(i, i, val)
    else:
        for k in range(-1,2):
            val = GetCoef(i+k,i,xmin,xmax,n)
            A.AssemblySetValue(i, i+k, val)

A.AssemblyEnd()
# ...
