import unittest

import numpy as np

from conf.testing import czifile, output_path, monopngfile, rgbpngfile, tifffile, monozstackmask
from model_server.accessors import CziImageFileAccessor, DataShapeError, generate_file_accessor, InMemoryDataAccessor, PngFileAccessor, write_accessor_data_to_file, TifSingleSeriesFileAccessor

class TestCziImageFileAccess(unittest.TestCase):

    def setUp(self) -> None:
        pass

    def test_tiffile_is_correct_shape(self):
        tf = generate_file_accessor(tifffile['path'])

        self.assertIsInstance(tf, TifSingleSeriesFileAccessor)
        self.assertEqual(tf.shape_dict['Y'], tifffile['h'])
        self.assertEqual(tf.shape_dict['X'], tifffile['w'])
        self.assertEqual(tf.chroma, tifffile['c'])
        self.assertTrue(tf.is_3d())
        self.assertEqual(len(tf.data.shape), 4)
        self.assertEqual(tf.shape[0], tifffile['h'])
        self.assertEqual(tf.shape[1], tifffile['w'])

    def test_czifile_is_correct_shape(self):
        cf = CziImageFileAccessor(czifile['path'])
        self.assertEqual(cf.shape_dict['Y'], czifile['h'])
        self.assertEqual(cf.shape_dict['X'], czifile['w'])
        self.assertEqual(cf.chroma, czifile['c'])
        self.assertFalse(cf.is_3d())
        self.assertEqual(len(cf.data.shape), 4)
        self.assertEqual(cf.shape[0], czifile['h'])
        self.assertEqual(cf.shape[1], czifile['w'])

    def test_get_single_channel_from_zstack(self):
        w = 256
        h = 512
        nc = 4
        nz = 11
        c = 3
        cf = InMemoryDataAccessor(np.random.rand(h, w, nc, nz))
        sc = cf.get_one_channel_data(c)
        self.assertEqual(sc.shape, (h, w, 1, nz))

    def test_write_single_channel_tif(self):
        ch = 4
        cf = CziImageFileAccessor(czifile['path'])
        mono = cf.get_one_channel_data(ch)
        self.assertTrue(
            write_accessor_data_to_file(
                output_path / f'{cf.fpath.stem}_ch{ch}.tif',
                mono
            )
        )
        self.assertEqual(cf.data.shape[0:2], mono.data.shape[0:2])
        self.assertEqual(cf.data.shape[3], mono.data.shape[2])

    def test_conform_data_shorter_than_xycz(self):
        h = 256
        w = 512
        data = np.random.rand(h, w, 1)
        acc = InMemoryDataAccessor(data)
        self.assertEqual(
            InMemoryDataAccessor.conform_data(data).shape,
            (256, 512, 1, 1)
        )
        self.assertEqual(
            acc.shape_dict,
            {'Y': 256, 'X': 512, 'C': 1, 'Z': 1}
        )

    def test_conform_data_longer_than_xycz(self):
        data = np.random.rand(256, 512, 12, 8, 3)
        with self.assertRaises(DataShapeError):
            acc = InMemoryDataAccessor(data)


    def test_write_multichannel_image_preserve_axes(self):
        h = 256
        w = 512
        c = 3
        nz = 10

        yxcz = (2**8 * np.random.rand(h, w, c, nz)).astype('uint8')
        acc = InMemoryDataAccessor(yxcz)
        fp = output_path / f'rand3d.tif'
        self.assertTrue(
            write_accessor_data_to_file(fp, acc)
        )
        # need to sort out x,y flipping since np convention yxcz flips axes in 3d tif
        self.assertEqual(acc.shape_dict['X'], w, acc.shape_dict)
        self.assertEqual(acc.shape_dict['Y'], h, acc.shape_dict)

        # re-open file and check axes order
        from tifffile import TiffFile
        fh = TiffFile(fp)
        self.assertEqual(len(fh.series), 1)
        se = fh.series[0]
        fh_shape_dict = {se.axes[i]: se.shape[i] for i in range(0, len(se.shape))}
        self.assertEqual(fh_shape_dict, acc.shape_dict, 'Axes are not preserved in TIF output')

    def test_read_png(self, pngfile=rgbpngfile):
        acc = PngFileAccessor(pngfile['path'])
        self.assertEqual(acc.hw, (pngfile['h'], pngfile['w']))
        self.assertEqual(acc.chroma, pngfile['c'])
        self.assertEqual(acc.nz, 1)

    def test_read_mono_png(self):
        return self.test_read_png(pngfile=monopngfile)

    def test_read_zstack_mono_mask(self):
        acc = generate_file_accessor(monozstackmask['path'])
        self.assertTrue(acc.is_mask())

    def test_read_in_pixel_scale_from_czi(self):
        cf = CziImageFileAccessor(czifile['path'])
        pxs = cf.pixel_scale_in_micrometers
        self.assertAlmostEqual(pxs['X'], czifile['um_per_pixel'], places=3)