#!/usr/bin/env python3

import os
import lab
import pickle
import hashlib
import unittest
import collections

TEST_DIRECTORY = os.path.dirname(__file__)


def object_hash(x):
    return hashlib.sha512(pickle.dumps(x)).hexdigest()

class Lab1Test(unittest.TestCase):
    # This class is not really a test; rather, it provides useful functions
    # we wish to share across the other test classes.
    def compare_greyscale_images(self, im1, im2):
        self.assertTrue(set(im1.keys()) == {'height', 'width', 'pixels'}, 'Incorrect keys in dictionary')
        self.assertEqual(im1['height'], im2['height'], 'Heights must match')
        self.assertEqual(im1['width'], im2['width'], 'Widths must match')
        self.assertEqual(len(im1['pixels']), im1['height']*im1['width'], 'Incorrect number of pixels')
        pix_incorrect = (None, None)
        for ix, (i, j) in enumerate(zip(im1['pixels'], im2['pixels'])):
            if i != j:
                pix_incorrect = (ix, abs(i-j))
        self.assertTrue(pix_incorrect == (None, None), 'Pixels must match.  Incorrect value at location %s (differs from expected by %s)' % pix_incorrect)

    def compare_color_images(self, im1, im2):
        self.assertTrue(set(im1.keys()) == {'height', 'width', 'pixels'}, 'Incorrect keys in dictionary')
        self.assertEqual(im1['height'], im2['height'], 'Heights must match')
        self.assertEqual(im1['width'], im2['width'], 'Widths must match')
        self.assertEqual(len(im1['pixels']), im1['height']*im1['width'], 'Incorrect number of pixels')
        self.assertTrue(all(isinstance(i, tuple) and len(i)==3 for i in im1['pixels']), 'Pixels must all be 3-tuples')
        self.assertTrue(all(0<=subi<=255 for i in im1['pixels'] for subi in i), 'Pixels values must all be in the range from [0, 255]')
        pix_incorrect = (None, None)
        for ix, (i, j) in enumerate(zip(im1['pixels'], im2['pixels'])):
            if i != j:
                tup_diff = tuple(abs(i[t]-j[t]) for t in {0,1,2})
                pix_incorrect = (ix, tup_diff)
        self.assertTrue(pix_incorrect == (None, None), f'Pixels must match.  Incorrect value at location %s (differs from expected by %s)' % pix_incorrect)

class TestSeamCarvingHelpers(Lab1Test):
    # These tests are not checked on the server; they are only
    # provided for debugging help, if you wish to use our 
    # suggested helper functions.
    def test_cumulative_energy(self):
        for fname in ('pattern', 'smallfrog', 'bluegill', 'twocats', 'tree'):
            infile = os.path.join(TEST_DIRECTORY, 'test_results', f'{fname}_energy.pickle')
            with open(infile, 'rb') as f:
                energy = pickle.load(f)
            result = lab.cumulative_energy_map(energy)

            expfile = os.path.join(TEST_DIRECTORY, 'test_results', f'{fname}_cumulative_energy.pickle')
            with open(expfile, 'rb') as f:
                cem = pickle.load(f)

            self.compare_greyscale_images(result, cem)

    def test_min_seam_indices(self):
        for fname in ('pattern', 'smallfrog', 'bluegill', 'twocats', 'tree'):
            infile = os.path.join(TEST_DIRECTORY, 'test_results', f'{fname}_cumulative_energy.pickle')
            with open(infile, 'rb') as f:
                cem = pickle.load(f)
            result = lab.minimum_energy_seam(cem)

            expfile = os.path.join(TEST_DIRECTORY, 'test_results', f'{fname}_minimum_energy_seam.pickle')
            with open(expfile, 'rb') as f:
                seam = pickle.load(f)

            self.assertEqual(len(result), len(seam))
            self.assertEqual(set(result), set(seam))

    def test_seam_removal(self):
        for fname in ('pattern', 'bluegill', 'twocats', 'tree'):
            infile = os.path.join(TEST_DIRECTORY, 'test_results', f'{fname}_minimum_energy_seam.pickle')
            with open(infile, 'rb') as f:
                seam = pickle.load(f)

            imfile = os.path.join(TEST_DIRECTORY, 'test_images', f'{fname}.png')
            result = lab.image_without_seam(lab.load_color_image(imfile), seam)

            expfile = os.path.join(TEST_DIRECTORY, 'test_results', f'{fname}_1seam.png')

            self.compare_color_images(result, lab.load_color_image(expfile))


class TestSeamCarving(Lab1Test):
    # These tests are those checked by the server to determine
    # your score.
    def test_endtoend_centeredpixel(self):
        inpfile = os.path.join(TEST_DIRECTORY, 'test_images', 'centered_pixel.png')

        im = lab.load_color_image(inpfile)
        oim = object_hash(im)

        for i in range(1, 11):
            result = lab.seam_carving(im, i)
            self.assertEqual(object_hash(im), oim, 'Be careful not to modify the original image!')

            expfile = os.path.join(TEST_DIRECTORY, 'test_results', 'seams_centered_pixel', f'{i:02d}.png')
            self.compare_color_images(result, lab.load_color_image(expfile))


    def test_endtoend_pattern(self):
        inpfile = os.path.join(TEST_DIRECTORY, 'test_images', 'pattern.png')

        im = lab.load_color_image(inpfile)
        oim = object_hash(im)

        for i in range(1, 9):
            result = lab.seam_carving(im, i)
            self.assertEqual(object_hash(im), oim, 'Be careful not to modify the original image!')

            expfile = os.path.join(TEST_DIRECTORY, 'test_results', 'seams_pattern', f'{i:02d}.png')
            self.compare_color_images(result, lab.load_color_image(expfile))


    def test_endtoend_smallfrog(self):
        inpfile = os.path.join(TEST_DIRECTORY, 'test_images', 'smallfrog.png')

        im = lab.load_color_image(inpfile)
        oim = object_hash(im)

        for i in range(1, 31):
            result = lab.seam_carving(im, i)
            self.assertEqual(object_hash(im), oim, 'Be careful not to modify the original image!')

            expfile = os.path.join(TEST_DIRECTORY, 'test_results', 'seams_smallfrog', f'{i:02d}.png')
            self.compare_color_images(result, lab.load_color_image(expfile))



if __name__ == '__main__':
    res = unittest.main(verbosity=3, exit=False)
