import unittest

from conf.testing import output_path

from extensions.chaeo.conf.testing import multichannel_zstack, pixel_classifier, pipeline_params
from extensions.chaeo.zstack import build_stack_mask
from model_server.accessors import generate_file_accessor, InMemoryDataAccessor, write_accessor_data_to_file
from extensions.ilastik.models import IlastikObjectClassifierModel, IlastikPixelClassifierModel

class TestZStackDerivedDataProducts(unittest.TestCase):

    def setUp(self) -> None:
        # need test data incl obj map
        self.zstack = generate_file_accessor(multichannel_zstack['path'])

        pxmodel = IlastikPixelClassifierModel(
            {'project_file': pixel_classifier['path']},
        )
        mip = InMemoryDataAccessor(self.zstack.get_one_channel_data(channel=0).data.max(axis=-1, keepdims=True))
        self.pxmap, result = pxmodel.infer(mip)

        write_accessor_data_to_file(output_path / 'pxmap.tif', self.pxmap)
        self.obmap = InMemoryDataAccessor((self.pxmap.data > pipeline_params['threshold']).astype('uint8'))
        write_accessor_data_to_file(output_path / 'obmap.tif', self.obmap)

    def test_zmask_makes_correct_boxes(self):
        zmask, meta = build_stack_mask(
            'test_zmask_with boxes',
            self.stack,
            self.obmap,
            mask_type='boxes',
        )
        zmask_acc = InMemoryDataAccessor(zmask)
        self.assertTrue(zmask_acc.is_object_map())

        # assert dimensionality of zmask
        self.assertEqual(zmask.shape_dict['Z'] > 1)
        self.assertEqual(zmask.shape_dict['C'] == 1)

        # assert non-trivial meta info in boxes
        pass

    def test_zmask_makes_correct_contours(self):
        pass