"""This module contains functionality to prepare lookups for the particle 
coupling coefficients, which allows to efficiently treat large numbers of 
particles. """

import sys
from tqdm import tqdm
import numpy as np
import scipy.special
import smuthi.fields
import smuthi.utility.math as sf
import smuthi.utility.cuda as cu
import smuthi.fields.transformations as trf
import smuthi.layers as lay
import smuthi.linearsystem.particlecoupling.prepare_lookup_cuda as cusrc

def volumetric_coupling_lookup_table(vacuum_wavelength, particle_list, layer_system, k_parallel='default', 
                                     resolution=None):
    """Prepare Sommerfeld integral lookup table to allow for a fast calculation of the coupling matrix by interpolation.
    This function is called when not all particles are on the same z-position.
    
    Args:
        vacuum_wavelength (float):  Vacuum wavelength in length units
        particle_list (list):       List of particle objects
        layer_system (smuthi.layers.LayerSystem):    Stratified medium
        k_parallel (numpy.ndarray or str):           In-plane wavenumber for Sommerfeld integrals.
                                                     If 'default', smuthi.fields.default_Sommerfeld_k_parallel_array
        resolution (float): Spatial resolution of lookup table in length units. (default: vacuum_wavelength / 100)        
                            Smaller means more accurate but higher memory footprint 
                            
    Returns:
        (tuple): tuple containing:

            w_pl (ndarray):  Coupling lookup for z1 + z2, indices are [rho, z, n1, n2]. Includes layer mediated coupling.
            w_mn (ndarray):  Coupling lookup for z1 + z2, indices are [rho, z, n1, n2]. Includes layer mediated and 
                             direct coupling.
            rho_array (ndarray):    Values for the radial distance considered for the lookup (starting from negative 
                                    numbers to allow for simpler cubic interpolation without distinction of cases 
                                    for lookup edges
            sz_array (ndarray):     Values for the sum of z-coordinates (z1 + z2) considered for the lookup
            dz_array (ndarray):     Values for the difference of z-coordinates (z1 - z2) considered for the lookup 
    """
    sys.stdout.write('Prepare 3D particle coupling lookup:\n')
    sys.stdout.flush()

    if resolution is None:
        resolution = vacuum_wavelength / 100
        sys.stdout.write('Setting lookup resolution to %f\n'%resolution)
        sys.stdout.flush()

    l_max = max([particle.l_max for particle in particle_list])
    m_max = max([particle.m_max for particle in particle_list])
    blocksize = smuthi.fields.blocksize(l_max, m_max)
    
    particle_x_array = np.array([particle.position[0] for particle in particle_list])
    particle_y_array = np.array([particle.position[1] for particle in particle_list])
    particle_z_array = np.array([particle.position[2] for particle in particle_list])
    particle_rho_array = np.sqrt((particle_x_array[:, None] - particle_x_array[None, :]) ** 2 
                                 + (particle_y_array[:, None] - particle_y_array[None, :]) ** 2)
    
    dz_min = particle_z_array.min() - particle_z_array.max()
    dz_max = particle_z_array.max() - particle_z_array.min()
    sz_min = 2 * particle_z_array.min()
    sz_max = 2 * particle_z_array.max()
    
    rho_array = np.arange(- 3 * resolution, particle_rho_array.max() + 3 * resolution, resolution)
    sz_array = np.arange(sz_min - 3 * resolution, sz_max + 3 * resolution, resolution)
    dz_array = np.arange(dz_min - 3 * resolution, dz_max + 3 * resolution, resolution)
    
    len_rho = len(rho_array)
    len_sz = len(sz_array)
    len_dz = len(dz_array)
    assert len_sz == len_dz
    
    i_s = layer_system.layer_number(particle_list[0].position[2])
    k_is = layer_system.wavenumber(i_s, vacuum_wavelength)
    z_is = layer_system.reference_z(i_s)

    wr_pl = np.zeros((len_rho, len_dz, blocksize, blocksize), dtype=np.complex64)
    wr_mn = np.zeros((len_rho, len_dz, blocksize, blocksize), dtype=np.complex64)

    # direct -----------------------------------------------------------------------------------------------------------
    w = np.zeros((len_rho, len_dz, blocksize, blocksize), dtype=np.complex64)
    sys.stdout.write('Lookup table memory footprint: ' + size_format(2 * w.nbytes) + '\n')
    sys.stdout.flush()

    r_array = np.sqrt(dz_array[None, :]**2 + rho_array[:, None]**2)
    r_array[r_array==0] = 1e-20
    ct = dz_array[None, :] / r_array
    st = rho_array[:, None] / r_array
    legendre, _, _ = sf.legendre_normalized(ct, st, 2 * l_max)
    
    bessel_h = []
    for dm in tqdm(range(2 * l_max + 1), desc='Spherical Hankel lookup   ', file=sys.stdout,
                   bar_format='{l_bar}{bar}| elapsed: {elapsed} remaining: {remaining}'):
        bessel_h.append(sf.spherical_hankel(dm, k_is * r_array))
    
    pbar = tqdm(total=blocksize**2, 
                desc='Direct coupling           ', 
                file=sys.stdout,
                bar_format='{l_bar}{bar}| elapsed: {elapsed} remaining: {remaining}')
            
    for m1 in range(-m_max, m_max+1):
        for m2 in range(-m_max, m_max+1):
            for l1 in range(max(1, abs(m1)), l_max + 1):
                for l2 in range(max(1, abs(m2)), l_max + 1):
                    A = np.zeros((len_rho, len_dz), dtype=complex)
                    B = np.zeros((len_rho, len_dz), dtype=complex)
                    for ld in range(max(abs(l1 - l2), abs(m1 - m2)), l1 + l2 + 1):  # if ld<abs(m1-m2) then P=0
                        a5, b5 = trf.ab5_coefficients(l2, m2, l1, m1, ld)    # remember that w = A.T
                        A += a5 * bessel_h[ld] * legendre[ld, abs(m1 - m2)]   # remember that w = A.T
                        B += b5 * bessel_h[ld] * legendre[ld, abs(m1 - m2)]   # remember that w = A.T
                    for tau1 in range(2):
                        n1 = smuthi.fields.multi_to_single_index(tau1, l1, m1, l_max, m_max)
                        for tau2 in range(2):
                            n2 = smuthi.fields.multi_to_single_index(tau2, l2, m2, l_max, m_max)
                            if tau1 == tau2:
                                w[:, :, n1, n2] = A
                            else:
                                w[:, :, n1, n2] = B
                            pbar.update()
    pbar.close()

    # switch off direct coupling contribution near rho=0:
    w[rho_array < particle_rho_array[~np.eye(particle_rho_array.shape[0],dtype=bool)].min() / 2, :, :, :] = 0  


    # Do not calculate the layer mediated lookup if there is no meaningful layer system!
    if layer_system.is_degenerate() == False:
        # layer mediated ---------------------------------------------------------------------------------------------------
        sys.stdout.write('Layer mediated coupling   : ...')
        sys.stdout.flush()
        if type(k_parallel) == str and k_parallel == 'default':
            k_parallel = smuthi.fields.default_Sommerfeld_k_parallel_array
        kz_is = smuthi.fields.k_z(k_parallel=k_parallel, k=k_is)
        len_kp = len(k_parallel)
    
        # phase factors
        epljksz = np.exp(1j * kz_is[None, :] * (sz_array[:, None] - 2 * z_is))  # z, k
        emnjksz = np.exp(- 1j * kz_is[None, :] * (sz_array[:, None] - 2 * z_is))
        epljkdz = np.exp(1j * kz_is[None, :] * dz_array[:, None])
        emnjkdz = np.exp(- 1j * kz_is[None, :] * dz_array[:, None])
           
        # layer response
        L = np.zeros((2, 2, 2, len_kp), dtype=complex)  # pol, pl/mn1, pl/mn2, kp
        for pol in range(2):
            L[pol, :, :, :] = lay.layersystem_response_matrix(pol, layer_system.thicknesses,
                                                              layer_system.refractive_indices, k_parallel,
                                                              smuthi.fields.angular_frequency(vacuum_wavelength), i_s, i_s)
       
        # transformation coefficients
        B_dag = np.zeros((2, 2, blocksize, len_kp), dtype=complex)  # pol, pl/mn, n, kp
        B = np.zeros((2, 2, blocksize, len_kp), dtype=complex)  # pol, pl/mn, n, kp
        ct_k = kz_is / k_is
        st_k = k_parallel / k_is
        _, pilm_pl, taulm_pl = sf.legendre_normalized(ct_k, st_k, l_max)
        _, pilm_mn, taulm_mn = sf.legendre_normalized(-ct_k, st_k, l_max)
        
        m_list = [None for i in range(blocksize)]
        for tau in range(2):
            for m in range(-m_max, m_max + 1):
                for l in range(max(1, abs(m)), l_max + 1):
                    n = smuthi.fields.multi_to_single_index(tau, l, m, l_max, m_max)
                    m_list[n] = m
                    for pol in range(2):
                        B_dag[pol, 0, n, :] = trf.transformation_coefficients_vwf(tau, l, m, pol, pilm_list=pilm_pl,
                                                                                        taulm_list=taulm_pl, dagger=True)
                        B_dag[pol, 1, n, :] = trf.transformation_coefficients_vwf(tau, l, m, pol, pilm_list=pilm_mn,
                                                                                        taulm_list=taulm_mn, dagger=True)
                        B[pol, 0, n, :] = trf.transformation_coefficients_vwf(tau, l, m, pol, pilm_list=pilm_pl,
                                                                                 taulm_list=taulm_pl, dagger=False)
                        B[pol, 1, n, :] = trf.transformation_coefficients_vwf(tau, l, m, pol, pilm_list=pilm_mn,
                                                                                 taulm_list=taulm_mn, dagger=False)
        
        # pairs of (n1, n2), listed by abs(m1-m2)
        n1n2_combinations = [[] for dm in range(2*m_max+1)]
        for n1 in range(blocksize):
            m1 = m_list[n1]
            for n2 in range(blocksize):
                m2 = m_list[n2]
                n1n2_combinations[abs(m1-m2)].append((n1,n2))
                       

        
        dkp = np.diff(k_parallel)
        if cu.use_gpu:
            re_dkp_d = cu.gpuarray.to_gpu(np.float32(dkp.real))
            im_dkp_d = cu.gpuarray.to_gpu(np.float32(dkp.imag))
            kernel_source_code = cusrc.volume_lookup_assembly_code %(blocksize, len_rho, len_sz, len_kp)
            helper_function = cu.SourceModule(kernel_source_code).get_function("helper")
            cuda_blocksize = 128
            cuda_gridsize = (len_rho * len_sz + cuda_blocksize - 1) // cuda_blocksize
    
            re_dwr_d = cu.gpuarray.to_gpu(np.zeros((len_rho, len_sz), dtype=np.float32))
            im_dwr_d = cu.gpuarray.to_gpu(np.zeros((len_rho, len_sz), dtype=np.float32))
    
        pbar = tqdm(total=blocksize**2, 
                    desc='Layer mediated coupling   ', 
                    file=sys.stdout,
                    bar_format='{l_bar}{bar}| elapsed: {elapsed} remaining: {remaining}')
                
        for dm in range(2*m_max+1):
            bessel = scipy.special.jv(dm, (k_parallel[None,:]*rho_array[:,None]))
            besjac = bessel * (k_parallel / (kz_is * k_is))[None,:]
    
            for n1n2 in n1n2_combinations[dm]:
                n1 = n1n2[0]
                m1 = m_list[n1]
                n2 = n1n2[1]
                m2 = m_list[n2]
        
                belbee_pl = np.zeros((len_dz, len_kp), dtype=complex)
                belbee_mn = np.zeros((len_dz, len_kp), dtype=complex)
                for pol in range(2):
                    belbee_pl += ((L[pol, 0, 1, :] * B_dag[pol, 0, n1, :] * B[pol, 1, n2, :])[None, :] * epljksz
                                  + (L[pol, 1, 0, :] * B_dag[pol, 1, n1, :] * B[pol, 0, n2, :])[None, :] * emnjksz)
                    belbee_mn += ((L[pol, 0, 0, :] * B_dag[pol, 0, n1, :] * B[pol, 0, n2, :])[None, :] * epljkdz
                                  + (L[pol, 1, 1, :] * B_dag[pol, 1, n1, :] * B[pol, 1, n2, :])[None, :] * emnjkdz)
                
                if cu.use_gpu:
                    re_belbee_pl_d = cu.gpuarray.to_gpu(np.float32(belbee_pl[None, :, :].real))
                    im_belbee_pl_d = cu.gpuarray.to_gpu(np.float32(belbee_pl[None, :, :].imag))
                    re_belbee_mn_d = cu.gpuarray.to_gpu(np.float32(belbee_mn[None, :, :].real))
                    im_belbee_mn_d = cu.gpuarray.to_gpu(np.float32(belbee_mn[None, :, :].imag))
    
                    re_besjac_d = cu.gpuarray.to_gpu(np.float32(besjac[:, None, :].real))
                    im_besjac_d = cu.gpuarray.to_gpu(np.float32(besjac[:, None, :].imag))
                    helper_function(re_besjac_d.gpudata, im_besjac_d.gpudata, re_belbee_pl_d.gpudata,
                                    im_belbee_pl_d.gpudata, re_dkp_d.gpudata, im_dkp_d.gpudata, re_dwr_d.gpudata, 
                                    im_dwr_d.gpudata, block=(cuda_blocksize, 1, 1), grid=(cuda_gridsize, 1))
                    wr_pl[:, :, n1, n2] = 4 * (1j)**abs(m2 - m1) * (re_dwr_d.get() + 1j * im_dwr_d.get()) 
                    
                    helper_function(re_besjac_d.gpudata, im_besjac_d.gpudata, re_belbee_mn_d.gpudata, 
                                    im_belbee_mn_d.gpudata, re_dkp_d.gpudata, im_dkp_d.gpudata, re_dwr_d.gpudata, 
                                    im_dwr_d.gpudata, block=(cuda_blocksize, 1, 1), grid=(cuda_gridsize, 1))
                    wr_mn[:, :, n1, n2] = 4 * (1j)**abs(m2 - m1) * (re_dwr_d.get() + 1j * im_dwr_d.get()) 
                else:
                    integrand = besjac[:, None, :] * belbee_pl[None, :, :] 
                    wr_pl[:, :, n1, n2] = 2 * (1j)**abs(m2 - m1) * ((integrand[:, :, :-1] + integrand[:, :, 1:]) 
                                                                    * dkp[None, None, :]).sum(axis=-1)   # trapezoidal rule
                    
                    integrand = besjac[:, None, :] * belbee_mn[None, :, :]
                    wr_mn[:, :, n1, n2] = 2 * (1j)**abs(m2 - m1) * ((integrand[:, :, :-1] + integrand[:, :, 1:])
                                                                    * dkp[None, None, :]).sum(axis=-1)
                pbar.update()
        pbar.close()
    else:
        sys.stdout.write('Omitting calculations of layer mediated coupling since no difference in layers was detected. \n')
        sys.stdout.flush()
    
    return wr_pl, w + wr_mn, rho_array, sz_array, dz_array


def radial_coupling_lookup_table(vacuum_wavelength, particle_list, layer_system, k_parallel='default', resolution=None):
    """Prepare Sommerfeld integral lookup table to allow for a fast calculation of the coupling matrix by interpolation.
    This function is called when all particles are on the same z-position.
    
    Args:
        vacuum_wavelength (float):  Vacuum wavelength in length units
        particle_list (list):       List of particle objects
        layer_system (smuthi.layers.LayerSystem):    Stratified medium
        k_parallel (numpy.ndarray or str):           In-plane wavenumber for Sommerfeld integrals.
                                                     If 'default', smuthi.fields.default_Sommerfeld_k_parallel_array
        resolution (float): Spatial resolution of lookup table in length units. (default: vacuum_wavelength / 100)       
                            Smaller means more accurate but higher memory footprint 
                            
    Returns:
        (tuple) tuple containing:
        
            lookup_table (ndarray):  Coupling lookup, indices are [rho, n1, n2].
            rho_array (ndarray):     Values for the radial distance considered for the lookup (starting from negative 
                                     numbers to allow for simpler cubic interpolation without distinction of cases 
                                     at rho=0)
    """

    sys.stdout.write('Prepare radial particle coupling lookup:\n')
    sys.stdout.flush()

    if resolution is None:
        resolution = vacuum_wavelength / 100
        sys.stdout.write('Setting lookup resolution to %f\n'%resolution)
        sys.stdout.flush()
   
    l_max = max([particle.l_max for particle in particle_list])
    m_max = max([particle.m_max for particle in particle_list])
    blocksize = smuthi.fields.blocksize(l_max, m_max)
    
    x_array = np.array([particle.position[0] for particle in particle_list])
    y_array = np.array([particle.position[1] for particle in particle_list])
    rho_array = np.sqrt((x_array[:, None] - x_array[None, :]) ** 2 + (y_array[:, None] - y_array[None, :]) ** 2)

    radial_distance_array = np.arange(- 3 * resolution, rho_array.max() + 3 * resolution, resolution)
    
    z = particle_list[0].position[2]
    i_s = layer_system.layer_number(z)
    k_is = layer_system.wavenumber(i_s, vacuum_wavelength)
    dz = z - layer_system.reference_z(i_s)
    
    len_rho = len(radial_distance_array)
        
    # direct -----------------------------------------------------------------------------------------------------------
    w = np.zeros((len_rho, blocksize, blocksize), dtype=np.complex64)
    sys.stdout.write('Memory footprint: ' + size_format(w.nbytes) + '\n')
    sys.stdout.flush()

    ct = np.array([0.0])
    st = np.array([1.0])
    bessel_h = []
    for n in range(2* l_max + 1):
        bessel_h.append(sf.spherical_hankel(n, k_is * radial_distance_array))
        bessel_h[-1][radial_distance_array <= 0] = np.nan
        
    legendre, _, _ = sf.legendre_normalized(ct, st, 2 * l_max)

    pbar = tqdm(total=blocksize**2, 
                desc='Direct coupling           ', 
                file=sys.stdout,
                bar_format='{l_bar}{bar}| elapsed: {elapsed} remaining: {remaining}')

    for m1 in range(-m_max, m_max+1):
        for m2 in range(-m_max, m_max+1):
            for l1 in range(max(1, abs(m1)), l_max + 1):
                for l2 in range(max(1, abs(m2)), l_max + 1):
                    A = np.zeros(len_rho, dtype=complex)
                    B = np.zeros(len_rho, dtype=complex)
                    for ld in range(max(abs(l1 - l2), abs(m1 - m2)), l1 + l2 + 1):  # if ld<abs(m1-m2) then P=0
                        a5, b5 = trf.ab5_coefficients(l2, m2, l1, m1, ld)
                        A = A + a5 * bessel_h[ld] * legendre[ld, abs(m1 - m2)]
                        B = B + b5 * bessel_h[ld] * legendre[ld, abs(m1 - m2)]
                    for tau1 in range(2):
                        n1 = smuthi.fields.multi_to_single_index(tau1, l1, m1, l_max, m_max)
                        for tau2 in range(2):
                            n2 = smuthi.fields.multi_to_single_index(tau2, l2, m2, l_max, m_max)
                            if tau1 == tau2:
                                w[:, n1, n2] = A  # remember that w = A.T
                            else:
                                w[:, n1, n2] = B
                            pbar.update()
    pbar.close()
    close_to_zero = radial_distance_array < rho_array[~np.eye(rho_array.shape[0],dtype=bool)].min() / 2 
    w[close_to_zero, :, :] = 0  # switch off direct coupling contribution near rho=0


    # Only calculate the layer mediated lookup if there is a meaningful layer system!
    if layer_system.is_degenerate() != True:
        # layer mediated ---------------------------------------------------------------------------------------------------
        sys.stdout.write('Layer mediated coupling   : ...')
        sys.stdout.flush()
        
        if type(k_parallel) == str and k_parallel == 'default':
            k_parallel = smuthi.fields.default_Sommerfeld_k_parallel_array
    
        kz_is = smuthi.fields.k_z(k_parallel=k_parallel, k=k_is)
        len_kp = len(k_parallel)
    
        # phase factors
        epl2jkz = np.exp(2j * kz_is * dz)
        emn2jkz = np.exp(-2j * kz_is * dz)
           
        # layer response
        L = np.zeros((2,2,2,len_kp), dtype=complex)  # pol, pl/mn1, pl/mn2, kp
        for pol in range(2):
            L[pol,:,:,:] = lay.layersystem_response_matrix(pol, layer_system.thicknesses,
                                                           layer_system.refractive_indices, k_parallel,
                                                           smuthi.fields.angular_frequency(vacuum_wavelength), i_s, i_s)
       
        # transformation coefficients
        B_dag = np.zeros((2, 2, blocksize, len_kp), dtype=complex)  # pol, pl/mn, n, kp
        B = np.zeros((2, 2, blocksize, len_kp), dtype=complex)  # pol, pl/mn, n, kp
        ct = kz_is / k_is
        st = k_parallel / k_is
        _, pilm_pl, taulm_pl = sf.legendre_normalized(ct, st, l_max)
        _, pilm_mn, taulm_mn = sf.legendre_normalized(-ct, st, l_max)
    
        m_list = [None for n in range(blocksize)]
        for tau in range(2):
            for m in range(-m_max, m_max + 1):
                for l in range(max(1, abs(m)), l_max + 1):
                    n = smuthi.fields.multi_to_single_index(tau, l, m, l_max, m_max)
                    m_list[n] = m
                    for pol in range(2):
                        B_dag[pol,0,n,:] = trf.transformation_coefficients_vwf(tau, l, m, pol, pilm_list=pilm_pl,
                                                                               taulm_list=taulm_pl, dagger=True)
                        B_dag[pol,1,n,:] = trf.transformation_coefficients_vwf(tau, l, m, pol, pilm_list=pilm_mn,
                                                                               taulm_list=taulm_mn, dagger=True)
                        B[pol,0,n,:] = trf.transformation_coefficients_vwf(tau, l, m, pol, pilm_list=pilm_pl,
                                                                           taulm_list=taulm_pl, dagger=False)
                        B[pol,1,n,:] = trf.transformation_coefficients_vwf(tau, l, m, pol, pilm_list=pilm_mn,
                                                                           taulm_list=taulm_mn, dagger=False)
        
        # pairs of (n1, n2), listed by abs(m1-m2)
        n1n2_combinations = [[] for dm in range(2*m_max+1)]
        for n1 in range(blocksize):
            m1 = m_list[n1]
            for n2 in range(blocksize):
                m2 = m_list[n2]
                n1n2_combinations[abs(m1-m2)].append((n1,n2))
                       
        wr = np.zeros((len_rho, blocksize, blocksize), dtype=complex)
        
        dkp = np.diff(k_parallel)
        if cu.use_gpu:
            re_dkp_d = cu.gpuarray.to_gpu(np.float32(dkp.real))
            im_dkp_d = cu.gpuarray.to_gpu(np.float32(dkp.imag))
            kernel_source_code = cusrc.radial_lookup_assembly_code %(blocksize, len_rho, len_kp)
            helper_function = cu.SourceModule(kernel_source_code).get_function("helper")
            cuda_blocksize = 128
            cuda_gridsize = (len_rho + cuda_blocksize - 1) // cuda_blocksize
    
            re_dwr_d = cu.gpuarray.to_gpu(np.zeros(len_rho, dtype=np.float32))
            im_dwr_d = cu.gpuarray.to_gpu(np.zeros(len_rho, dtype=np.float32))
        n1n2_combinations = [[] for dm in range(2*m_max+1)]
        for n1 in range(blocksize):
            m1 = m_list[n1]
            for n2 in range(blocksize):
                m2 = m_list[n2]
                n1n2_combinations[abs(m1-m2)].append((n1,n2))
        
        pbar = tqdm(total=blocksize**2, 
                    desc='Layer mediated coupling   ', 
                    file=sys.stdout,
                    bar_format='{l_bar}{bar}| elapsed: {elapsed} remaining: {remaining}')
                
        for dm in range(2*m_max+1):
            bessel = scipy.special.jv(dm, (k_parallel[None,:]*radial_distance_array[:,None]))
            besjac = bessel * (k_parallel / (kz_is * k_is))[None,:]
            for n1n2 in n1n2_combinations[dm]:
                n1 = n1n2[0]
                m1 = m_list[n1]
                n2 = n1n2[1]
                m2 = m_list[n2]
                
                belbe = np.zeros(len_kp, dtype=complex)  # n1, n2, kp
                for pol in range(2):
                    belbe += L[pol,0,0,:] * B_dag[pol,0,n1,:] * B[pol,0,n2,:]
                    belbe += L[pol,1,0,:] * B_dag[pol,1,n1,:] * B[pol,0,n2,:] * emn2jkz
                    belbe += L[pol,0,1,:] * B_dag[pol,0,n1,:] * B[pol,1,n2,:] * epl2jkz
                    belbe += L[pol,1,1,:] * B_dag[pol,1,n1,:] * B[pol,1,n2,:]
                
                if cu.use_gpu:
                    re_belbe_d = cu.gpuarray.to_gpu(np.float32(belbe[None, :].real))
                    im_belbe_d = cu.gpuarray.to_gpu(np.float32(belbe[None, :].imag))
    
                    re_besjac_d = cu.gpuarray.to_gpu(np.float32(besjac.real))
                    im_besjac_d = cu.gpuarray.to_gpu(np.float32(besjac.imag))
                    helper_function(re_besjac_d.gpudata, im_besjac_d.gpudata, 
                                    re_belbe_d.gpudata, im_belbe_d.gpudata, 
                                    re_dkp_d.gpudata, im_dkp_d.gpudata, 
                                    re_dwr_d.gpudata, im_dwr_d.gpudata, 
                                    block=(cuda_blocksize, 1, 1), grid=(cuda_gridsize, 1))
                    
                    wr[:,n1,n2] = 4 * (1j)**abs(m2-m1) * (re_dwr_d.get() + 1j*im_dwr_d.get()) 
    
                else:
                    integrand = besjac * belbe[None, :]  # rho, kp 
                    wr[:,n1,n2] = 2 * (1j)**abs(m2-m1) * ((integrand[:,:-1] + integrand[:,1:]) 
                                                                * dkp[None,:]).sum(axis=-1)  # trapezoidal rule
                pbar.update()
        pbar.close()
        w = w + wr
    else:
        sys.stdout.write('Omitting calculations of layer mediated coupling since no difference in layers was detected. \n')
        sys.stdout.flush()
        
    return w , radial_distance_array


def size_format(b):
    if b < 1000:
              return '%i' % b + 'B'
    elif 1000 <= b < 1000000:
        return '%.1f' % float(b/1000) + 'KB'
    elif 1000000 <= b < 1000000000:
        return '%.1f' % float(b/1000000) + 'MB'
    elif 1000000000 <= b < 1000000000000:
        return '%.1f' % float(b/1000000000) + 'GB'
    elif 1000000000000 <= b:
        return '%.1f' % float(b/1000000000000) + 'TB'



def radial_direct_pwe_mediated_coupling_lookup_table(vacuum_wavelength, rho_max, l_max, k_is, k_parallel='default', resolution=None):
    """Prepare Sommerfeld integral lookup table to allow for a fast calculation of the coupling matrix by interpolation.
    This function is called when all particles are on the same z-position.
    
    This function uses PVWF representation of the coupling operator and is only
    for the direct particle-particle lookup table for particles in close vicinity. 
    
    
    Args:
        vacuum_wavelength (float):  Vacuum wavelength in length units
        rho_max (float):            Maximal horizontal particle displacement for which 
                                    direct coupling is evaluated via PVWFs
        l_max (int):                Maximal multipole degree
        k_is (complex):             Wavenumber of the scattering layer
        k_parallel (numpy.ndarray or str):           In-plane wavenumber for Sommerfeld integrals.
                                                     If 'default', smuthi.fields.default_Sommerfeld_k_parallel_array
        resolution (float): Spatial resolution of lookup table in length units. (default: vacuum_wavelength / 100)       
                            Smaller means more accurate but higher memory footprint 
                            
    Returns:
        (tuple) tuple containing:
        
            lookup_table (ndarray):  Coupling lookup, indices are [rho, n1, n2].
            rho_array (ndarray):     Values for the radial distance considered for the lookup (starting from negative 
                                     numbers to allow for simpler cubic interpolation without distinction of cases 
                                     at rho=0)
    """

    sys.stdout.write('Prepare radial particle coupling lookup:\n')
    sys.stdout.flush()

    if resolution is None:
        resolution = vacuum_wavelength / 100
        sys.stdout.write('Setting lookup resolution to %f\n'%resolution)
        sys.stdout.flush()
    
    m_max = l_max # necessary for pwe coupling
    blocksize = smuthi.fields.blocksize(l_max, m_max)
    
    radial_distance_array = np.arange(- 3 * resolution, rho_max + 3 * resolution, resolution)   
    len_rho = len(radial_distance_array)
    
    # pwe-coupling, similar to layer mediated coupling -------------------------------------------------------------------------------------
    sys.stdout.write('PWE-coupling correction for direct coupling: ...')
    sys.stdout.flush()
    
    if type(k_parallel) == str and k_parallel == 'default':
        k_parallel = smuthi.fields.default_Sommerfeld_k_parallel_array

    kz_is = smuthi.fields.k_z(k_parallel=k_parallel, k=k_is)
    len_kp = len(k_parallel)

       
    # transformation coefficients
    B_dag = np.zeros((2, blocksize, len_kp), dtype=complex)  # pol, n, kp (pl only)
    B = np.zeros((2, blocksize, len_kp), dtype=complex)  # pol, n, kp (pl only)
    ct = kz_is / k_is
    st = k_parallel / k_is
    _, pilm_pl, taulm_pl = sf.legendre_normalized(ct, st, l_max)
    _, pilm_mn, taulm_mn = sf.legendre_normalized(-ct, st, l_max)

    m_list = [None for n in range(blocksize)]
    for tau in range(2):
        for m in range(-m_max, m_max + 1):
            for l in range(max(1, abs(m)), l_max + 1):
                n = smuthi.fields.multi_to_single_index(tau, l, m, l_max, m_max)
                m_list[n] = m
                for pol in range(2):
                    B_dag[pol,n,:] = trf.transformation_coefficients_vwf(tau, l, m, pol, pilm_list=pilm_pl,
                                                                           taulm_list=taulm_pl, dagger=True)
                    B[pol,n,:] = trf.transformation_coefficients_vwf(tau, l, m, pol, pilm_list=pilm_pl,
                                                                       taulm_list=taulm_pl, dagger=False)

    
    # pairs of (n1, n2), listed by abs(m1-m2)
    n1n2_combinations = [[] for dm in range(2*m_max+1)]
    for n1 in range(blocksize):
        m1 = m_list[n1]
        for n2 in range(blocksize):
            m2 = m_list[n2]
            n1n2_combinations[abs(m1-m2)].append((n1,n2))
    
    # in the rotated coordinate system, the radial distance rho is represented 
    # as a vertical distance dz
    dz_array = radial_distance_array 
    epljkdz = np.exp(1j * kz_is[None, :] * dz_array[:, None])               
    w = np.zeros((len_rho, blocksize, blocksize), dtype=complex)
    
    dkp = np.diff(k_parallel)
    if cu.use_gpu:
        re_dkp_d = cu.gpuarray.to_gpu(np.float32(dkp.real))
        im_dkp_d = cu.gpuarray.to_gpu(np.float32(dkp.imag))
        kernel_source_code = cusrc.radial_lookup_assembly_code_PVWF %(blocksize, len_rho, len_kp)
        helper_function = cu.SourceModule(kernel_source_code).get_function("helper")
        cuda_blocksize = 128
        cuda_gridsize = (len_rho + cuda_blocksize - 1) // cuda_blocksize

        re_dw_d = cu.gpuarray.to_gpu(np.zeros(len_rho, dtype=np.float32))
        im_dw_d = cu.gpuarray.to_gpu(np.zeros(len_rho, dtype=np.float32))
    n1n2_combinations = [[] for dm in range(2*m_max+1)]
    for n1 in range(blocksize):
        m1 = m_list[n1]
        for n2 in range(blocksize):
            m2 = m_list[n2]
            n1n2_combinations[abs(m1-m2)].append((n1,n2))
    
    pbar = tqdm(total=blocksize**2, 
                desc='PWE direct coupling   ', 
                file=sys.stdout,
                bar_format='{l_bar}{bar}| elapsed: {elapsed} remaining: {remaining}')
            
    for dm in range(2*m_max+1):
        # in the rotated coordinate system, rho is always zero
        bessel = scipy.special.jv(dm, (k_parallel[None,:]*np.zeros(len_rho)[:,None]))
        besjac = bessel * (k_parallel / (kz_is * k_is))[None,:]
        for n1n2 in n1n2_combinations[dm]:
            n1 = n1n2[0]
            m1 = m_list[n1]
            n2 = n1n2[1]
            m2 = m_list[n2]
            
            bebe = np.zeros((len_rho, len_kp), dtype=complex) # n1, n2, rho, kp
            for pol in range(2):
                # always use B+ (all kz > 0)
                bebe += ((B_dag[pol, n1, :] * B[pol, n2, :])[None, :] * epljkdz) 
                
            if cu.use_gpu:
                re_bebe_d = cu.gpuarray.to_gpu(np.float32(bebe.real)) 
                im_bebe_d = cu.gpuarray.to_gpu(np.float32(bebe.imag)) 

                re_besjac_d = cu.gpuarray.to_gpu(np.float32(besjac.real))
                im_besjac_d = cu.gpuarray.to_gpu(np.float32(besjac.imag))
                helper_function(re_besjac_d.gpudata, im_besjac_d.gpudata, 
                                re_bebe_d.gpudata, im_bebe_d.gpudata, 
                                re_dkp_d.gpudata, im_dkp_d.gpudata, 
                                re_dw_d.gpudata, im_dw_d.gpudata, 
                                block=(cuda_blocksize, 1, 1), grid=(cuda_gridsize, 1))
                
                w[:,n1,n2] = 4 * (1j)**abs(m2-m1) * (re_dw_d.get() + 1j*im_dw_d.get()) 

            else:
                integrand = besjac * bebe  # rho, kp 
                w[:,n1,n2] = 2 * (1j)**abs(m2-m1) * ((integrand[:,:-1] + integrand[:,1:]) 
                                                            * dkp[None,:]).sum(axis=-1)  # trapezoidal rule
            pbar.update()
    pbar.close()
    
    alpha = 0
    beta = -np.pi / 2
    rot_mat_1 = np.transpose(trf.block_rotation_matrix_D_svwf(l_max, m_max, 0, beta, alpha))
    rot_mat_2 = np.transpose(trf.block_rotation_matrix_D_svwf(l_max, m_max, -alpha, -beta, 0))
    w = np.dot(np.einsum('ij,ljk->lik', rot_mat_1, w), rot_mat_2)
    
    return w, radial_distance_array
