# (C) British Crown Copyright 2010 - 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 first so that some things can be initialised before importing anything else
import iris.tests as tests

import os
import tempfile

import cf_units
import numpy as np

import iris
import iris.coords
import iris.coord_systems
import iris.fileformats.pp as ff_pp
from iris.fileformats.pp import PPField3
import iris.io
from iris.tests import mock
import iris.tests.pp as pp
import iris.util
import iris.tests.stock as stock


def itab_callback(cube, field, filename):
    cube.add_aux_coord(iris.coords.AuxCoord([field.lbrel], long_name='MOUMHeaderReleaseNumber', units='no_unit')) 
    cube.add_aux_coord(iris.coords.AuxCoord([field.lbexp], long_name='ExperimentNumber(ITAB)', units='no_unit')) 


@tests.skip_data
class TestPPSave(tests.IrisTest, pp.PPTest):
    def test_no_forecast_time(self):
        cube = stock.lat_lon_cube()
        coord = iris.coords.DimCoord(np.array([24], dtype=np.int64),
                                     standard_name='time',
                                     units='hours since epoch')
        cube.add_aux_coord(coord)
        self.assertCML(cube, ['cube_to_pp', 'no_forecast_time.cml'])

        reference_txt_path = tests.get_result_path(('cube_to_pp', 'no_forecast_time.txt'))
        with self.cube_save_test(reference_txt_path, reference_cubes=cube) as temp_pp_path:
            iris.save(cube, temp_pp_path)

    def test_no_forecast_period(self):
        cube = stock.lat_lon_cube()
        # Add a bounded scalar time coord and a forecast_reference_time.
        time_coord = iris.coords.DimCoord(
            10.958333, standard_name='time',
            units='days since 2013-05-10 12:00',
            bounds=[10.916667, 11.0])
        cube.add_aux_coord(time_coord)
        forecast_reference_time = iris.coords.DimCoord(
            2.0, standard_name='forecast_reference_time',
            units='weeks since 2013-05-07')
        cube.add_aux_coord(forecast_reference_time)

        self.assertCML(cube, ['cube_to_pp', 'no_forecast_period.cml'])
        reference_txt_path = tests.get_result_path(('cube_to_pp',
                                                    'no_forecast_period.txt'))
        with self.cube_save_test(reference_txt_path, reference_cubes=cube) as \
                temp_pp_path:
            iris.save(cube, temp_pp_path)

    def test_pp_save_rules(self):
        # Test pp save rules without user rules.

        #read
        in_filename = tests.get_data_path(('PP', 'simple_pp', 'global.pp'))
        cubes = iris.load(in_filename, callback=itab_callback)

        reference_txt_path = tests.get_result_path(('cube_to_pp', 'simple.txt'))
        with self.cube_save_test(reference_txt_path, reference_cubes=cubes) as temp_pp_path:
            iris.save(cubes, temp_pp_path)

    def test_user_pp_save_rules(self):
        # Test pp save rules with user rules.
        
        #create a user rules file
        user_rules_filename = iris.util.create_temp_filename(suffix='.txt')
        try:
            with open(user_rules_filename, "wt") as user_rules_file:
                user_rules_file.write("IF\ncm.standard_name == 'air_temperature'\nTHEN\npp.lbuser[3] = 9222")
            with iris.fileformats.rules._disable_deprecation_warnings():
                iris.fileformats.pp.add_save_rules(user_rules_filename)
            try:
                #read pp
                in_filename = tests.get_data_path(('PP', 'simple_pp', 'global.pp'))
                cubes = iris.load(in_filename, callback=itab_callback)

                reference_txt_path = tests.get_result_path(('cube_to_pp', 'user_rules.txt'))
                with self.cube_save_test(reference_txt_path, reference_cubes=cubes) as temp_pp_path:
                    iris.save(cubes, temp_pp_path)

            finally:
                iris.fileformats.pp.reset_save_rules()
        finally:
            os.remove(user_rules_filename)

    def test_pp_append_singles(self):
        # Test pp append saving - single cubes.
        
        # load 2 arrays of >2D cubes
        cube = stock.simple_pp()
        
        reference_txt_path = tests.get_result_path(('cube_to_pp', 'append_single.txt'))
        with self.cube_save_test(reference_txt_path, reference_cubes=[cube, cube]) as temp_pp_path:
            iris.save(cube, temp_pp_path)                # Create file
            iris.save(cube, temp_pp_path, append=True)   # Append to file

        reference_txt_path = tests.get_result_path(('cube_to_pp', 'replace_single.txt'))
        with self.cube_save_test(reference_txt_path, reference_cubes=cube) as temp_pp_path:
            iris.save(cube, temp_pp_path)                # Create file
            iris.save(cube, temp_pp_path)                # Replace file

    def test_pp_append_lists(self):
        # Test PP append saving - lists of cubes.
        # For each of the first four time-steps in the 4D cube,
        # pull out the bottom two levels.
        cube_4d = stock.realistic_4d()
        cubes = [cube_4d[i, :2, :, :] for i in range(4)]

        reference_txt_path = tests.get_result_path(('cube_to_pp', 'append_multi.txt'))
        with self.cube_save_test(reference_txt_path, reference_cubes=cubes) as temp_pp_path:
            iris.save(cubes[:2], temp_pp_path)
            iris.save(cubes[2:], temp_pp_path, append=True)

        reference_txt_path = tests.get_result_path(('cube_to_pp', 'replace_multi.txt'))
        with self.cube_save_test(reference_txt_path, reference_cubes=cubes[2:]) as temp_pp_path:
            iris.save(cubes[:2], temp_pp_path)
            iris.save(cubes[2:], temp_pp_path)

    def add_coords_to_cube_and_test(self, coord1, coord2):
        # a wrapper for creating arbitrary 2d cross-sections and run pp-saving tests
        dataarray = np.arange(16, dtype='>f4').reshape(4, 4)
        cm = iris.cube.Cube(data=dataarray)

        cm.add_dim_coord(coord1, 0)
        cm.add_dim_coord(coord2, 1)
        
        # TODO: This is the desired line of code...
        # reference_txt_path = tests.get_result_path(('cube_to_pp', '%s.%s.pp.txt' % (coord1.name(), coord2.name())))
        # ...but this is required during the CF change, to maintain the original filename.
        coord1_name = coord1.name().replace("air_", "")
        coord2_name = coord2.name().replace("air_", "") 
        reference_txt_path = tests.get_result_path(('cube_to_pp', '%s.%s.pp.txt' % (coord1_name, coord2_name))) 

        # test with name
        with self.cube_save_test(reference_txt_path, reference_cubes=cm, 
                field_coords=[coord1.name(), coord2.name()]) as temp_pp_path:
            iris.save(cm, temp_pp_path, field_coords=[coord1.name(), coord2.name()])
        # test with coord
        with self.cube_save_test(reference_txt_path, reference_cubes=cm, 
                field_coords=[coord1, coord2]) as temp_pp_path:
            iris.save(cm, temp_pp_path, field_coords=[coord1, coord2])

    def test_non_standard_cross_sections(self):
        #ticket #1037, the five variants being dealt with are
        #    'pressure.latitude',
        #    'depth.latitude',
        #    'eta.latitude',
        #    'pressure.time',
        #    'depth.time',

        f = FakePPEnvironment()

        self.add_coords_to_cube_and_test(
            iris.coords.DimCoord(f.z, long_name='air_pressure', units='hPa', bounds=f.z_bounds),
            iris.coords.DimCoord(f.y, standard_name='latitude', units='degrees', bounds=f.y_bounds, coord_system=f.geog_cs()))
            
        self.add_coords_to_cube_and_test(
            iris.coords.DimCoord(f.z, long_name='depth', units='m', bounds=f.z_bounds),
            iris.coords.DimCoord(f.y, standard_name='latitude', units='degrees', bounds=f.y_bounds, coord_system=f.geog_cs()))
            
        self.add_coords_to_cube_and_test(
            iris.coords.DimCoord(f.z, long_name='eta', units='1', bounds=f.z_bounds),
            iris.coords.DimCoord(f.y, standard_name='latitude', units='degrees', bounds=f.y_bounds, coord_system=f.geog_cs()))
            
        self.add_coords_to_cube_and_test(
            iris.coords.DimCoord(f.z, long_name='air_pressure', units='hPa', bounds=f.z_bounds),
            iris.coords.DimCoord(f.y, standard_name='time', units=cf_units.Unit('days since 0000-01-01 00:00:00', calendar=cf_units.CALENDAR_360_DAY), bounds=f.y_bounds))
            
        self.add_coords_to_cube_and_test(
            iris.coords.DimCoord(f.z, standard_name='depth', units='m', bounds=f.z_bounds),
            iris.coords.DimCoord(f.y, standard_name='time', units=cf_units.Unit('days since 0000-01-01 00:00:00', calendar=cf_units.CALENDAR_360_DAY), bounds=f.y_bounds))

    def test_365_calendar_export(self):
        # test for 365 day calendar export
        cube = stock.simple_pp()
        new_unit = cf_units.Unit('hours since 1970-01-01 00:00:00',
                                  calendar=cf_units.CALENDAR_365_DAY)
        cube.coord('time').units = new_unit
        pp_field = mock.MagicMock(spec=PPField3)
        iris.fileformats.pp._ensure_save_rules_loaded()
        iris.fileformats.pp._save_rules.verify(cube, pp_field)
        self.assertEqual(pp_field.lbtim.ic, 4)

            
class FakePPEnvironment(object):
    ''' fake a minimal PP environment for use in cross-section coords, as in PP save rules '''
    y = [1, 2, 3, 4]
    z = [111, 222, 333, 444]
    y_bounds = [[0.9, 1.1], [1.9, 2.1], [2.9, 3.1], [3.9, 4.1]]
    z_bounds = [[110.9, 111.1], [221.9, 222.1], [332.9, 333.1], [443.9, 444.1]]

    def geog_cs(self):
        """Return a GeogCS for this PPField.

        Returns:
            A GeogCS with the appropriate earth shape, meridian and pole position.
        """
        return iris.coord_systems.GeogCS(6371229.0)


class TestPPSaveRules(tests.IrisTest, pp.PPTest):
    def test_default_coord_system(self):
        GeogCS = iris.coord_systems.GeogCS
        cube = iris.tests.stock.lat_lon_cube()
        reference_txt_path = tests.get_result_path(('cube_to_pp',
                                                    'default_coord_system.txt'))
        # Remove all coordinate systems.
        for coord in cube.coords():
            coord.coord_system = None
        # Ensure no coordinate systems available.
        self.assertIsNone(cube.coord_system(GeogCS))
        self.assertIsNone(cube.coord_system(None))
        with self.cube_save_test(reference_txt_path, reference_cubes=cube) as \
                temp_pp_path:
            # Save cube to PP with no coordinate system.
            iris.save(cube, temp_pp_path)
            pp_cube = iris.load_cube(temp_pp_path)
            # Ensure saved cube has the default coordinate system.
            self.assertIsInstance(pp_cube.coord_system(GeogCS),
                                  iris.coord_systems.GeogCS)
            self.assertIsNotNone(pp_cube.coord_system(None))
            self.assertIsInstance(pp_cube.coord_system(None),
                                  iris.coord_systems.GeogCS)
            self.assertIsNotNone(pp_cube.coord_system())
            self.assertIsInstance(pp_cube.coord_system(),
                                  iris.coord_systems.GeogCS)

    def lbproc_from_pp(self, filename):
        # Gets the lbproc field from the ppfile
        pp_file = iris.fileformats.pp.load(filename)
        field = next(pp_file)
        return field.lbproc

    def test_pp_save_rules(self):
        # Test single process flags
        for _, process_desc in iris.fileformats.pp.LBPROC_PAIRS[1:]:
            # Get basic cube and set process flag manually
            ll_cube = stock.lat_lon_cube()
            ll_cube.attributes["ukmo__process_flags"] = (process_desc,)
    
            # Save cube to pp
            temp_filename = iris.util.create_temp_filename(".pp")
            iris.save(ll_cube, temp_filename)
     
            # Check the lbproc is what we expect
            self.assertEqual(self.lbproc_from_pp(temp_filename),
                             iris.fileformats.pp.lbproc_map[process_desc])

            os.remove(temp_filename)

        # Test mutiple process flags
        multiple_bit_values = ((128, 64), (4096, 1024), (8192, 1024))
        
        # Maps lbproc value to the process flags that should be created
        multiple_map = {sum(bits) : [iris.fileformats.pp.lbproc_map[bit] for bit in bits] for bits in multiple_bit_values}

        for lbproc, descriptions in six.iteritems(multiple_map):
            ll_cube = stock.lat_lon_cube()
            ll_cube.attributes["ukmo__process_flags"] = descriptions
            
            # Save cube to pp
            temp_filename = iris.util.create_temp_filename(".pp")
            iris.save(ll_cube, temp_filename)
            
            # Check the lbproc is what we expect
            self.assertEqual(self.lbproc_from_pp(temp_filename), lbproc)

            os.remove(temp_filename)
            
    def test_lbvc(self):
        cube = stock.realistic_4d_no_derived()[0, :4, ...]
        
        v_coord = iris.coords.DimCoord(standard_name='depth', 
                                       units='m', points=[-5, -10, -15, -20])
        
        cube.remove_coord('level_height')
        cube.remove_coord('sigma')
        cube.remove_coord('surface_altitude')
        cube.add_aux_coord(v_coord, 0)
    
        expected = ([2, 1, -5.0],
                    [2, 2, -10.0],
                    [2, 3, -15.0],
                    [2, 4, -20.0])
                    
        for field, (lbvc, lblev, blev) in zip(fields_from_cube(cube), expected):
            self.assertEqual(field.lbvc, lbvc)
            self.assertEqual(field.lblev, lblev)
            self.assertEqual(field.blev, blev)
        
        
def fields_from_cube(cubes):
    """
    Return an iterator of PP fields generated from saving the given cube(s)
    to a temporary file, and then subsequently loading them again 
    """
    with tempfile.NamedTemporaryFile('w+b', suffix='.pp') as tmp_file:
        if six.PY2:
            fh = tmp_file.file
        else:
            fh = tmp_file
        iris.save(cubes, fh, saver='pp')
        
        # make sure the fh is written to disk, and move it back to the
        # start of the file 
        fh.flush()
        os.fsync(fh)        
        fh.seek(0)
        
        # load in the saved pp fields and check the appropriate metadata
        for field in ff_pp.load(tmp_file.name):
            yield field
            

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