<commit_before>import filecmp
import itertools
import json
import os
import warnings
from os.path import dirname, join
from pathlib import Path
from typing import Any, Tuple, cast, Generator

import pytest

import numpy as np
from shutil import rmtree, copytree

from wkw import wkw
from wkw.wkw import WKWException

from wkcuber.api.dataset import Dataset, DatasetViewConfiguration
from wkcuber.api.bounding_box import BoundingBox
from os import makedirs

from wkcuber.api.layer import (
    Layer,
    LayerCategories,
    SegmentationLayer,
    LayerViewConfiguration,
)
from wkcuber.api.mag_view import MagView
from wkcuber.api.properties.dataset_properties import Properties
from wkcuber.api.properties.layer_properties import SegmentationLayerProperties
from wkcuber.api.properties.resolution_properties import Resolution
from wkcuber.api.view import View
from wkcuber.compress import compress_mag_inplace
from wkcuber.mag import Mag
from wkcuber.utils import get_executor_for_args, named_partial, _snake_to_camel_case

TESTDATA_DIR = Path("testdata")
TESTOUTPUT_DIR = Path("testoutput")


def delete_dir(relative_path: Path) -> None:
    if relative_path.exists() and relative_path.is_dir():
        rmtree(relative_path)


def chunk_job(args: Tuple[View, int]) -> None:
    (view, i) = args
    # increment the color value of each voxel
    data = view.read(size=view.size)
    if data.shape[0] == 1:
        data = data[0, :, :, :]
    data += 50
    view.write(data)


def advanced_chunk_job(args: Tuple[View, int], type: type) -> None:
    view, i = args

    # write different data for each chunk (depending on the global_offset of the chunk)
    data = view.read(size=view.size)
    data = np.ones(data.shape, dtype=type) * type(sum(view.global_offset))
    view.write(data)


def for_each_chunking_with_wrong_chunk_size(view: View) -> None:
    with get_executor_for_args(None) as executor:
        with pytest.raises(AssertionError):
            view.for_each_chunk(
                chunk_job,
                chunk_size=(0, 64, 64),
                executor=executor,
            )
        with pytest.raises(AssertionError):
            view.for_each_chunk(
                chunk_job,
                chunk_size=(16, 64, 64),
                executor=executor,
            )
        with pytest.raises(AssertionError):
            view.for_each_chunk(
                chunk_job,
                chunk_size=(100, 64, 64),
                executor=executor,
            )


def for_each_chunking_advanced(ds: Dataset, view: View) -> None:
    chunk_size = (64, 64, 64)
    with get_executor_for_args(None) as executor:
        func = named_partial(advanced_chunk_job, type=np.uint8)
        view.for_each_chunk(
            func,
            chunk_size=chunk_size,
            executor=executor,
        )

    for offset, size in [
        ((10, 10, 10), (54, 54, 54)),
        ((10, 64, 10), (54, 64, 54)),
        ((10, 128, 10), (54, 32, 54)),
        ((64, 10, 10), (64, 54, 54)),
        ((64, 64, 10), (64, 64, 54)),
        ((64, 128, 10), (64, 32, 54)),
        ((128, 10, 10), (32, 54, 54)),
        ((128, 64, 10), (32, 64, 54)),
        ((128, 128, 10), (32, 32, 54)),
    ]:
        chunk = ds.get_layer("color").get_mag("1").get_view(size=size, offset=offset)
        chunk_data = chunk.read(size=chunk.size)
        assert np.array_equal(
            np.ones(chunk_data.shape, dtype=np.uint8)
            * np.uint8(sum(chunk.global_offset)),
            chunk_data,
        )


def copy_and_transform_job(args: Tuple[View, View, int], name: str, val: int) -> None:
    (source_view, target_view, i) = args
    # This method simply takes the data from the source_view, transforms it and writes it to the target_view

    # These assertions are just to demonstrate how the passed parameters can be accessed inside this method
    assert name == "foo"
    assert val == 42

    # increment the color value of each voxel
    data = source_view.read(size=source_view.size)
    if data.shape[0] == 1:
        data = data[0, :, :, :]
    data += 50
    target_view.write(data)


def get_multichanneled_data(dtype: type) -> np.ndarray:
    data = np.zeros((3, 250, 200, 10), dtype=dtype)
    for h in range(10):
        for i in range(250):
            for j in range(200):
                data[0, i, j, h] = i * 256
                data[1, i, j, h] = j * 256
                data[2, i, j, h] = 100 * 256
    return data


def test_create_dataset_with_layer_and_mag() -> None:
    delete_dir(TESTOUTPUT_DIR / "wk_dataset")

    ds = Dataset.create(TESTOUTPUT_DIR / "wk_dataset", scale=(1, 1, 1))
    ds.add_layer("color", "color")

    ds.get_layer("color").add_mag("1")
    ds.get_layer("color").add_mag("2-2-1")

    assert (TESTOUTPUT_DIR / "wk_dataset" / "color" / "1").exists()
    assert (TESTOUTPUT_DIR / "wk_dataset" / "color" / "2-2-1").exists()

    assert len(ds.properties.data_layers) == 1
    assert len(ds.properties.data_layers["color"].wkw_magnifications) == 2


def test_create_dataset_with_explicit_header_fields() -> None:
    delete_dir(TESTOUTPUT_DIR / "wk_dataset_advanced")

    ds = Dataset.create(TESTOUTPUT_DIR / "wk_dataset_advanced", scale=(1, 1, 1))
    ds.add_layer(
        "color", LayerCategories.COLOR_TYPE, dtype_per_layer="uint48", num_channels=3
    )

    ds.get_layer("color").add_mag("1", block_len=64, file_len=64)
    ds.get_layer("color").add_mag("2-2-1")

    assert (TESTOUTPUT_DIR / "wk_dataset_advanced" / "color" / "1").exists()
    assert (TESTOUTPUT_DIR / "wk_dataset_advanced" / "color" / "2-2-1").exists()

    assert len(ds.properties.data_layers) == 1
    assert len(ds.properties.data_layers["color"].wkw_magnifications) == 2

    assert ds.properties.data_layers["color"].element_class == "uint48"
    assert (
        cast(
            Resolution, ds.properties.data_layers["color"].wkw_magnifications[0]
        ).cube_length
        == 64 * 64
    )  # mag "1"
    assert ds.properties.data_layers["color"].wkw_magnifications[0].mag == Mag("1")
    assert (
        cast(
            Resolution, ds.properties.data_layers["color"].wkw_magnifications[1]
        ).cube_length
        == 32 * 32
    )  # mag "2-2-1" (defaults are used)
    assert ds.properties.data_layers["color"].wkw_magnifications[1].mag == Mag("2-2-1")


def test_open_dataset() -> None:
    ds = Dataset(TESTDATA_DIR / "simple_wk_dataset")

    assert len(ds.properties.data_layers) == 1
    assert len(ds.properties.data_layers["color"].wkw_magnifications) == 1


def test_view_read_with_open() -> None:
    wk_view = (
        Dataset(TESTDATA_DIR / "simple_wk_dataset")
        .get_layer("color")
        .get_mag("1")
        .get_view(size=(16, 16, 16))
    )

    assert not wk_view._is_opened

    with wk_view.open():
        assert wk_view._is_opened

        data = wk_view.read(size=(10, 10, 10))
        assert data.shape == (3, 10, 10, 10)  # three channel

    assert not wk_view._is_opened


def test_view_read_without_open() -> None:
    wk_view = (
        Dataset(TESTDATA_DIR / "simple_wk_dataset")
        .get_layer("color")
        .get_mag("1")
        .get_view(size=(16, 16, 16))
    )
<commit_msg>[ds-api] Modification of dataset with float-type layer results in JSON dump error (#358)<commit_after>import filecmp
import itertools
import json
import os
import warnings
from os.path import dirname, join
from pathlib import Path
from typing import Any, Tuple, cast, Generator

import pytest

import numpy as np
from shutil import rmtree, copytree

from wkw import wkw
from wkw.wkw import WKWException

from wkcuber.api.dataset import Dataset, DatasetViewConfiguration
from wkcuber.api.bounding_box import BoundingBox
from os import makedirs

from wkcuber.api.layer import (
    Layer,
    LayerCategories,
    SegmentationLayer,
    LayerViewConfiguration,
)
from wkcuber.api.mag_view import MagView
from wkcuber.api.properties.dataset_properties import Properties
from wkcuber.api.properties.layer_properties import SegmentationLayerProperties
from wkcuber.api.properties.resolution_properties import Resolution
from wkcuber.api.view import View
from wkcuber.compress import compress_mag_inplace
from wkcuber.mag import Mag
from wkcuber.utils import get_executor_for_args, named_partial, _snake_to_camel_case

TESTDATA_DIR = Path("testdata")
TESTOUTPUT_DIR = Path("testoutput")


def delete_dir(relative_path: Path) -> None:
    if relative_path.exists() and relative_path.is_dir():
        rmtree(relative_path)


def chunk_job(args: Tuple[View, int]) -> None:
    (view, i) = args
    # increment the color value of each voxel
    data = view.read(size=view.size)
    if data.shape[0] == 1:
        data = data[0, :, :, :]
    data += 50
    view.write(data)


def advanced_chunk_job(args: Tuple[View, int], type: type) -> None:
    view, i = args

    # write different data for each chunk (depending on the global_offset of the chunk)
    data = view.read(size=view.size)
    data = np.ones(data.shape, dtype=type) * type(sum(view.global_offset))
    view.write(data)


def for_each_chunking_with_wrong_chunk_size(view: View) -> None:
    with get_executor_for_args(None) as executor:
        with pytest.raises(AssertionError):
            view.for_each_chunk(
                chunk_job,
                chunk_size=(0, 64, 64),
                executor=executor,
            )
        with pytest.raises(AssertionError):
            view.for_each_chunk(
                chunk_job,
                chunk_size=(16, 64, 64),
                executor=executor,
            )
        with pytest.raises(AssertionError):
            view.for_each_chunk(
                chunk_job,
                chunk_size=(100, 64, 64),
                executor=executor,
            )


def for_each_chunking_advanced(ds: Dataset, view: View) -> None:
    chunk_size = (64, 64, 64)
    with get_executor_for_args(None) as executor:
        func = named_partial(advanced_chunk_job, type=np.uint8)
        view.for_each_chunk(
            func,
            chunk_size=chunk_size,
            executor=executor,
        )

    for offset, size in [
        ((10, 10, 10), (54, 54, 54)),
        ((10, 64, 10), (54, 64, 54)),
        ((10, 128, 10), (54, 32, 54)),
        ((64, 10, 10), (64, 54, 54)),
        ((64, 64, 10), (64, 64, 54)),
        ((64, 128, 10), (64, 32, 54)),
        ((128, 10, 10), (32, 54, 54)),
        ((128, 64, 10), (32, 64, 54)),
        ((128, 128, 10), (32, 32, 54)),
    ]:
        chunk = ds.get_layer("color").get_mag("1").get_view(size=size, offset=offset)
        chunk_data = chunk.read(size=chunk.size)
        assert np.array_equal(
            np.ones(chunk_data.shape, dtype=np.uint8)
            * np.uint8(sum(chunk.global_offset)),
            chunk_data,
        )


def copy_and_transform_job(args: Tuple[View, View, int], name: str, val: int) -> None:
    (source_view, target_view, i) = args
    # This method simply takes the data from the source_view, transforms it and writes it to the target_view

    # These assertions are just to demonstrate how the passed parameters can be accessed inside this method
    assert name == "foo"
    assert val == 42

    # increment the color value of each voxel
    data = source_view.read(size=source_view.size)
    if data.shape[0] == 1:
        data = data[0, :, :, :]
    data += 50
    target_view.write(data)


def get_multichanneled_data(dtype: type) -> np.ndarray:
    data = np.zeros((3, 250, 200, 10), dtype=dtype)
    for h in range(10):
        for i in range(250):
            for j in range(200):
                data[0, i, j, h] = i * 256
                data[1, i, j, h] = j * 256
                data[2, i, j, h] = 100 * 256
    return data


def test_create_dataset_with_layer_and_mag() -> None:
    delete_dir(TESTOUTPUT_DIR / "wk_dataset")

    ds = Dataset.create(TESTOUTPUT_DIR / "wk_dataset", scale=(1, 1, 1))
    ds.add_layer("color", "color")

    ds.get_layer("color").add_mag("1")
    ds.get_layer("color").add_mag("2-2-1")

    assert (TESTOUTPUT_DIR / "wk_dataset" / "color" / "1").exists()
    assert (TESTOUTPUT_DIR / "wk_dataset" / "color" / "2-2-1").exists()

    assert len(ds.properties.data_layers) == 1
    assert len(ds.properties.data_layers["color"].wkw_magnifications) == 2


def test_create_dataset_with_explicit_header_fields() -> None:
    delete_dir(TESTOUTPUT_DIR / "wk_dataset_advanced")

    ds = Dataset.create(TESTOUTPUT_DIR / "wk_dataset_advanced", scale=(1, 1, 1))
    ds.add_layer(
        "color", LayerCategories.COLOR_TYPE, dtype_per_layer="uint48", num_channels=3
    )

    ds.get_layer("color").add_mag("1", block_len=64, file_len=64)
    ds.get_layer("color").add_mag("2-2-1")

    assert (TESTOUTPUT_DIR / "wk_dataset_advanced" / "color" / "1").exists()
    assert (TESTOUTPUT_DIR / "wk_dataset_advanced" / "color" / "2-2-1").exists()

    assert len(ds.properties.data_layers) == 1
    assert len(ds.properties.data_layers["color"].wkw_magnifications) == 2

    assert ds.properties.data_layers["color"].element_class == "uint48"
    assert (
        cast(
            Resolution, ds.properties.data_layers["color"].wkw_magnifications[0]
        ).cube_length
        == 64 * 64
    )  # mag "1"
    assert ds.properties.data_layers["color"].wkw_magnifications[0].mag == Mag("1")
    assert (
        cast(
            Resolution, ds.properties.data_layers["color"].wkw_magnifications[1]
        ).cube_length
        == 32 * 32
    )  # mag "2-2-1" (defaults are used)
    assert ds.properties.data_layers["color"].wkw_magnifications[1].mag == Mag("2-2-1")


def test_open_dataset() -> None:
    ds = Dataset(TESTDATA_DIR / "simple_wk_dataset")

    assert len(ds.properties.data_layers) == 1
    assert len(ds.properties.data_layers["color"].wkw_magnifications) == 1
    
    
def test_modify_existing_dataset() -> None:
    delete_dir(TESTOUTPUT_DIR / "simple_wk_dataset")
    ds1 = Dataset.create(TESTOUTPUT_DIR / "simple_wk_dataset", scale=(1, 1, 1))
    ds1.add_layer(
        "color", LayerCategories.COLOR_TYPE, dtype_per_layer="float", num_channels=1
    )
    
    ds2 = Dataset(TESTOUTPUT_DIR / "simple_wk_dataset")
    ds2.add_layer(
        'segmentation', LayerCategories.SEGMENTATION_TYPE, 'uint8', largest_segment_id=100000
    )
    
    assert((TESTOUTPUT_DIR / "simple_wk_dataset" / "segmentation").is_dir())


def test_view_read_with_open() -> None:
    wk_view = (
        Dataset(TESTDATA_DIR / "simple_wk_dataset")
        .get_layer("color")
        .get_mag("1")
        .get_view(size=(16, 16, 16))
    )

    assert not wk_view._is_opened

    with wk_view.open():
        assert wk_view._is_opened

        data = wk_view.read(size=(10, 10, 10))
        assert data.shape == (3, 10, 10, 10)  # three channel

    assert not wk_view._is_opened


def test_view_read_without_open() -> None:
    wk_view = (
        Dataset(TESTDATA_DIR / "simple_wk_dataset")
        .get_layer("color")
        .get_mag("1")
        .get_view(size=(16, 16, 16))
    )

{'fraction_non_alphanumeric': 0.12541694462975317, 'fraction_numerical': 0.03722481654436291, 'mean_word_length': 2.8816675297773173, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 28, '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': '4587710', 'n_tokens_mistral': 5531, 'n_tokens_neox': 4950, 'n_words': 1147}