from __future__ import division, print_function

import unittest

import os

import dicom

import numpy as np

from rvseg import patient

class TestPatientData(unittest.TestCase):

    def setUp(self):

        # test data has frames 0,1,..21, with the 20th frame labeled

        self.directory = "../test-assets/patient09/"

        self.dicom_path = self.directory + "P09dicom"

    def test_patient_data(self):

        p = patient.PatientData(self.directory)

        self.assertEqual(p.index, 9)

        self.assertEqual(p.dicom_path, self.dicom_path)

        self.assertEqual(len(p.images), 3)

        self.assertEqual(len(p.dicoms), 3)

        self.assertEqual(len(p.all_images), 24)

        self.assertEqual(len(p.all_dicoms), 24)

        self.assertEqual(p.image_width, 256)

        self.assertEqual(p.image_height, 216)

        self.assertEqual(p.rotated, True)

        self.assertEqual(p.labeled, [20, 22, 23])

        self.assertEqual(len(p.endocardium_masks), 3)

        self.assertEqual(len(p.epicardium_masks), 3)

        self.assertEqual(len(p.endocardium_contours), 3)

        self.assertEqual(len(p.epicardium_contours), 3)

        # check dicom MRI image

        plan = dicom.read_file(self.directory + "P09dicom/P09-0000.dcm")

        np.testing.assert_array_equal(p.all_dicoms[0].pixel_array,

                                      plan.pixel_array)

        self.assertEqual(p.images[0].dtype, 'uint16')

        # check endo- and epicardium masks

        endo_mask = np.loadtxt(self.directory + "endocardium-p09-0020.mask")

        np.testing.assert_array_equal(p.endocardium_masks[0], endo_mask)

        self.assertEqual(p.endocardium_masks[0].dtype, 'uint8')

        self.assertSetEqual(set(p.endocardium_masks[0].flatten()), set([0, 1]))

        epi_mask = np.loadtxt(self.directory + "epicardium-p09-0020.mask")

        np.testing.assert_array_equal(p.epicardium_masks[0], epi_mask)

        self.assertEqual(p.epicardium_masks[0].dtype, 'uint8')

        self.assertSetEqual(set(p.epicardium_masks[0].flatten()), set([0, 1]))

    @unittest.skip("Skipping video write: OpenCV is a pain to install in test env.")

    def test_write_video(self):

        p = patient.PatientData(self.directory)

        outfile = "test_write_video.mp4"

        path = os.path.join(os.getcwd(), outfile)

        p.write_video(path)

        # simply check file exists and isn't empty

        self.assertTrue(os.path.isfile(path))

        self.assertGreater(os.path.getsize(path), 0)

        os.remove(path)