# -*- 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()
    

[((1464, 1493), 'numpy.load', 'np.load', (['"""./AL_dist_flat.npy"""'], {}), "('./AL_dist_flat.npy')\n", (1471, 1493), True, 'import numpy as np\n'), ((1582, 1605), 'numpy.logspace', 'np.logspace', (['(-2)', '(3)', '(100)'], {}), '(-2, 3, 100)\n', (1593, 1605), True, 'import numpy as np\n'), ((1782, 1807), 'multiprocessing.Value', 'mp.Value', (['ctypes.c_int', 'n'], {}), '(ctypes.c_int, n)\n', (1790, 1807), True, 'import multiprocessing as mp\n'), ((1820, 1845), 'multiprocessing.Value', 'mp.Value', (['ctypes.c_int', 'm'], {}), '(ctypes.c_int, m)\n', (1828, 1845), True, 'import multiprocessing as mp\n'), ((1976, 2061), 'multiprocessing.Pool', 'mp.Pool', (['(20)'], {'initializer': 'parallelinit', 'initargs': '(AL_dist_flat_glo, n_glo, m_glo)'}), '(20, initializer=parallelinit, initargs=(AL_dist_flat_glo, n_glo, m_glo)\n    )\n', (1983, 2061), True, 'import multiprocessing as mp\n'), ((2071, 2082), 'time.time', 'time.time', ([], {}), '()\n', (2080, 2082), False, 'import time\n'), ((2165, 2176), 'time.time', 'time.time', ([], {}), '()\n', (2174, 2176), False, 'import time\n'), ((2208, 2244), 'numpy.save', 'np.save', (['"""./AL_results.npy"""', 'results'], {}), "('./AL_results.npy', results)\n", (2215, 2244), True, 'import numpy as np\n'), ((2654, 2680), 'matplotlib.pyplot.figure', 'plt.figure', ([], {'figsize': '(8, 6)'}), '(figsize=(8, 6))\n', (2664, 2680), True, 'import matplotlib.pyplot as plt\n'), ((2684, 2731), 'matplotlib.pyplot.plot', 'plt.plot', (['q_range', 'Pq'], {'lw': '(3)', 'color': '"""tab:orange"""'}), "(q_range, Pq, lw=3, color='tab:orange')\n", (2692, 2731), True, 'import matplotlib.pyplot as plt\n'), ((2736, 2753), 'matplotlib.pyplot.xscale', 'plt.xscale', (['"""log"""'], {}), "('log')\n", (2746, 2753), True, 'import matplotlib.pyplot as plt\n'), ((2758, 2775), 'matplotlib.pyplot.yscale', 'plt.yscale', (['"""log"""'], {}), "('log')\n", (2768, 2775), True, 'import matplotlib.pyplot as plt\n'), ((2780, 2810), 'matplotlib.pyplot.xlabel', 'plt.xlabel', (['"""$q$"""'], {'fontsize': '(15)'}), "('$q$', fontsize=15)\n", (2790, 2810), True, 'import matplotlib.pyplot as plt\n'), ((2815, 2848), 'matplotlib.pyplot.ylabel', 'plt.ylabel', (['"""$P(q)$"""'], {'fontsize': '(15)'}), "('$P(q)$', fontsize=15)\n", (2825, 2848), True, 'import matplotlib.pyplot as plt\n'), ((2853, 2871), 'matplotlib.pyplot.tight_layout', 'plt.tight_layout', ([], {}), '()\n', (2869, 2871), True, 'import matplotlib.pyplot as plt\n'), ((2876, 2945), 'matplotlib.pyplot.savefig', 'plt.savefig', (['"""./AL_form_factor_log.pdf"""'], {'dpi': '(300)', 'bbox_inches': '"""tight"""'}), "('./AL_form_factor_log.pdf', dpi=300, bbox_inches='tight')\n", (2887, 2945), True, 'import matplotlib.pyplot as plt\n'), ((2951, 2961), 'matplotlib.pyplot.show', 'plt.show', ([], {}), '()\n', (2959, 2961), True, 'import matplotlib.pyplot as plt\n'), ((618, 641), 'numpy.logspace', 'np.logspace', (['(-2)', '(3)', '(100)'], {}), '(-2, 3, 100)\n', (629, 641), True, 'import numpy as np\n'), ((660, 734), 'numpy.subtract', 'np.subtract', (['AL_dist_flat_glo_s[args[1]]', 'AL_dist_flat_glo_s[1 + args[1]:]'], {}), '(AL_dist_flat_glo_s[args[1]], AL_dist_flat_glo_s[1 + args[1]:])\n', (671, 734), True, 'import numpy as np\n'), ((1508, 1530), 'numpy.shape', 'np.shape', (['AL_dist_flat'], {}), '(AL_dist_flat)\n', (1516, 1530), True, 'import numpy as np\n'), ((1542, 1564), 'numpy.shape', 'np.shape', (['AL_dist_flat'], {}), '(AL_dist_flat)\n', (1550, 1564), True, 'import numpy as np\n'), ((1187, 1215), 'numpy.array', 'np.array', (['[cosx, cosy, cosz]'], {}), '([cosx, cosy, cosz])\n', (1195, 1215), True, 'import numpy as np\n'), ((763, 782), 'numpy.array', 'np.array', (['[1, 0, 0]'], {}), '([1, 0, 0])\n', (771, 782), True, 'import numpy as np\n'), ((817, 836), 'numpy.array', 'np.array', (['[0, 1, 0]'], {}), '([0, 1, 0])\n', (825, 836), True, 'import numpy as np\n'), ((871, 890), 'numpy.array', 'np.array', (['[0, 0, 1]'], {}), '([0, 0, 1])\n', (879, 890), True, 'import numpy as np\n'), ((2279, 2296), 'numpy.array', 'np.array', (['results'], {}), '(results)\n', (2287, 2296), True, 'import numpy as np\n')]