import unittest

import numpy as np

from conf.testing import output_path

from extensions.chaeo.conf.testing import multichannel_zstack, pixel_classifier, pipeline_params
from extensions.chaeo.zmask import build_zmask_from_object_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.stack = generate_file_accessor(multichannel_zstack['path'])

        pxmodel = IlastikPixelClassifierModel(
            {'project_file': pixel_classifier['path']},
        )
        mip = InMemoryDataAccessor(self.stack.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'])
        # write_accessor_data_to_file(output_path / 'obmap.tif', self.obmap)

    def test_zmask_makes_correct_boxes(self, mask_type='boxes', filters=None, expand_box_by=None):
        zmask, meta = build_zmask_from_object_mask(
            self.obmap.get_one_channel_data(0),
            self.stack.get_one_channel_data(0),
            mask_type=mask_type,
            filters=filters,
            expand_box_by=expand_box_by,
        )
        zmask_acc = InMemoryDataAccessor(zmask)
        self.assertTrue(zmask_acc.is_mask())

        # assert dimensionality of zmask
        self.assertGreater(zmask_acc.shape_dict['Z'], 1)
        self.assertEqual(zmask_acc.shape_dict['C'], 1)
        write_accessor_data_to_file(output_path / 'zmask.tif', zmask_acc)

        # mask values are not just all True or all False
        self.assertTrue(np.any(zmask))
        self.assertFalse(np.all(zmask))

        # assert non-trivial meta info in boxes
        self.assertGreater(len(meta), 1)
        sh = meta[1]['mask'].shape
        ar = meta[1]['info'].area
        self.assertGreaterEqual(sh[0] * sh[1], ar)

    def test_zmask_makes_correct_contours(self):
        return self.test_zmask_makes_correct_boxes(mask_type='contours')

    def test_zmask_makes_correct_boxes_with_filters(self):
        return self.test_zmask_makes_correct_boxes(filters={'area': (1e3, 1e4)})

    def test_zmask_makes_correct_expanded_boxes(self):
        return self.test_zmask_makes_correct_boxes(expand_box_by=(64, 2))