import numpy as np
import scipy.sparse
%load_ext cython

"""
Explanation: A comparison of dense and sparse matrix vector multiplication routines
End of explanation
"""


p = 0.01
Nc, Na = 10000, 200
c = np.ones(Nc)
a = np.ones(Na)
K = np.random.random((Nc, Na)) < p

"""
Explanation: Setup
End of explanation
"""


%timeit K.dot(a)

%timeit c.dot(K)

"""
Explanation: Dense matrix vector multiplication
End of explanation
"""


Ksp = scipy.sparse.csr_matrix(K)

%timeit scipy.sparse.csr_matrix(K)

np.all(Ksp.dot(a) == K.dot(a))

%timeit Ksp.dot(a)

np.all(Ksp.transpose(copy=False).dot(c) == c.dot(K))

%timeit Ksp.transpose(copy=False).dot(c)

csp = scipy.sparse.csr_matrix(c)
%timeit csp.dot(Ksp)

"""
Explanation: Sparse matrix vector multiplication
End of explanation
"""


# inspired by http://stackoverflow.com/questions/17158893/does-scipy-support-multithreading-for-sparse-matrix-multiplication-when-using-mk
# and https://github.com/afedynitch/MCEq/blob/master/MCEq/kernels.py

from ctypes import POINTER,c_void_p,c_int,c_char,c_double,byref,cdll

def SpMV_viaMKL(A, x, trans=False):
    """
    Wrapper to Intel's Sparse Matrix-Vector multiplaction routine.
    Handles rectangular matrices
    """
 
    mkl = cdll.LoadLibrary("libmkl_rt.so")
    mkl.mkl_set_num_threads(byref(c_int(4)))

    SpMV = mkl.mkl_dcsrmv
    (m, k) = A.shape

    data    = A.data.ctypes.data_as(POINTER(c_double))
    pb = A.indptr[:-1].ctypes.data_as(POINTER(c_int))
    pe = A.indptr[1:].ctypes.data_as(POINTER(c_int))
    indices = A.indices.ctypes.data_as(POINTER(c_int))

    # Allocate output, using same conventions as input
    insize = m if trans else k
    outsize = k if trans else m
    y = np.empty(outsize, dtype=np.double, order='F')
    if x.size != insize:
        raise Exception("x must have n entries. x.size is %d, n is %d" % (x.size, outsize))

    # Check input
    if x.dtype.type is not np.double:
        x = x.astype(np.double, copy=True)

    np_x = x.ctypes.data_as(POINTER(c_double))
    np_y = y.ctypes.data_as(POINTER(c_double))
    # now call MKL. This returns the answer in np_y, which links to y
    alpha = c_double(1.0)
    beta = c_double(0.0)
    npmatd = np.chararray(6)
    npmatd[0] = 'G'
    npmatd[3] = 'C'
    matdescra = npmatd.ctypes.data_as(POINTER(c_char))
    SpMV(byref(c_char("T" if trans else "N")), byref(c_int(m)), byref(c_int(k)), byref(alpha),
         matdescra, data, indices, pb, pe, np_x, byref(beta), np_y ) 
    return y

Kfloat = K.astype(np.float)
Kfloatsp = scipy.sparse.csr_matrix(Kfloat)

np.all(SpMV_viaMKL(Kfloatsp, a) == Ksp.dot(a))

%timeit SpMV_viaMKL(Kfloatsp, a)

np.all(SpMV_viaMKL(Kfloatsp, c, True)  == c.dot(K))

%timeit SpMV_viaMKL(Kfloatsp, c, True)

%prun [SpMV_viaMKL(Kfloatsp, c, True) for i in range(1000)]

%%cython -l mkl_core -l mkl_intel_lp64
cimport numpy as np
import numpy as np

cdef extern from "mkl_types.h":
    ctypedef MKL_INT

cdef extern from "mkl.h" nogil:
    double cblas_dasum (MKL_INT n, double *x, MKL_INT incx);

def cythonSpMV_viaMKL(np.ndarray[np.double_t] x):
    """
    Wrapper to Intel's Sparse Matrix-Vector multiplaction routine.
    Handles rectangular matrices
    """
    
    #cdef MKL_INT n = x.shape[0]
    #cdef MKL_INT incx = 1
    return 2#cblas_dasum(n, &x[0], incx)

%timeit cythonSpMV_viaMKL(Kfloatsp, c, True)

"""
Explanation: Sparse matrix vector multiplication using MKL
Intel's Math Kernel Library (MKL) provides BLAS routines for sparse matrices. In the following we are going to use those to compare their speed to using scipy's sparse routines (which are implemented directly even if scipy is linked against MKL).
End of explanation
"""


