#!/usr/bin/env python
from multiprocessing import Lock, Process, Queue, current_process
import time
import queue # imported for using queue.Empty exception
import numpy as np
import time
from scipy._lib._util import _asarray_validated
import sys

try:
    import skcuda.magma as magma

    useScipy = False
except Exception as err:
    print("#", err)
    print("# Cannot import scikit-cuda. Fall back to scipy.linalg")
    import scipy.linalg
    useScipy = True

typedict = {'s': np.float32, 'd': np.float64, 'c': np.complex64, 'z': np.complex128}
typedict_ = {v: k for k, v in typedict.items()}


def gpu_eig(lock, a, left=False, right=True, check_finite=True, verbose=True, *args, **kwargs):
    """
        Eigenvalue solver using Magma GPU library (variants of Lapack's geev).

        Note:
        - not a generalized eigenvalue solver
    """
    if useScipy:
        return scipy.linalg.eig(a, left=left, right=right, check_finite=check_finite)
    else:
        if len(a.shape) != 2:
            raise ValueError("M needs to be a rank 2 square array for eig.")

        a = _asarray_validated(a, check_finite=check_finite)

        lock.acquire()
        #magma.magma_init()
        dtype = type(a[0, 0])
        t = typedict_[dtype]
        N = a.shape[0]

        # Set up output buffers:
        if t in ['s', 'd']:
            wr = np.zeros((N,), dtype)  # eigenvalues
            wi = np.zeros((N,), dtype)  # eigenvalues
        elif t in ['c', 'z']:
            w = np.zeros((N,), dtype)  # eigenvalues

        if left:
            vl = np.zeros((N, N), dtype)
            jobvl = 'V'
        else:
            vl = np.zeros((1, 1), dtype)
            jobvl = 'N'
        if right:
            vr = np.zeros((N, N), dtype)
            jobvr = 'V'
        else:
            vr = np.zeros((1, 1), dtype)
            jobvr = 'N'

        # Set up workspace:
        if t == 's':
            nb = magma.magma_get_sgeqrf_nb(N, N)
        if t == 'd':
            nb = magma.magma_get_dgeqrf_nb(N, N)
        if t == 'c':
            nb = magma.magma_get_cgeqrf_nb(N, N)
        if t == 'z':
            nb = magma.magma_get_zgeqrf_nb(N, N)

        lwork = N * (1 + 2 * nb)
        work = np.zeros((lwork,), dtype)
        if t in ['c', 'z']:
            rwork = np.zeros((2 * N,), dtype)

        # Compute:
        gpu_time = time.time();
        if t == 's':
            status = magma.magma_sgeev(jobvl, jobvr, N, a.ctypes.data, N,
                                       wr.ctypes.data, wi.ctypes.data,
                                       vl.ctypes.data, N, vr.ctypes.data, N,
                                       work.ctypes.data, lwork)
        if t == 'd':
            status = magma.magma_dgeev(jobvl, jobvr, N, a.ctypes.data, N,
                                       wr.ctypes.data, wi.ctypes.data,
                                       vl.ctypes.data, N, vr.ctypes.data, N,
                                       work.ctypes.data, lwork)
        if t == 'c':
            status = magma.magma_cgeev(jobvl, jobvr, N, a.ctypes.data, N,
                                       w.ctypes.data, vl.ctypes.data, N, vr.ctypes.data, N,
                                       work.ctypes.data, lwork, rwork.ctypes.data)
        if t == 'z':
            status = magma.magma_zgeev(jobvl, jobvr, N, a.ctypes.data, N,
                                       w.ctypes.data, vl.ctypes.data, N, vr.ctypes.data, N,
                                       work.ctypes.data, lwork, rwork.ctypes.data)
        gpu_time = time.time() - gpu_time;

        if verbose:
            print("Time for eig: ", gpu_time)

        if t in ['s', 'd']:
            w_gpu = wr + 1j * wi
        else:
            w_gpu = w

        lock.release()
        #magma.magma_finalize()
        return w_gpu, vr

def get_eigsys(lock, mat, tasks_to_accomplish, tasks_that_are_done):
    while True:
        try:
            '''
                try to get task from the queue. get_nowait() function will
                raise queue.Empty exception if the queue is empty.
                queue(False) function would do the same task also.
            '''
            task_ampl = tasks_to_accomplish.get_nowait()
            w, v = gpu_eig(lock, mat, left=False, right=True)

        except queue.Empty:

            break
        else:
            '''
                if no exception has been raised, add the task completion
                message to task_that_are_done queue
            '''
            tasks_that_are_done.put((task_ampl,w,v))
            time.sleep(.5)
    return True


def main():
    size = 3
    ampls = np.linspace(10.0, 30.0, 2)
    number_of_task = len(ampls)
    number_of_processes = number_of_task
    tasks_to_accomplish = Queue()
    tasks_that_are_done = Queue()
    processes = []
    magma.magma_init()

    for h in ampls:
        tasks_to_accomplish.put(h)

    # creating processes
    lock = Lock()
    for h in ampls:
        #Generate Floquet Matrix here. Random for now
        floquet_matrix = np.random.random((size,size)) + (1j) * np.random.random((size,size))
        p = Process(target=get_eigsys, args=(lock, floquet_matrix, tasks_to_accomplish, tasks_that_are_done))
        processes.append(p)
        p.start()

    # completing process
    for p in processes:
        p.join()

    magma.magma_finalize()

    # print the output
    while not tasks_that_are_done.empty():
        print(tasks_that_are_done.get())

    return True


if __name__ == '__main__':
    main()