#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Feb 10 09:24:58 2021
"""

#%%
# Loading the required libraries! 

from itertools import repeat
import numpy as np
from multiprocessing.shared_memory import SharedMemory
from multiprocessing import Pool, freeze_support
#%%
# Making the forest! 

toy_forest = [[list(range(2)), list(range(3))]]
np_arr = np.array(toy_forest, dtype=object)

# Create a shared memory instance
shm_f = SharedMemory(create=True, size=np_arr.nbytes)
b_shm = np.ndarray(np_arr.shape, dtype=np_arr.dtype, buffer=shm_f.buf)

# Copy the original data into shared memory
b_shm[:] = np_arr[:]

# Writing the function to map over! 


def dummy_function(i, shm_f, shape_f, dtype_f):
    forest = np.ndarray(shape_f, dtype= dtype_f, buffer=shm_f.buf)
    a = forest[0][0][0]
    return a

#%%
# Doing the multiprocessing! 



if __name__ == '__main__':
    n_test = 2
    maxworkers = 2 
    x = 6
    with Pool(processes = maxworkers) as pool:
        result = pool.starmap(func = dummy_function, 
                              iterable = zip(list(range(n_test)),
                                              repeat(shm_f),
                                              repeat(b_shm.shape),
                                              repeat(b_shm.dtype)))
    print("done")