from neuroseg.augmentations.brightness import Brightness
from neuroseg.augmentations.occlusion import Occlusion
from neuroseg.augmentations.duplicate import Duplicate
import unittest
from neuroseg.datasets.filetypes import TiffVolume
from neuroseg.datasets.dataset import AlignedVolume
import tifffile as tif
import os.path
import pytest
from neuroseg.datasets.datatypes import BoundingBox, Vector

IMAGE_PATH = "./tests/images/"


class TestAugmentations(unittest.TestCase):
    def test_brightness(self):
        input_dataset = TiffVolume(os.path.join(IMAGE_PATH,
                                                "inputs.tif"),
                                   BoundingBox(Vector(0, 0, 0),
                                               Vector(1024, 512, 50)))
        label_dataset = TiffVolume(os.path.join(IMAGE_PATH,
                                                "labels.tif"),
                                   BoundingBox(Vector(0, 0, 0),
                                               Vector(1024, 512, 50)))
        input_dataset.__enter__()
        label_dataset.__enter__()
        training_dataset = AlignedVolume((input_dataset, label_dataset),
                                         iteration_size=BoundingBox(Vector(0, 0, 0), Vector(128, 128, 20)),
                                         stride=Vector(128, 128, 20))

        brightness_dataset = Brightness(training_dataset)

        tif.imsave(os.path.join(IMAGE_PATH, "test_brightness_input.tif"),
                   brightness_dataset[10][0].getArray())
        tif.imsave(os.path.join(IMAGE_PATH, "test_brightness_label.tif"),
                   brightness_dataset[10][1].getArray()*255)

    def test_occlusion(self):
        input_dataset = TiffVolume(os.path.join(IMAGE_PATH,
                                                "inputs.tif"),
                                   BoundingBox(Vector(0, 0, 0),
                                               Vector(1024, 512, 50)))
        label_dataset = TiffVolume(os.path.join(IMAGE_PATH,
                                                "labels.tif"), 
                                   BoundingBox(Vector(0, 0, 0),
                                               Vector(1024, 512, 50)))
        input_dataset.__enter__()
        label_dataset.__enter__()
        training_dataset = AlignedVolume((input_dataset, label_dataset),
                                         iteration_size=BoundingBox(Vector(0, 0, 0), Vector(128, 128, 20)),
                                         stride=Vector(128, 128, 20))

        occlusion_dataset = Occlusion(training_dataset)

        tif.imsave(os.path.join(IMAGE_PATH, "test_occlusion_input.tif"),
                   occlusion_dataset[10][0].getArray())
        tif.imsave(os.path.join(IMAGE_PATH, "test_occlusion_label.tif"),
                   occlusion_dataset[10][1].getArray()*255)

    def test_duplication(self):
        input_dataset = TiffVolume(os.path.join(IMAGE_PATH,
                                                "inputs.tif"),
                                   BoundingBox(Vector(0, 0, 0),
                                               Vector(1024, 512, 50)))
        
        label_dataset = TiffVolume(os.path.join(IMAGE_PATH,
                                                "labels.tif"),
                                   BoundingBox(Vector(0, 0, 0),
                                               Vector(1024, 512, 50)))
        input_dataset.__enter__()
        label_dataset.__enter__()
        training_dataset = AlignedVolume((input_dataset, label_dataset),
                                         iteration_size=BoundingBox(Vector(0, 0, 0), Vector(128, 128, 20)),
                                         stride=Vector(128, 128, 20))

        duplicate_dataset = Duplicate(training_dataset)

        tif.imsave(os.path.join(IMAGE_PATH, "test_duplicate_input.tif"),
                   duplicate_dataset[10][0].getArray())
        tif.imsave(os.path.join(IMAGE_PATH, "test_duplicate_label.tif"),
                   duplicate_dataset[10][1].getArray()*255)
