# -*- coding: utf-8 -*-
"""
Created on Mon Sep  7 10:59:00 2020

@author: user
"""

import numpy as np
import multiprocessing as mp
import matplotlib.pyplot as plt
import time
import itertools
import ctypes

def formfactor(args):
    # with AL_dist_flat_glo.get_lock:
    AL_dist_flat_glo_r = np.frombuffer(AL_dist_flat_glo.get_obj())
    AL_dist_flat_glo_s = AL_dist_flat_glo_r.reshape((n_glo.value,m_glo.value))
    # ffq = np.sum(np.cos(np.dot(np.logspace(-2,3,100)[args[0]]*np.array([1,0,0]), 
    #                            np.subtract(AL_dist_flat_glo_s[args[1]], AL_dist_flat_glo_s[1+args[1]:]).T)))
    qr = np.logspace(-2,3,100)[args[0]]
    rvec = np.subtract(AL_dist_flat_glo_s[args[1]], AL_dist_flat_glo_s[1+args[1]:]).T
    cosx = np.cos(np.dot(qr*np.array([1,0,0]), rvec))
    cosy = np.cos(np.dot(qr*np.array([0,1,0]), rvec))
    cosz = np.cos(np.dot(qr*np.array([0,0,1]), rvec))
    # cosxy = np.cos(np.dot(qr*np.array([0.707,0.707,0]), rvec))
    # cosyz = np.cos(np.dot(qr*np.array([0,0.707,0.707]), rvec))
    # cosxz = np.cos(np.dot(qr*np.array([0.707,0,0.707]), rvec))
    # cosxyz = np.cos(np.dot(qr*np.array([0.577,0.577,0.577]), rvec))
    ffq = np.sum(np.mean(np.array([cosx, cosy, cosz]), axis=0))
    return ffq

def parallelinit(AL_dist_flat_glo_, n_glo_, m_glo_):
    global AL_dist_flat_glo, n_glo, m_glo
    AL_dist_flat_glo = AL_dist_flat_glo_
    n_glo = n_glo_
    m_glo = m_glo_

if __name__ == '__main__': 
    AL_dist_flat = np.load(r'./AL_dist_flat.npy')
    
    n = np.shape(AL_dist_flat)[0]
    m = np.shape(AL_dist_flat)[1]
    q_range = np.logspace(-2,3,100)
    # r_x = np.array([1, 0, 0])
    
    # q_range_glo = mp.Array(ctypes.c_double, q_range)
    AL_dist_flat_glo =  mp.Array(ctypes.c_double, AL_dist_flat.flatten())
    n_glo = mp.Value(ctypes.c_int, n)
    m_glo = mp.Value(ctypes.c_int, m)
    # r_x_glo = mp.Array(ctypes.c_double, r_x)
    
    paramlist = list(itertools.product(range(100), range(n)))
    
    pool = mp.Pool(20, initializer=parallelinit, initargs=(AL_dist_flat_glo, n_glo, m_glo))
    
    t1 = time.time()
    
    results = pool.map(formfactor, paramlist)
    pool.close()
    
    t2 = time.time()
    
    print(t2-t1)
    
    np.save(r'./AL_results.npy', results)
    
    Pq = 2*np.divide(np.sum(np.array(results).reshape(100, n), axis=1), n)
    
    # fig = plt.figure(figsize=(8,6))
    # plt.plot(q_range, Pq, lw=3, color='tab:orange')
    # plt.xscale('log')
    # plt.xlabel('$q$', fontsize=15)
    # plt.ylabel('$P(q)$', fontsize=15)
    # plt.tight_layout()
    # plt.savefig(r'./AL_form_factor.pdf', dpi=300, bbox_inches='tight')
    # plt.show()
    
    fig = plt.figure(figsize=(8,6))
    plt.plot(q_range, Pq, lw=3, color='tab:orange')
    plt.xscale('log')
    plt.yscale('log')
    plt.xlabel('$q$', fontsize=15)
    plt.ylabel('$P(q)$', fontsize=15)
    plt.tight_layout()
    plt.savefig(r'./AL_form_factor_log.pdf', dpi=300, bbox_inches='tight')
    plt.show()
    

[((40, 19, 40, 49), 'numpy.load', 'np.load', ({(40, 27, 40, 48): '"""./AL_dist_flat.npy"""'}, {}), "('./AL_dist_flat.npy')", True, 'import numpy as np\n'), ((44, 14, 44, 35), 'numpy.logspace', 'np.logspace', ({(44, 26, 44, 28): '-2', (44, 29, 44, 30): '3', (44, 31, 44, 34): '100'}, {}), '(-2, 3, 100)', True, 'import numpy as np\n'), ((49, 12, 49, 37), 'multiprocessing.Value', 'mp.Value', ({(49, 21, 49, 33): 'ctypes.c_int', (49, 35, 49, 36): 'n'}, {}), '(ctypes.c_int, n)', True, 'import multiprocessing as mp\n'), ((50, 12, 50, 37), 'multiprocessing.Value', 'mp.Value', ({(50, 21, 50, 33): 'ctypes.c_int', (50, 35, 50, 36): 'm'}, {}), '(ctypes.c_int, m)', True, 'import multiprocessing as mp\n'), ((55, 11, 55, 91), 'multiprocessing.Pool', 'mp.Pool', (), '', True, 'import multiprocessing as mp\n'), ((57, 9, 57, 20), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((62, 9, 62, 20), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((66, 4, 66, 41), 'numpy.save', 'np.save', ({(66, 12, 66, 31): '"""./AL_results.npy"""', (66, 33, 66, 40): 'results'}, {}), "('./AL_results.npy', results)", True, 'import numpy as np\n'), ((79, 10, 79, 35), 'matplotlib.pyplot.figure', 'plt.figure', (), '', True, 'import matplotlib.pyplot as plt\n'), ((80, 4, 80, 51), 'matplotlib.pyplot.plot', 'plt.plot', (), '', True, 'import matplotlib.pyplot as plt\n'), ((81, 4, 81, 21), 'matplotlib.pyplot.xscale', 'plt.xscale', ({(81, 15, 81, 20): '"""log"""'}, {}), "('log')", True, 'import matplotlib.pyplot as plt\n'), ((82, 4, 82, 21), 'matplotlib.pyplot.yscale', 'plt.yscale', ({(82, 15, 82, 20): '"""log"""'}, {}), "('log')", True, 'import matplotlib.pyplot as plt\n'), ((83, 4, 83, 34), 'matplotlib.pyplot.xlabel', 'plt.xlabel', (), '', True, 'import matplotlib.pyplot as plt\n'), ((84, 4, 84, 37), 'matplotlib.pyplot.ylabel', 'plt.ylabel', (), '', True, 'import matplotlib.pyplot as plt\n'), ((85, 4, 85, 22), 'matplotlib.pyplot.tight_layout', 'plt.tight_layout', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((86, 4, 86, 74), 'matplotlib.pyplot.savefig', 'plt.savefig', (), '', True, 'import matplotlib.pyplot as plt\n'), ((87, 4, 87, 14), 'matplotlib.pyplot.show', 'plt.show', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((21, 9, 21, 30), 'numpy.logspace', 'np.logspace', ({(21, 21, 21, 23): '(-2)', (21, 24, 21, 25): '(3)', (21, 26, 21, 29): '(100)'}, {}), '(-2, 3, 100)', True, 'import numpy as np\n'), ((22, 11, 22, 83), 'numpy.subtract', 'np.subtract', ({(22, 23, 22, 50): 'AL_dist_flat_glo_s[args[1]]', (22, 52, 22, 82): 'AL_dist_flat_glo_s[1 + args[1]:]'}, {}), '(AL_dist_flat_glo_s[args[1]], AL_dist_flat_glo_s[1 + args[1]:])', True, 'import numpy as np\n'), ((42, 8, 42, 30), 'numpy.shape', 'np.shape', ({(42, 17, 42, 29): 'AL_dist_flat'}, {}), '(AL_dist_flat)', True, 'import numpy as np\n'), ((43, 8, 43, 30), 'numpy.shape', 'np.shape', ({(43, 17, 43, 29): 'AL_dist_flat'}, {}), '(AL_dist_flat)', True, 'import numpy as np\n'), ((30, 25, 30, 53), 'numpy.array', 'np.array', ({(30, 34, 30, 52): '[cosx, cosy, cosz]'}, {}), '([cosx, cosy, cosz])', True, 'import numpy as np\n'), ((23, 28, 23, 45), 'numpy.array', 'np.array', ({(23, 37, 23, 44): '[1, 0, 0]'}, {}), '([1, 0, 0])', True, 'import numpy as np\n'), ((24, 28, 24, 45), 'numpy.array', 'np.array', ({(24, 37, 24, 44): '[0, 1, 0]'}, {}), '([0, 1, 0])', True, 'import numpy as np\n'), ((25, 28, 25, 45), 'numpy.array', 'np.array', ({(25, 37, 25, 44): '[0, 0, 1]'}, {}), '([0, 0, 1])', True, 'import numpy as np\n'), ((68, 28, 68, 45), 'numpy.array', 'np.array', ({(68, 37, 68, 44): 'results'}, {}), '(results)', True, 'import numpy as np\n')]