#! /usr/bin/python

import numpy as np
from scipy import sparse


# .....................
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
# .....................
def computeLocalnnz(nodelist, dof=1):
    # compute the number of non-zeros;
    nnzeros = 0
    for i in nodelist:
        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
    return nnzeros
# .....................

# .....................
def computeNeighbors():
    list_ineighbor = [[]]*n
    for i in range(1, n+1):
        list_j = []
        if (i == 1) or (i == n):
            list_j.append(i)
        else:
            for k in range(-1,2):
                list_j.append(i+k)
        list_ineighbor[i-1] = list_j

#    print list_ineighbor
    return list_ineighbor
# .....................

dof = 1 # to move to sys
n = 10 # to move to sys
xmin = 0.
xmax = 1.

list_i = [] ; list_j = [] ; list_v = []

# Dirichlet boundary condition
i = 1 ; j = 1 ; val = GetCoef(i,i,xmin,xmax,n)
list_i.append(i) ; list_j.append(j) ; list_v.append(val)

i = n ; j = n ; val = GetCoef(i,i,xmin,xmax,n)
list_i.append(i) ; list_j.append(j) ; list_v.append(val)

# Interior
for i in range(2, n):
    for k in range(-1,2):
        j = i + k
        val = GetCoef(j,i,xmin,xmax,n)
        list_i.append(i) ; list_j.append(j) ; list_v.append(val)

I = np.array(list_i) - 1
J = np.array(list_j) - 1
V = np.array(list_v)

coo_mat = sparse.coo_matrix((V,(I,J)),shape=(n,n))
csr_mat = coo_mat.tocsr()
coo_mat = csr_mat.tocoo()
