# -*- 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, 'numpy.load', 'np.load', (['"""./AL_dist_flat.npy"""'], {}), True, 'import numpy as np\n'), (44, 'numpy.logspace', 'np.logspace', (['(-2)', '(3)', '(100)'], {}), True, 'import numpy as np\n'), (49, 'multiprocessing.Value', 'mp.Value', (['ctypes.c_int', 'n'], {}), True, 'import multiprocessing as mp\n'), (50, 'multiprocessing.Value', 'mp.Value', (['ctypes.c_int', 'm'], {}), True, 'import multiprocessing as mp\n'), (55, 'multiprocessing.Pool', 'mp.Pool', (['(20)'], {'initializer': 'parallelinit', 'initargs': '(AL_dist_flat_glo, n_glo, m_glo)'}), True, 'import multiprocessing as mp\n'), (57, 'time.time', 'time.time', ([], {}), False, 'import time\n'), (62, 'time.time', 'time.time', ([], {}), False, 'import time\n'), (66, 'numpy.save', 'np.save', (['"""./AL_results.npy"""', 'results'], {}), True, 'import numpy as np\n'), (79, 'matplotlib.pyplot.figure', 'plt.figure', ([], {'figsize': '(8, 6)'}), True, 'import matplotlib.pyplot as plt\n'), (80, 'matplotlib.pyplot.plot', 'plt.plot', (['q_range', 'Pq'], {'lw': '(3)', 'color': '"""tab:orange"""'}), True, 'import matplotlib.pyplot as plt\n'), (81, 'matplotlib.pyplot.xscale', 'plt.xscale', (['"""log"""'], {}), True, 'import matplotlib.pyplot as plt\n'), (82, 'matplotlib.pyplot.yscale', 'plt.yscale', (['"""log"""'], {}), True, 'import matplotlib.pyplot as plt\n'), (83, 'matplotlib.pyplot.xlabel', 'plt.xlabel', (['"""$q$"""'], {'fontsize': '(15)'}), True, 'import matplotlib.pyplot as plt\n'), (84, 'matplotlib.pyplot.ylabel', 'plt.ylabel', (['"""$P(q)$"""'], {'fontsize': '(15)'}), True, 'import matplotlib.pyplot as plt\n'), (85, 'matplotlib.pyplot.tight_layout', 'plt.tight_layout', ([], {}), True, 'import matplotlib.pyplot as plt\n'), (86, 'matplotlib.pyplot.savefig', 'plt.savefig', (['"""./AL_form_factor_log.pdf"""'], {'dpi': '(300)', 'bbox_inches': '"""tight"""'}), True, 'import matplotlib.pyplot as plt\n'), (87, 'matplotlib.pyplot.show', 'plt.show', ([], {}), True, 'import matplotlib.pyplot as plt\n'), (21, 'numpy.logspace', 'np.logspace', (['(-2)', '(3)', '(100)'], {}), True, 'import numpy as np\n'), (22, 'numpy.subtract', 'np.subtract', (['AL_dist_flat_glo_s[args[1]]', 'AL_dist_flat_glo_s[1 + args[1]:]'], {}), True, 'import numpy as np\n'), (42, 'numpy.shape', 'np.shape', (['AL_dist_flat'], {}), True, 'import numpy as np\n'), (43, 'numpy.shape', 'np.shape', (['AL_dist_flat'], {}), True, 'import numpy as np\n'), (30, 'numpy.array', 'np.array', (['[cosx, cosy, cosz]'], {}), True, 'import numpy as np\n'), (23, 'numpy.array', 'np.array', (['[1, 0, 0]'], {}), True, 'import numpy as np\n'), (24, 'numpy.array', 'np.array', (['[0, 1, 0]'], {}), True, 'import numpy as np\n'), (25, 'numpy.array', 'np.array', (['[0, 0, 1]'], {}), True, 'import numpy as np\n'), (68, 'numpy.array', 'np.array', (['results'], {}), True, 'import numpy as np\n')]