# (C) British Crown Copyright 2014 - 2016, Met Office
#
# This file is part of Iris.
#
# Iris is free software: you can redistribute it and/or modify it under
# the terms of the GNU Lesser General Public License as published by the
# Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Iris is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with Iris.  If not, see <http://www.gnu.org/licenses/>.

from __future__ import (absolute_import, division, print_function)
from six.moves import (filter, input, map, range, zip)  # noqa
import six

import iris.tests as tests
import iris

import numpy as np
import PIL.Image


@tests.skip_gdal
@tests.skip_data
class TestGeoTiffExport(tests.IrisTest):
    def check_tiff_header(self, tiff_filename, expect_keys, expect_entries):
        """
        Checks the given tiff file's metadata contains the expected keys,
        and some matching values (not all).

        """
        with open(tiff_filename, 'rb') as fh:
            im = PIL.Image.open(fh)
            file_keys = im.tag.keys()

            missing_keys = sorted(set(expect_keys) - set(file_keys))
            msg_nokeys = "Tiff header has missing keys : {}."
            self.assertEqual(missing_keys, [],
                             msg_nokeys.format(missing_keys))

            extra_keys = sorted(set(file_keys) - set(expect_keys))
            msg_extrakeys = "Tiff header has extra unexpected keys : {}."
            self.assertEqual(extra_keys, [],
                             msg_extrakeys.format(extra_keys))

            msg_badval = "Tiff header entry {} has value {} != {}."
            for key, value in expect_entries.items():
                content = im.tag[key]
                self.assertEqual(content, value,
                                 msg_badval.format(key, content, value))

    def check_tiff(self, cube, header_keys, header_items):
        # Check that the cube saves correctly to TIFF :
        #   * the header contains expected keys and (some) values
        #   * the data array retrives correctly
        import iris.experimental.raster
        with self.temp_filename('.tif') as temp_filename:
            iris.experimental.raster.export_geotiff(cube, temp_filename)

            # Check the metadata is correct.
            self.check_tiff_header(temp_filename, header_keys, header_items)

            # Ensure that north is at the top then check the data is correct.
            coord_y = cube.coord(axis='Y', dim_coords=True)
            data = cube.data
            if np.diff(coord_y.bounds[0]) > 0:
                data = cube.data[::-1, :]
            im = PIL.Image.open(temp_filename)
            im_data = np.array(im)
            # Currently we only support writing 32-bit tiff, when comparing
            # the data ensure that it is also 32-bit
            np.testing.assert_array_equal(im_data,
                                          data.astype(np.float32))

    def _check_tiff_export(self, masked, inverted=False):
        tif_header = 'SMALL_total_column_co2.nc.tif_header.txt'
        tif_header_keys = [256, 257, 258, 259, 262, 273,
                           277, 278, 279, 284, 339, 33550, 33922]
        tif_header_entries = {
            256: (160,),
            257: (159,),
            258: (32,),
            259: (1,),
            262: (1,),
            # Skip this one: behaviour is not consistent across gdal versions.
            # 273: (354, 8034, 15714, 23394, 31074, 38754, 46434,
            #       54114, 61794, 69474, 77154, 84834, 92514, 100194),
            277: (1,),
            278: (12,),
            279: (7680, 7680, 7680, 7680, 7680, 7680, 7680,
                  7680, 7680, 7680, 7680, 7680, 7680, 1920),
            284: (1,),
            339: (3,),
            33550: (1.125, 1.125, 0.0),
            33922: (0.0, 0.0, 0.0, -0.5625, 89.4375, 0.0)
        }
        fin = tests.get_data_path(('NetCDF', 'global', 'xyt',
                                   'SMALL_total_column_co2.nc'))
        cube = iris.load_cube(fin)[0]
        # PIL doesn't support float64
        cube.data = cube.data.astype('f4')

        # Ensure longitude values are continuous and monotonically increasing,
        # and discard the 'half cells' at the top and bottom of the UM output
        # by extracting a subset.
        east = iris.Constraint(longitude=lambda cell: cell < 180)
        non_edge = iris.Constraint(latitude=lambda cell: -90 < cell < 90)
        cube = cube.extract(east & non_edge)
        cube.coord('longitude').guess_bounds()
        cube.coord('latitude').guess_bounds()

        if masked:
            # Mask some of the data + expect a slightly different header...
            cube.data = np.ma.masked_where(cube.data <= 380, cube.data)

            # There is an additional key..
            tif_header_keys += [42113]
            # Don't add a check entry for this, as coding changed between gdal
            # version 1 and 2, *and* between Python2 and Python3.
            # tif_header_entries[42113] = (u'1e+20',)

        if inverted:
            # Check with the latitude coordinate (and the corresponding
            # cube.data) inverted.
            # The output should be exactly the same.
            coord = cube.coord('latitude')
            coord.points = coord.points[::-1]
            coord.bounds = None
            coord.guess_bounds()
            cube.data = cube.data[::-1, :]

        self.check_tiff(cube, tif_header_keys, tif_header_entries)

    def test_unmasked(self):
        self._check_tiff_export(masked=False)

    def test_masked(self):
        self._check_tiff_export(masked=True)

    def test_inverted(self):
        self._check_tiff_export(masked=False, inverted=True)


if __name__ == "__main__":
    tests.main()
