import os
import unittest
import warnings

import numpy as np
import vtk
from ccpi.viewer.utils.conversion import Converter, calculate_target_downsample_shape, \
    cilRawResampleReader, cilTIFFResampleReader, cilNumpyMETAImageWriter, cilMetaImageResampleReader,\
    cilNumpyResampleReader


class TestResampleReaders(unittest.TestCase):

    def setUp(self):
        # Generate random 3D array and write to HDF5:
        np.random.seed(1)
        bits = 16
        self.bytes_per_element = int(bits / 8)
        shape = (5, 10, 6)
        self.size_to_resample_to = 100
        self.size_greater_than_input_size = 10000
        self.input_3D_array = np.random.randint(10, size=shape, dtype=eval(f"np.uint{bits}"))
        self.input_3D_array = np.reshape(np.arange(self.input_3D_array.size),
                                         newshape=shape).astype(dtype=eval(f"np.uint{bits}"))
        bytes_3D_array = bytes(self.input_3D_array)
        self.raw_filename_3D = 'test_3D_data.raw'
        with open(self.raw_filename_3D, 'wb') as f:
            f.write(bytes_3D_array)

        # write header to go with raw file
        self.mhd_filename_3D = 'raw_header.mhd'
        typecode = f'uint{bits}'
        big_endian = False
        header_length = 0
        shape = np.shape(self.input_3D_array)
        shape_to_write = shape[::-1]  # because it is not a fortran order array we have to swap
        cilNumpyMETAImageWriter.WriteMETAImageHeader(self.raw_filename_3D,
                                                     self.mhd_filename_3D,
                                                     typecode,
                                                     big_endian,
                                                     header_length,
                                                     shape_to_write,
                                                     spacing=(1., 1., 1.),
                                                     origin=(0., 0., 0.))

        self.numpy_filename_3D = 'test_3D_data.npy'
        np.save(self.numpy_filename_3D, self.input_3D_array)

        self.meta_filename_3D = 'test_3D_data.mha'
        vtk_image = Converter.numpy2vtkImage(self.input_3D_array)
        writer = vtk.vtkMetaImageWriter()
        writer.SetFileName(self.meta_filename_3D)
        writer.SetInputData(vtk_image)
        writer.SetCompression(False)
        writer.Write()

        # Create TIFF Files
        fnames = []
        arr = self.input_3D_array
        from PIL import Image
        for i in range(arr.shape[0]):
            fname = 'tiff_test_file_{:03d}.tiff'.format(i)
            fnames.append(os.path.abspath(fname))
            # Using vtk the Y axis gets reversed
            # vtk_image = Converter.numpy2vtkImage(np.expand_dims(arr[i,:,:], axis=0))
            # twriter.SetFileName(fnames[-1])
            # twriter.SetInputData(vtk_image)
            # twriter.Write()
            im = Image.fromarray(arr[i])
            im.save(fnames[-1])

        self.tiff_fnames = fnames

    def resample_reader_test1(self, reader, target_size, expected_array=None):
        # Tests image with correct target size is generated by resample reader:
        # Not a great test, but at least checks the resample reader runs
        # without crashing
        # TODO: improve this test

        if expected_array is None:
            expected_array = self.input_3D_array

        reader.SetTargetSize(target_size)
        reader.Modified()
        reader.Update()
        og_shape = np.shape(expected_array)
        raw_type_code = str(expected_array.dtype)

        if target_size < self.input_3D_array.size * self.bytes_per_element:
            # Check raw type code was set correctly:
            self.assertEqual(raw_type_code, reader.GetTypeCodeName())
            image = reader.GetOutput()
            extent = image.GetExtent()
            resulting_shape = (extent[1] + 1, (extent[3] + 1), (extent[5] + 1))
            og_shape = (og_shape[2], og_shape[1], og_shape[0])
            og_size = og_shape[0] * og_shape[1] * og_shape[2] * self.bytes_per_element
            expected_shape = calculate_target_downsample_shape(target_size, og_size, og_shape)
            self.assertEqual(resulting_shape, expected_shape)
            # Now test if we get the correct z extent if we set that we
            # have acquisition data
            reader.SetIsAcquisitionData(True)
            # why do we need this?
            reader.Modified()
            reader.Update()
            image = reader.GetOutput()
            extent2 = image.GetExtent()
            extent = extent2
            shape_acquisition = calculate_target_downsample_shape(target_size, og_size, og_shape, acq=True)
            expected_size = shape_acquisition[0] * \
                shape_acquisition[1]*shape_acquisition[2]
            resulting_shape = (extent[1] + 1, (extent[3] + 1), (extent[5] + 1))
            resulting_size = resulting_shape[0] * \
                resulting_shape[1]*resulting_shape[2]
            # angle (z direction) is first index in numpy array, and in cil
            # but it is the last in vtk.
            resulting_z_shape = extent[5] + 1
            og_z_shape = og_shape[2]
            self.assertEqual(resulting_size, expected_size)
            self.assertEqual(resulting_z_shape, og_z_shape)
        else:
            reader.Modified()
            reader.Update()
            image = reader.GetOutput()
            resulting_shape = image.GetDimensions()
            expected_shape = np.shape(expected_array)
            self.assertEqual(resulting_shape, expected_shape[::-1])
            resulting_array = Converter.vtk2numpy(image)
            np.testing.assert_array_equal(np.asfortranarray(expected_array), resulting_array)

    def test_raw_resample_reader(self):
        og_shape = np.shape(self.input_3D_array)
        reader = cilRawResampleReader()
        reader.SetFileName(self.raw_filename_3D)
        reader.SetBigEndian(False)
        reader.SetIsFortran(False)
        raw_type_code = str(self.input_3D_array.dtype)
        reader.SetTypeCodeName(raw_type_code)
        reader.SetStoredArrayShape(og_shape)
        self.resample_reader_test1(reader, self.size_to_resample_to)

    def test_raw_resample_reader_when_resampling_not_needed(self):
        og_shape = np.shape(self.input_3D_array)
        reader = cilRawResampleReader()
        reader.SetFileName(self.raw_filename_3D)
        reader.SetBigEndian(False)
        reader.SetIsFortran(False)
        raw_type_code = str(self.input_3D_array.dtype)
        reader.SetTypeCodeName(raw_type_code)
        reader.SetStoredArrayShape(og_shape)
        self.resample_reader_test1(reader, self.size_greater_than_input_size)

    def _setup_tiff_resample_reader(self):
        reader = cilTIFFResampleReader()
        reader.SetFileName(self.tiff_fnames)
        return reader

    def test_tiff_resample_reader(self):
        reader = self._setup_tiff_resample_reader()
        self.resample_reader_test1(reader, self.size_to_resample_to)

    def test_tiff_resample_reader_with_orientation_type_set(self):
        reader = self._setup_tiff_resample_reader()
        reader.SetOrientationType(4)  # this flips the y axis
        expected_array = np.flip(np.copy(self.input_3D_array), axis=1)
        self.resample_reader_test1(reader, self.size_to_resample_to, expected_array)

    def test_tiff_resample_reader_when_resampling_not_needed(self):
        reader = self._setup_tiff_resample_reader()
        self.resample_reader_test1(reader, self.size_greater_than_input_size)

    def test_tiff_resample_reader_with_orientation_type_set_when_resampling_not_needed(self):
        reader = self._setup_tiff_resample_reader()
        reader.SetOrientationType(4)  # this flips the y axis
        expected_array = np.flip(np.copy(self.input_3D_array), axis=1)
        self.resample_reader_test1(reader, self.size_greater_than_input_size, expected_array)

    def test_meta_resample_reader_mha(self):
        reader = cilMetaImageResampleReader()
        reader.SetFileName(self.meta_filename_3D)
        self.resample_reader_test1(reader, self.size_to_resample_to)

    def test_meta_resample_reader_mha_when_resampling_not_needed(self):
        reader = cilMetaImageResampleReader()
        reader.SetFileName(self.meta_filename_3D)
        self.resample_reader_test1(reader, self.size_greater_than_input_size)

    def test_meta_resample_reader_mhd(self):
        reader = cilMetaImageResampleReader()
        reader.SetFileName(self.mhd_filename_3D)
        self.resample_reader_test1(reader, self.size_to_resample_to)

    def test_meta_resample_reader_mhd_when_resampling_not_needed(self):
        reader = cilMetaImageResampleReader()
        reader.SetFileName(self.mhd_filename_3D)
        self.resample_reader_test1(reader, self.size_greater_than_input_size)

    def test_npy_resample_reader(self):
        reader = cilNumpyResampleReader()
        reader.SetFileName(self.numpy_filename_3D)
        self.resample_reader_test1(reader, self.size_to_resample_to)
        self.resample_reader_test1(reader, self.size_greater_than_input_size)

    def test_npy_resample_reader_when_resampling_not_needed(self):
        reader = cilNumpyResampleReader()
        reader.SetFileName(self.numpy_filename_3D)
        self.resample_reader_test1(reader, self.size_to_resample_to)
        self.resample_reader_test1(reader, self.size_greater_than_input_size)

    def tearDown(self):
        files = [self.raw_filename_3D, self.numpy_filename_3D, self.meta_filename_3D
                 ] + self.tiff_fnames + [self.mhd_filename_3D]
        for f in files:
            # print (f'removing {f}')
            os.remove(f)
            if os.path.exists(f):
                warnings.warn(f'file {f} not deleted!')


if __name__ == '__main__':
    unittest.main()
