#! /usr/bin/python

# test file based on Murge-Fortran.F90

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

# .....................
def GetCoef(i,j,xmin,xmax,n):
    """
    Computes the value for a given coefficient.

    Parameters:
      val  - Value to set
      i    - Row of the coefficient.
      j    - Column of the coefficient.
      xmin - Minimum value of the interval.
      xmax - Maximum value of the interval.
      n    - Number of points in the interval.
    """
    dx_1 = (n-1.)/ (xmax - xmin)
    val = 0.

    if (i==j):
        if (i==1) or (i == n):
            # Boundary Condition (Dirichlet)
            val = 1.
        else:
            # Interior diagonnal part
            val = -2 * dx_1
    else:
        val = dx_1

    return val
# .....................

# .....................
sin = np.sin
pi = np.pi
def GetRhs(i,xmin,xmax,n):
    """
        computes the value of a coefficient of the Right-hand-side member.

        Parameters:
         val  - Value to set.
         i    - Index of the value.
         xmin - Minimum value of the interval.
         xmax - Maximum value of the interval.
         n    - Number of points in the interval.
    """
    dx = (xmax - xmin) / (n-1.)
    x = xmin + (i-1)*dx

    val = -4 * pi**2 * np.sin(2 * pi * x)
    # Boundary Condition (Dirichlet)
    if (i == n) or (i==1):
       val = 0.
    else:
       val = dx*val

    return val
# .....................


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

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

dof = 1 # to move to sys
n = 10 # to move to sys

root = -1
zero = 0
one = 1
base = 1
tol = 1.e-7
xmin = 0.
xmax = 1.

murge = MURGE(nmatrices = 1)

id = 0
murge.SetDefaultOptions(id, zero)

murge.SetOptionInt(id, MURGE.IParam.DOF, dof)
murge.SetOptionInt(id, MURGE.IParam.SYM, MURGE.Boolean.FALSE)
murge.SetOptionInt(id, MURGE.IParam.BASEVAL, base)

murge.SetOptionReal(id, MURGE.RParam.EPSILON_ERROR, tol)

#
# Set the graph : all processor enter some edge of the
# graph that corresponds to non-zeros location in the matrix
#
#
# ****************************************
# ** Enter the matrix non-zero pattern  **
# ** you can use any distribution       **
# ****************************************
#
# this processor enters the A(myfirstrow:mylastrow, *)
# part of the matrix non-zero pattern
#
if rank == 0:
    edgenbr = 3*n-4
    murge.GraphBegin(id, n, edgenbr)
    # Dirichlet boundary condition
    murge.GraphEdge(id, one, one)
    murge.GraphEdge(id, n, n)
    # Interior
    for i in range(2, n):
        for j in range(-1,2):
            murge.GraphEdge(id, i, i+j)
else:
    edgenbr = 0
    murge.GraphBegin(id, n, edgenbr)

murge.GraphEnd(id)
# ...

# Get Local nodes
localnodenbr = np.zeros(1, dtype=np.int32)
murge.GetLocalNodeNbr(id, localnodenbr)
localnodenbr = localnodenbr[0]
print 'localnodenbr : ' , localnodenbr

nodelist = np.zeros(localnodenbr, dtype=np.int32)
murge.GetLocalNodeList(id, nodelist)
#print "nodelist = ", nodelist
mode = 0

# job = 0 : initialize the matrix with this coefficient
# job = 1 : these coefficients are added to the matrix
job = 0

# compute the number of non-zeros;
nnzeros = 0
for m in range(0, localnodenbr):
    i = nodelist[m]
    if (i == 1) or (i == n):
        # Boundaries
        nnzeros = nnzeros + 1
    else:
        # Interior
        for k in range(-1, 2):
            nnzeros = nnzeros + 1

# We are using dof so a non zero is in fact a block of size dof**2
nnzeros = nnzeros * dof**2

# You can enter only coefficient (i,j) that are in A(nodelist, nodelist)
# on this processor
# We enter the lower and upper triangular part of the matrix so sym = 0
# expand is the identity matrix of size 'dof' stored by line
expand = np.zeros(dof**2)

k = 0
for i in range(1,dof+1):
    for j in range(1,dof+1):
        if (i == j):
            expand[k] = 1.
        k = k + 1

# ...
# Assembling
# ...
murge.AssemblyBegin( id, n, nnzeros \
                   , MURGE.AssemblyOp.OVW, MURGE.AssemblyOp.OVW \
                   , MURGE.AssemblyMode.FOOL, MURGE.Boolean.FALSE)

for m in range(0, localnodenbr):
    i = nodelist[m]
#    print "i = ", i
    if (i == 1) or (i == n):
        # Boundaries
        val = GetCoef(i,i,xmin,xmax,n)
#        print val
#        murge.AssemblySetValue(id, i, i, val)
        murge.AssemblySetNodeValues(id, i, i, val*expand)
    else:
        for k in range(-1,2):
            val = GetCoef(i+k,i,xmin,xmax,n)
#            murge.AssemblySetValue(id, i, i+k, val)
            murge.AssemblySetNodeValues(id, i, i+k, val*expand)
#            print val

murge.AssemblyEnd(id)
# ...



# We expand the rhs
lrhs    = np.zeros(localnodenbr*dof, dtype=np.double)
globrhs = np.zeros(n*dof, dtype=np.double)
k = 0
for m in range(0,localnodenbr):
    i = nodelist[m]
    val = GetRhs(i,xmin,xmax,n)
    globrhs[(i-1)*dof:(i-1)*dof+dof] = val
    lrhs[k:k+dof] = val
    k = k + dof

globrhs_recv = np.zeros(n*dof, dtype=np.double)

comm.Reduce([globrhs, MPI.DOUBLE], [globrhs_recv, MPI.DOUBLE], op=MPI.SUM, root=0)

murge.SetLocalRHS(id, lrhs, MURGE.AssemblyOp.OVW, MURGE.AssemblyOp.OVW)

# Get the global solution
globx = np.zeros(n*dof, dtype=np.double)
murge.GetGlobalSolution(id, globx, root)

murge.SetGlobalRHS(id, globx, -one, MURGE.AssemblyOp.OVW)

globprod = np.zeros(n*dof, dtype=np.double)
murge.GetGlobalProduct(id, globprod, -one)

print "||AX - B||/||AX||  :", np.sqrt(np.sum((globprod - globrhs_recv)**2)/np.sum(globprod**2))

# Store in a file
#if Me == 0:
#    store(globx,xmin,xmax,dof)
#murge.Save(id, "data")

murge.Clean(id)
