#! /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
from time import time

execfile('utils.py')

id = 0
root = -1
one = 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
print "Rank ", rank," has Local nnz = ", nnzeros


list_ineighbors = computeNeighbors()

# ...
# Assembling
# ...
if True:
    t_start = time()
    A.AssemblyBegin()
    for i in nodelist:
        list_j = list_ineighbors[i-1]
        for j in list_j:
            val = GetCoef(i,j,xmin,xmax,n)
            A.AssemblySetValue(i, j, val)
    A.AssemblyEnd()
    t_end = time()
    if rank==0:
        print "Native - Matrix Assembly elapsed time : ", t_end-t_start
# ...
if True:
    t_start = time()
    def _GetCoef(i,j):
        return GetCoef(i,j,xmin,xmax,n)

    A.Assembly(_GetCoef, list_ineighbors=list_ineighbors)
    t_end = time()
    if rank==0:
        print "CallBack - Matrix Assembly elapsed time : ", t_end-t_start
# ...

# ...
# We expand the rhs
# ...
if True:
    def AssemblyLocalRHS():
        localnodenbr = len(nodelist)
        lrhs    = np.zeros(localnodenbr*dof, dtype=np.double)
        k = 0
        for i in nodelist:
            val = GetRhs(i,xmin,xmax,n)
            lrhs[k:k+dof] = val
            k = k + dof
        return lrhs
    t_start = time()
    lrhs = AssemblyLocalRHS()
    t_end = time()
    if rank==0:
        print "Native - Local RHS Assembly elapsed time : ", t_end-t_start
# ...
if True:
    t_start = time()
    def _GetRhs(i):
        return GetRhs(i,xmin,xmax,n)

    lrhs = A.AssemblyRHS(_GetRhs, local=True)
    t_end = time()
    if rank==0:
        print "CallBack - Local RHS Assembly elapsed time : ", t_end-t_start
# ...

# ...
if True:
    t_start = time()
    def AssemblyGlobalRHS():
        globrhs = np.zeros(n*dof, dtype=np.double)
        for i in nodelist:
            val = GetRhs(i,xmin,xmax,n)
            globrhs[(i-1)*dof:(i-1)*dof+dof] = val

        globrhs_recv = np.zeros(n*dof, dtype=np.double)
        comm.Reduce([globrhs, MPI.DOUBLE], [globrhs_recv, MPI.DOUBLE], op=MPI.SUM, root=0)
        return globrhs_recv

    globrhs_recv = AssemblyGlobalRHS()
    t_end = time()
    if rank==0:
        print "Native - Global RHS Assembly elapsed time : ", t_end-t_start
# ...
if True:
    t_start = time()
    def _GetRhs(i):
        return GetRhs(i,xmin,xmax,n)
    globrhs_ = A.AssemblyRHS(_GetRhs, local=False)
    t_end = time()
    if rank==0:
        print "CallBack - Global RHS Assembly elapsed time : ", t_end-t_start
# ...

# Set the local RHS
A.SetRHS(lrhs, local=True)

# Get the global solution
X = A.GetSolution(local=False, root=root)

Y = A.dot(X, local=False, root=root)
if rank==0:
    print np.linalg.norm(globrhs_-Y)
#    print np.linalg.norm(globrhs_recv-Y)

