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.products import export_patches_from_zstack
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', **kwargs):
        zmask, meta, df, interm = build_zmask_from_object_mask(
            self.obmap.get_one_channel_data(0),
            self.stack.get_one_channel_data(0),
            mask_type=mask_type,
            **kwargs,
        )
        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)

        # assert dimensionality of intermediate data products
        self.assertEqual(interm['label_map'].shape, zmask.shape[0:2])
        self.assertEqual(interm['argmax'].shape, zmask.shape[0:2])

        return zmask, meta

    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))

    def test_make_2d_patches_from_zmask(self):
        zmask, meta = self.test_zmask_makes_correct_boxes(
            filters={'area': (1e3, 1e4)},
            expand_box_by=(64, 2)
        )
        files = export_patches_from_zstack(
            output_path / '2d_patches',
            self.stack.get_one_channel_data(channel=1),
            meta,
            draw_bounding_box=True,
        )
        self.assertGreaterEqual(len(files), 1)

    def test_make_3d_patches_from_zmask(self):
        zmask, meta = self.test_zmask_makes_correct_boxes(
            filters={'area': (1e3, 1e4)},
            expand_box_by=(64, 2),
        )
        files = export_patches_from_zstack(
            output_path / '3d_patches',
            self.stack.get_one_channel_data(channel=1),
            meta,
            make_3d=True)
        self.assertGreaterEqual(len(files), 1)

    def test_flatten_image(self):
        zmask, meta, df, interm = build_zmask_from_object_mask(
            self.obmap.get_one_channel_data(0),
            self.stack.get_one_channel_data(1),
            mask_type='boxes',
        )

        from extensions.chaeo.zmask import project_stack_from_focal_points

        dff = df[df['keeper'] == True]

        img = project_stack_from_focal_points(
            dff['centroid-0'].to_numpy(),
            dff['centroid-1'].to_numpy(),
            dff['zi'].to_numpy(),
            self.stack.get_one_channel_data(1),
        )

        self.assertEqual(img.shape[0:2], self.stack.shape[0:2])

        write_accessor_data_to_file(
            output_path / 'flattened.tif',
            InMemoryDataAccessor(img)
        )