#! /usr/bin/python

# test file based on Murge-Fortran.F90

from mpi4py import MPI
from pymurge import MURGE
from pymurge import Matrix
import numpy as np
from time import time

execfile('utils.py')

root = -1

# ... 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, comm=comm)

# ... Graph initialization
edgenbr = 3*n-4

A.GraphBegin(edgenbr, rank=rank)
# Dirichlet boundary condition
A.GraphEdge(1, 1, rank=rank)
A.GraphEdge(n, n, rank=rank)
# Interior
for i in range(2, n):
    for k in range(-1,2):
        A.GraphEdge(i, i+k, rank=rank)

A.GraphEnd()
# ...

# Get Local nodes
nodelist = A.GetLocalNodeList()
# compute the number of non-zeros;
nnzeros = computeLocalnnz(nodelist)
A.localnnz = nnzeros

# ...
# Assembling
# ...
list_ineighbors = computeNeighbors()
def _GetCoef(i,j):
    return GetCoef(i,j,xmin,xmax,n)

A.Assembly(_GetCoef, list_ineighbors=list_ineighbors)
# ...

# ...
# We expand the rhs
# ...
def _GetRhs(i):
    return GetRhs(i,xmin,xmax,n)

lrhs = A.AssemblyRHS(_GetRhs, local=True)
globrhs = A.AssemblyRHS(_GetRhs, local=False)
# ...

## Set the local RHS
#A.SetRHS(lrhs, local=True)
#
## Get the global solution
#X = A.GetSolution(local=False, root=root)

X = A.Solve(lrhs, local=True, root=root)

# Compute the product A X
Y = A.dot(X, local=False, root=root)
if rank==0:
    print np.linalg.norm(globrhs-Y)

