<commit_before>from os.path import join

from wkw import wkw
import numpy as np

import wkcuber.api as api
from wkcuber.api.View import WKView, TiffView
from wkcuber.api.TiffData.TiffMag import TiffMagHeader
from wkcuber.mag import Mag


class MagDataset:
    def __init__(self, layer, name):
        self.layer = layer
        self.name = name
        self.header = self.get_header()

        self.view = self.get_view(is_bounded=False)

    def open(self):
        self.view.initialize()

    def close(self):
        self.view.close()

    def read(self, size, offset=(0, 0, 0)) -> np.array:
        return self.view.read(size, offset)

    def write(self, data, offset=(0, 0, 0)):
        self._assert_valid_num_channels(data.shape)
        self.view.write(data, offset)
        layer_properties = self.layer.dataset.properties.data_layers[self.layer.name]
        current_offset_in_mag1 = layer_properties.get_bounding_box_offset()
        current_size_in_mag1 = layer_properties.get_bounding_box_size()

        mag = Mag(self.name)
        mag_np = mag.as_np()

        offset_in_mag1 = tuple(np.array(offset) * mag_np)

        new_offset_in_mag1 = (
            offset_in_mag1
            if current_offset_in_mag1 == (-1, -1, -1)
            else tuple(min(x) for x in zip(current_offset_in_mag1, offset_in_mag1))
        )

        old_end_offset_in_mag1 = np.array(current_offset_in_mag1) + np.array(
            current_size_in_mag1
        )
        new_end_offset_in_mag1 = (np.array(offset) + np.array(data.shape[-3:])) * mag_np
        max_end_offset_in_mag1 = np.array(
            [old_end_offset_in_mag1, new_end_offset_in_mag1]
        ).max(axis=0)
        total_size_in_mag1 = max_end_offset_in_mag1 - np.array(new_offset_in_mag1)
        total_size = total_size_in_mag1 / mag_np

        self.view.size = tuple(total_size)

        self.layer.dataset.properties._set_bounding_box_of_layer(
            self.layer.name, tuple(new_offset_in_mag1), tuple(total_size_in_mag1)
        )

    def get_header(self):
        raise NotImplementedError

    def get_view(self, size=None, offset=None, is_bounded=True):
        size_in_properties = self.layer.dataset.properties.data_layers[
            self.layer.name
        ].get_bounding_box_size()

        if offset is None:
            offset = (0, 0, 0)

        if size is None:
            size = size_in_properties

        # assert that the parameter size is valid
        for s1, s2, off in zip(size_in_properties, size, offset):
            if s2 + off > s1 and is_bounded:
                raise AssertionError(
                    f"The combination of the passed parameter 'size' {size} and {offset} are not compatible with the "
                    f"size ({size_in_properties}) from the properties.json."
                )

        mag_file_path = join(self.layer.dataset.path, self.layer.name, self.name)
        offset_in_properties = self.layer.dataset.properties.data_layers[
            self.layer.name
        ].get_bounding_box_offset()
        dataset_offset = (
            (0, 0, 0) if offset_in_properties == (-1, -1, -1) else offset_in_properties
        )
        global_offset = np.array(dataset_offset) + np.array(offset)
        return self._get_view_type()(
            mag_file_path, self.header, size, global_offset, is_bounded
        )

    def _get_view_type(self):
        raise NotImplementedError

    def _assert_valid_num_channels(self, write_data_shape):
        num_channels = self.layer.num_channels
        if len(write_data_shape) == 3:
            assert (
                num_channels == 1
            ), f"The number of channels of the dataset ({num_channels}) does not match the number of channels of the passed data (1)"
        else:
            assert (
                num_channels == 3
            ), f"The number of channels of the dataset ({num_channels}) does not match the number of channels of the passed data ({write_data_shape[0]})"


class WKMagDataset(MagDataset):
    def __init__(self, layer, name, block_len, file_len, block_type):
        self.block_len = block_len
        self.file_len = file_len
        self.block_type = block_type
        super().__init__(layer, name)

    def get_header(self) -> wkw.Header:
        return wkw.Header(
            voxel_type=self.layer.dtype,
            num_channels=self.layer.num_channels,
            version=1,
            block_len=self.block_len,
            file_len=self.file_len,
            block_type=self.block_type,
        )

    @classmethod
    def create(cls, layer, name, block_len, file_len, block_type):
        mag_dataset = cls(layer, name, block_len, file_len, block_type)
        wkw.Dataset.create(<commit_msg>fix channel checking for dataset-write (#208)<commit_after>from os.path import join

from wkw import wkw
import numpy as np

import wkcuber.api as api
from wkcuber.api.View import WKView, TiffView
from wkcuber.api.TiffData.TiffMag import TiffMagHeader
from wkcuber.mag import Mag


class MagDataset:
    def __init__(self, layer, name):
        self.layer = layer
        self.name = name
        self.header = self.get_header()

        self.view = self.get_view(is_bounded=False)

    def open(self):
        self.view.initialize()

    def close(self):
        self.view.close()

    def read(self, size, offset=(0, 0, 0)) -> np.array:
        return self.view.read(size, offset)

    def write(self, data, offset=(0, 0, 0)):
        self._assert_valid_num_channels(data.shape)
        self.view.write(data, offset)
        layer_properties = self.layer.dataset.properties.data_layers[self.layer.name]
        current_offset_in_mag1 = layer_properties.get_bounding_box_offset()
        current_size_in_mag1 = layer_properties.get_bounding_box_size()

        mag = Mag(self.name)
        mag_np = mag.as_np()

        offset_in_mag1 = tuple(np.array(offset) * mag_np)

        new_offset_in_mag1 = (
            offset_in_mag1
            if current_offset_in_mag1 == (-1, -1, -1)
            else tuple(min(x) for x in zip(current_offset_in_mag1, offset_in_mag1))
        )

        old_end_offset_in_mag1 = np.array(current_offset_in_mag1) + np.array(
            current_size_in_mag1
        )
        new_end_offset_in_mag1 = (np.array(offset) + np.array(data.shape[-3:])) * mag_np
        max_end_offset_in_mag1 = np.array(
            [old_end_offset_in_mag1, new_end_offset_in_mag1]
        ).max(axis=0)
        total_size_in_mag1 = max_end_offset_in_mag1 - np.array(new_offset_in_mag1)
        total_size = total_size_in_mag1 / mag_np

        self.view.size = tuple(total_size)

        self.layer.dataset.properties._set_bounding_box_of_layer(
            self.layer.name, tuple(new_offset_in_mag1), tuple(total_size_in_mag1)
        )

    def get_header(self):
        raise NotImplementedError

    def get_view(self, size=None, offset=None, is_bounded=True):
        size_in_properties = self.layer.dataset.properties.data_layers[
            self.layer.name
        ].get_bounding_box_size()

        if offset is None:
            offset = (0, 0, 0)

        if size is None:
            size = size_in_properties

        # assert that the parameter size is valid
        for s1, s2, off in zip(size_in_properties, size, offset):
            if s2 + off > s1 and is_bounded:
                raise AssertionError(
                    f"The combination of the passed parameter 'size' {size} and {offset} are not compatible with the "
                    f"size ({size_in_properties}) from the properties.json."
                )

        mag_file_path = join(self.layer.dataset.path, self.layer.name, self.name)
        offset_in_properties = self.layer.dataset.properties.data_layers[
            self.layer.name
        ].get_bounding_box_offset()
        dataset_offset = (
            (0, 0, 0) if offset_in_properties == (-1, -1, -1) else offset_in_properties
        )
        global_offset = np.array(dataset_offset) + np.array(offset)
        return self._get_view_type()(
            mag_file_path, self.header, size, global_offset, is_bounded
        )

    def _get_view_type(self):
        raise NotImplementedError

    def _assert_valid_num_channels(self, write_data_shape):
        num_channels = self.layer.num_channels
        if len(write_data_shape) == 3:
            assert (
                num_channels == 1
            ), f"The number of channels of the dataset ({num_channels}) does not match the number of channels of the passed data (1)"
        else:
            assert (
                num_channels == write_data_shape[0]
            ), f"The number of channels of the dataset ({num_channels}) does not match the number of channels of the passed data ({write_data_shape[0]})"


class WKMagDataset(MagDataset):
    def __init__(self, layer, name, block_len, file_len, block_type):
        self.block_len = block_len
        self.file_len = file_len
        self.block_type = block_type
        super().__init__(layer, name)

    def get_header(self) -> wkw.Header:
        return wkw.Header(
            voxel_type=self.layer.dtype,
            num_channels=self.layer.num_channels,
            version=1,
            block_len=self.block_len,
            file_len=self.file_len,
            block_type=self.block_type,
        )

    @classmethod
    def create(cls, layer, name, block_len, file_len, block_type):
        mag_dataset = cls(layer, name, block_len, file_len, block_type)
        wkw.Dataset.create(
{'fraction_non_alphanumeric': 0.09142736486486487, 'fraction_numerical': 0.011085304054054054, 'mean_word_length': 2.4310032596885187, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '326902', 'n_tokens_mistral': 3055, 'n_tokens_neox': 2974, 'n_words': 646}