"""Basic functions for async executions."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../src/utils/async_fn.ipynb.

# %% auto #0
__all__ = ['run_parallel_jobs', 'MemoryMappedArray']

# %% ../../src/utils/async_fn.ipynb #06272f6f-b4e3-4504-a90a-feebbf6ad821
from ..imports import *
from joblib import Parallel, delayed

from tensordict.tensordict import MemoryMappedTensor
import tempfile

# %% ../../src/utils/async_fn.ipynb #7c3dbea0-513a-4754-b24b-23bfde48d423
def run_parallel_jobs(f: callable, loop_set, n_jobs: int = 1):      
    if n_jobs > 1: res = Parallel(n_jobs=n_jobs)(delayed(f)(x) for x in loop_set)         
    else:          res = [f(x) for x in loop_set]      
    return res 

# %% ../../src/utils/async_fn.ipynb #141766e2-4e5f-45d8-bda8-f606d56e6774
class MemoryMappedArray():
    def __init__(self, obj, type="tensor"):
        self.obj  = obj
        self.type = type
        assert type in ["tensor", "numpy"]
        
        with tempfile.NamedTemporaryFile(delete=False) as file:   
            # Note can bes simplified with python 3.12 as we can set delete=true, and delete_on_close=True, so it will be kept and we dont need to delete
            # see https://docs.python.org/3.12/library/tempfile.html
            
            self.temporaryFileName = file.name
            file.close()

        if self.type == "numpy":
            self.obj_memmap    = np.memmap(filename=self.temporaryFileName, dtype=obj.dtype, mode='w+', shape=obj.shape)
            self.obj_memmap[:] = self.obj[:]
            self.obj_memmap.flush()
        
        elif self.type == "tensor":  
            self.obj_memmap = MemoryMappedTensor.from_tensor(self.obj.cpu(), filename=self.temporaryFileName, existsok=True) 
        
        else: 
            raise NotImplementedError()

    def get_obj(self):
        if self.type == "numpy":
            self.obj = self.obj_memmap.copy()
            
        elif self.type == "tensor":
            self.obj = self.obj_memmap.contiguous().clone().to(self.obj.device)

        del self.obj_memmap
        return self.obj, self.temporaryFileName

    @staticmethod
    def clean(temp_files):
        for temp_file in temp_files:
            try: os.remove(temp_file)
            except Exception as e: print(f"[ERROR]: {e}") 
