# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/ffcv.writer.ipynb.

# %% auto 0
__all__ = ['DatasetWriter']

# %% ../../nbs/ffcv.writer.ipynb 1
# Contains code from:
# FFCV - Apache License 2.0 - Copyright (c) 2022 FFCV Team

# %% ../../nbs/ffcv.writer.ipynb 4
from typing import Mapping
import numpy as np
from multiprocessing import shared_memory, cpu_count

from ffcv.utils import is_power_of_2
from ffcv.types import get_metadata_type, HeaderType, FieldDescType, CURRENT_VERSION
from ffcv.fields import Field, BytesField, IntField, FloatField, NDArrayField, JSONField, TorchTensorField
from ffcv.fields import RGBImageField as _RGBImageField
from ffcv.writer import DatasetWriter as _DatasetWriter
from ffcv.writer import MIN_PAGE_SIZE, MAX_PAGE_SIZE

from .fields import RGBImageField

# %% ../../nbs/ffcv.writer.ipynb 6
HANDLER_TYPE_ID = {
    None: 255, 
    FloatField: 0, 
    IntField: 1, 
    _RGBImageField: 2, 
    RGBImageField: 2,
    BytesField: 3, 
    NDArrayField: 4, 
    JSONField: 5, 
    TorchTensorField: 6
}

# %% ../../nbs/ffcv.writer.ipynb 8
class DatasetWriter(_DatasetWriter):
    """Writes given dataset into FFCV format (.beton).
    Supports indexable objects (e.g., PyTorch Datasets) and webdataset.

    Parameters
    ----------
    fname: str
        File name to store dataset in FFCV format (.beton)
    fields : Mapping[str, Field]
        Map from keys to Field's (order matters!)
    page_size : int
        Page size used internally
    num_workers : int
        Number of processes to use
    """
    def __init__(self, fname: str, fields: Mapping[str, Field],
                 page_size: int = 4 * MIN_PAGE_SIZE, num_workers: int = -1):
        self.fields = fields
        self.fname = fname
        self.metadata_type = get_metadata_type(list(self.fields.values()))

        self.num_workers = num_workers
        # We use all cores by default
        if self.num_workers < 1:
            self.num_workers = cpu_count()

        if not is_power_of_2(page_size):
            raise ValueError(f'page_size isnt a power of 2')
        if page_size < MIN_PAGE_SIZE:
            raise ValueError(f"page_size can't be lower than{MIN_PAGE_SIZE}")
        if page_size >= MAX_PAGE_SIZE:
            raise ValueError(f"page_size can't be bigger(or =) than{MAX_PAGE_SIZE}")

        self.page_size = page_size


    def prepare(self):

        with open(self.fname, 'wb') as fp:
            # Prepare the header data
            header = np.zeros(1, dtype=HeaderType)[0]
            header['version'] = CURRENT_VERSION
            header['num_samples'] = self.num_samples
            header['num_fields'] = len(self.fields)
            header['page_size'] = self.page_size
            self.header = header

            # We will write the header at the end because we need to know where
            # The memory allocation table is in the file
            # We still write it here to make space for it later
            fp.write(self.header.tobytes())


            # Writes the information about the fields
            fields_descriptor = np.zeros(len(self.fields), dtype=FieldDescType)

            fieldname_max_len = fields_descriptor[0]['name'].shape[0]

            for i, (name, field) in enumerate(self.fields.items()):
                type_id = HANDLER_TYPE_ID.get(type(field), 255)
                encoded_name = name.encode('ascii')
                encoded_name = np.frombuffer(encoded_name, dtype='<u1')
                actual_length = min(fieldname_max_len, len(encoded_name))
                fields_descriptor[i]['type_id'] = type_id
                fields_descriptor[i]['name'][:actual_length] = (
                    encoded_name[:actual_length])
                fields_descriptor[i]['arguments'][:] = field.to_binary()[0]

            fp.write(fields_descriptor.tobytes())

        total_metadata_size = self.num_samples * self.metadata_type.itemsize

        # Shared memory for all the writers to fill the information
        self.metadata_sm = 3
        self.metadata_start = (HeaderType.itemsize + fields_descriptor.nbytes)

        self.metadata_sm = shared_memory.SharedMemory(create=True,
                                                      size=total_metadata_size)

        self.data_region_start = self.metadata_start + total_metadata_size
