# Source code for nes.nes_formats.wrf_chem_format

    #!/usr/bin/env python

    import numpy as np
    import nes
    from netCDF4 import Dataset
    from mpi4py import MPI
    from copy import deepcopy
    from datetime import datetime

    GLOBAL_ATTRIBUTES_ORDER = [
        'TITLE', 'START_DATE', 'WEST-EAST_GRID_DIMENSION', 'SOUTH-NORTH_GRID_DIMENSION', 'BOTTOM-TOP_GRID_DIMENSION', 'DX',
        'DY', 'GRIDTYPE', 'DIFF_OPT', 'KM_OPT', 'DAMP_OPT', 'DAMPCOEF', 'KHDIF', 'KVDIF', 'MP_PHYSICS', 'RA_LW_PHYSICS',
        'RA_SW_PHYSICS', 'SF_SFCLAY_PHYSICS', 'SF_SURFACE_PHYSICS', 'BL_PBL_PHYSICS', 'CU_PHYSICS', 'SF_LAKE_PHYSICS',
        'SURFACE_INPUT_SOURCE', 'SST_UPDATE', 'GRID_FDDA', 'GFDDA_INTERVAL_M', 'GFDDA_END_H', 'GRID_SFDDA',
        'SGFDDA_INTERVAL_M', 'SGFDDA_END_H', 'WEST-EAST_PATCH_START_UNSTAG', 'WEST-EAST_PATCH_END_UNSTAG',
        'WEST-EAST_PATCH_START_STAG', 'WEST-EAST_PATCH_END_STAG', 'SOUTH-NORTH_PATCH_START_UNSTAG',
        'SOUTH-NORTH_PATCH_END_UNSTAG', 'SOUTH-NORTH_PATCH_START_STAG', 'SOUTH-NORTH_PATCH_END_STAG',
        'BOTTOM-TOP_PATCH_START_UNSTAG', 'BOTTOM-TOP_PATCH_END_UNSTAG', 'BOTTOM-TOP_PATCH_START_STAG',
        'BOTTOM-TOP_PATCH_END_STAG', 'GRID_ID', 'PARENT_ID', 'I_PARENT_START', 'J_PARENT_START', 'PARENT_GRID_RATIO', 'DT',
        'CEN_LAT', 'CEN_LON', 'TRUELAT1', 'TRUELAT2', 'MOAD_CEN_LAT', 'STAND_LON', 'POLE_LAT', 'POLE_LON', 'GMT', 'JULYR',
        'JULDAY', 'MAP_PROJ', 'MMINLU', 'NUM_LAND_CAT', 'ISWATER', 'ISLAKE', 'ISICE', 'ISURBAN', 'ISOILWATER']

    # noinspection DuplicatedCode

    def to_netcdf_wrf_chem(self, path, chunking=False, keep_open=False):
        """
        Create the NetCDF using netcdf4-python methods.

        Parameters
        ----------
        self : nes.Nes
            Source projection Nes Object.
        path : str
            Path to the output netCDF file.
        chunking: bool
            Indicates if you want to chunk the output netCDF.
        keep_open : bool
            Indicates if you want to keep open the NetCDH to fill the data by time-step.
        """

        self.to_dtype(np.float32)

        set_global_attributes(self)
        change_variable_attributes(self)

        # Open NetCDF
        if self.info:
            print("Rank {0:03d}: Creating {1}".format(self.rank, path))
        if self.size > 1:
            netcdf = Dataset(path, format="NETCDF4", mode='w', parallel=True, comm=self.comm, info=MPI.Info())
        else:
            netcdf = Dataset(path, format="NETCDF4", mode='w', parallel=False)
        if self.info:
            print("Rank {0:03d}: NetCDF ready to write".format(self.rank))

        # Create dimensions
        create_dimensions(self, netcdf)

        create_dimension_variables(self, netcdf)
        if self.info:
            print("Rank {0:03d}: Dimensions done".format(self.rank))

        # Create variables
        create_variables(self, netcdf)

        for att_name in GLOBAL_ATTRIBUTES_ORDER:
            netcdf.setncattr(att_name, self.global_attrs[att_name])

        # Close NetCDF
        if keep_open:
            self.netcdf = netcdf
        else:
            netcdf.close()

        return None

    def change_variable_attributes(self):
        """
        Modify the emission list to be consistent to use the output as input for WRF-CHEM model.

        Parameters
        ----------
        self : nes.Nes
            Nes Object.
        """

        for var_name in self.variables.keys():
            if self.variables[var_name]['units'] == 'mol.h-1.km-2':
                self.variables[var_name]['FieldType'] = np.int32(104)
                self.variables[var_name]['MemoryOrder'] = "XYZ"
                self.variables[var_name]['description'] = "EMISSIONS"
                self.variables[var_name]['units'] = "mol km^-2 hr^-1"
                self.variables[var_name]['stagger'] = ""
                self.variables[var_name]['coordinates'] = "XLONG XLAT"

            elif self.variables[var_name]['units'] == 'ug.s-1.m-2':
                self.variables[var_name]['FieldType'] = np.int32(104)
                self.variables[var_name]['MemoryOrder'] = "XYZ"
                self.variables[var_name]['description'] = "EMISSIONS"
                self.variables[var_name]['units'] = "ug/m3 m/s"
                self.variables[var_name]['stagger'] = ""
                self.variables[var_name]['coordinates'] = "XLONG XLAT"

            else:
                raise TypeError("The unit '{0}' of specie {1} is not defined correctly. ".format(
                    self.variables[var_name]['units'], var_name) + "Should be 'mol.h-1.km-2' or 'ug.s-1.m-2'")

            if 'long_name' in self.variables[var_name].keys():
                del self.variables[var_name]['long_name']

        return None

    def to_wrf_chem_units(self):
        """
        Change the data values according to the WRF-CHEM conventions.

        Parameters
        ----------
        self : nes.Nes
            Nes Object.

        Returns
        -------
        dict
            Variable in the MONARCH units.
        """

        self.calculate_grid_area(overwrite=False)
        for var_name in self.variables.keys():
            if isinstance(self.variables[var_name]['data'], np.ndarray):
                if self.variables[var_name]['units'] == 'mol.h-1.km-2':
                    # 10**6 -> from m2 to km2
                    # 10**3 -> from kmol to mol
                    # 3600 -> from s to h
                    self.variables[var_name]['data'] = np.array(
                        self.variables[var_name]['data'] * 10 ** 6 * 10 ** 3 * 3600, dtype=np.float32)
                elif self.variables[var_name]['units'] == 'ug.s-1.m-2':
                    # 10**9 -> from kg to ug
                    self.variables[var_name]['data'] = np.array(
                        self.variables[var_name]['data'] * 10 ** 9, dtype=np.float32)

                else:
                    raise TypeError("The unit '{0}' of specie {1} is not defined correctly. ".format(
                        self.variables[var_name]['units'], var_name) + "Should be 'mol.h-1.km-2' or 'ug.s-1.m-2'")
            self.variables[var_name]['dtype'] = np.float32

        return self.variables

    def create_times_var(self):
        """
        Create the content of the WRF-CHEM variable times.

        Parameters
        ----------
        self : nes.Nes
            Nes Object.

        Returns
        -------
        numpy.ndarray
            Array with the content of TFLAG.
        """

        aux_times = np.chararray((len(self.time), 19), itemsize=1)

        for i_d, aux_date in enumerate(self.time):
            aux_times[i_d] = list(aux_date.strftime("%Y-%m-%d_%H:%M:%S"))

        return aux_times

    def set_global_attributes(self):
        """
        Set the NetCDF global attributes

        Parameters
        ----------
        self : nes.Nes
            Nes Object.
        """

        now = datetime.now()
        if len(self.time) > 1:
            tstep = ((self.time[1] - self.time[0]).seconds // 3600) * 10000
        else:
            tstep = 1 * 10000

        current_attributes = deepcopy(self.global_attrs)
        del self.global_attrs

        self.global_attrs = {'TITLE': None,
                             'START_DATE': self.time[0].strftime("%Y-%m-%d_%H:%M:%S"),
                             'WEST-EAST_GRID_DIMENSION': None,  # Projection dependent attributes
                             'SOUTH-NORTH_GRID_DIMENSION': None,  # Projection dependent attributes
                             'BOTTOM-TOP_GRID_DIMENSION': np.int32(45),
                             'DX': None,  # Projection dependent attributes
                             'DY': None,  # Projection dependent attributes
                             'GRIDTYPE': 'C',
                             'DIFF_OPT': np.int32(1),
                             'KM_OPT': np.int32(4),
                             'DAMP_OPT': np.int32(3),
                             'DAMPCOEF': np.float32(0.2),
                             'KHDIF': np.float32(0.),
                             'KVDIF': np.float32(0.),
                             'MP_PHYSICS': np.int32(6),
                             'RA_LW_PHYSICS': np.int32(4),
                             'RA_SW_PHYSICS': np.int32(4),
                             'SF_SFCLAY_PHYSICS': np.int32(2),
                             'SF_SURFACE_PHYSICS': np.int32(2),
                             'BL_PBL_PHYSICS': np.int32(8),
                             'CU_PHYSICS': np.int32(0),
                             'SF_LAKE_PHYSICS': np.int32(0),
                             'SURFACE_INPUT_SOURCE': None,  # Projection dependent attributes
                             'SST_UPDATE': np.int32(0),
                             'GRID_FDDA': np.int32(0),
                             'GFDDA_INTERVAL_M': np.int32(0),
                             'GFDDA_END_H': np.int32(0),
                             'GRID_SFDDA': np.int32(0),
                             'SGFDDA_INTERVAL_M': np.int32(0),
                             'SGFDDA_END_H': np.int32(0),
                             'WEST-EAST_PATCH_START_UNSTAG': None,  # Projection dependent attributes
                             'WEST-EAST_PATCH_END_UNSTAG': None,  # Projection dependent attributes
                             'WEST-EAST_PATCH_START_STAG': None,  # Projection dependent attributes
                             'WEST-EAST_PATCH_END_STAG': None,  # Projection dependent attributes
                             'SOUTH-NORTH_PATCH_START_UNSTAG': None,  # Projection dependent attributes
                             'SOUTH-NORTH_PATCH_END_UNSTAG': None,  # Projection dependent attributes
                             'SOUTH-NORTH_PATCH_START_STAG': None,  # Projection dependent attributes
                             'SOUTH-NORTH_PATCH_END_STAG': None,  # Projection dependent attributes
                             'BOTTOM-TOP_PATCH_START_UNSTAG': None,
                             'BOTTOM-TOP_PATCH_END_UNSTAG': None,
                             'BOTTOM-TOP_PATCH_START_STAG': None,
                             'BOTTOM-TOP_PATCH_END_STAG': None,
                             'GRID_ID': np.int32(1),
                             'PARENT_ID': np.int32(0),
                             'I_PARENT_START': np.int32(1),
                             'J_PARENT_START': np.int32(1),
                             'PARENT_GRID_RATIO': np.int32(1),
                             'DT': np.float32(18.),
                             'CEN_LAT': None,  # Projection dependent attributes
                             'CEN_LON': None,  # Projection dependent attributes
                             'TRUELAT1': None,  # Projection dependent attributes
                             'TRUELAT2': None,  # Projection dependent attributes
                             'MOAD_CEN_LAT': None,  # Projection dependent attributes
                             'STAND_LON': None,  # Projection dependent attributes
                             'POLE_LAT': None,  # Projection dependent attributes
                             'POLE_LON': None,  # Projection dependent attributes
                             'GMT': np.float32(self.time[0].hour),
                             'JULYR': np.int32(self.time[0].year),
                             'JULDAY': np.int32(self.time[0].strftime("%j")),
                             'MAP_PROJ': None,  # Projection dependent attributes
                             'MMINLU': 'MODIFIED_IGBP_MODIS_NOAH',
                             'NUM_LAND_CAT': np.int32(41),
                             'ISWATER': np.int32(17),
                             'ISLAKE': np.int32(-1),
                             'ISICE': np.int32(15),
                             'ISURBAN': np.int32(13),
                             'ISOILWATER': np.int32(14),
                             'HISTORY': "",  # Editable
                             }

        # Editable attributes
        float_atts = ['DAMPCOEF', 'KHDIF', 'KVDIF', 'CEN_LAT', 'CEN_LON', 'DT']
        int_atts = ['BOTTOM-TOP_GRID_DIMENSION', 'DIFF_OPT', 'KM_OPT', 'DAMP_OPT',
                    'MP_PHYSICS', 'RA_LW_PHYSICS', 'RA_SW_PHYSICS', 'SF_SFCLAY_PHYSICS', 'SF_SURFACE_PHYSICS',
                    'BL_PBL_PHYSICS', 'CU_PHYSICS', 'SF_LAKE_PHYSICS', 'SURFACE_INPUT_SOURCE', 'SST_UPDATE',
                    'GRID_FDDA', 'GFDDA_INTERVAL_M', 'GFDDA_END_H', 'GRID_SFDDA', 'SGFDDA_INTERVAL_M', 'SGFDDA_END_H',
                    'BOTTOM-TOP_PATCH_START_UNSTAG', 'BOTTOM-TOP_PATCH_END_UNSTAG', 'BOTTOM-TOP_PATCH_START_STAG',
                    'BOTTOM-TOP_PATCH_END_STAG', 'GRID_ID', 'PARENT_ID', 'I_PARENT_START', 'J_PARENT_START',
                    'PARENT_GRID_RATIO', 'NUM_LAND_CAT', 'ISWATER', 'ISLAKE', 'ISICE', 'ISURBAN', 'ISOILWATER']
        str_atts = ['GRIDTYPE', 'MMINLU', 'HISTORY']
        for att_name, att_value in current_attributes.items():
            if att_name in int_atts:
                self.global_attrs[att_name] = np.int32(att_value)
            elif att_name in float_atts:
                self.global_attrs[att_name] = np.float32(att_value)
            elif att_name in str_atts:
                self.global_attrs[att_name] = str(att_value)

        # Projection dependent attributes
        if isinstance(self, nes.LCCNes) or isinstance(self, nes.MercatorNes):
            self.global_attrs['WEST-EAST_GRID_DIMENSION'] = np.int32(len(self._x['data']) + 1)
            self.global_attrs['SOUTH-NORTH_GRID_DIMENSION'] = np.int32(len(self._y['data']) + 1)
            self.global_attrs['DX'] = np.float32(self._x['data'][1] - self._x['data'][0])
            self.global_attrs['DY'] = np.float32(self._y['data'][1] - self._y['data'][0])
            self.global_attrs['SURFACE_INPUT_SOURCE'] = np.int32(1)
            self.global_attrs['WEST-EAST_PATCH_START_UNSTAG'] = np.int32(1)
            self.global_attrs['WEST-EAST_PATCH_END_UNSTAG'] = np.int32(len(self._x['data']))
            self.global_attrs['WEST-EAST_PATCH_START_STAG'] = np.int32(1)
            self.global_attrs['WEST-EAST_PATCH_END_STAG'] = np.int32(len(self._x['data']) + 1)
            self.global_attrs['SOUTH-NORTH_PATCH_START_UNSTAG'] = np.int32(1)
            self.global_attrs['SOUTH-NORTH_PATCH_END_UNSTAG'] = np.int32(len(self._y['data']))
            self.global_attrs['SOUTH-NORTH_PATCH_START_STAG'] = np.int32(1)
            self.global_attrs['SOUTH-NORTH_PATCH_END_STAG'] = np.int32(len(self._y['data']) + 1)

            self.global_attrs['POLE_LAT'] = np.float32(90)
            self.global_attrs['POLE_LON'] = np.float32(0)

            if isinstance(self, nes.LCCNes):
                self.global_attrs['MAP_PROJ'] = np.int32(1)
                self.global_attrs['TRUELAT1'] = np.float32(self.projection_data['standard_parallel'][0])
                self.global_attrs['TRUELAT2'] = np.float32(self.projection_data['standard_parallel'][1])
                self.global_attrs['MOAD_CEN_LAT'] = np.float32(self.projection_data['latitude_of_projection_origin'])
                self.global_attrs['STAND_LON'] = np.float32(self.projection_data['longitude_of_central_meridian'])
                self.global_attrs['CEN_LAT'] = np.float32(self.projection_data['latitude_of_projection_origin'])
                self.global_attrs['CEN_LON'] = np.float32(self.projection_data['longitude_of_central_meridian'])
            elif isinstance(self, nes.MercatorNes):
                self.global_attrs['MAP_PROJ'] = np.int32(3)
                self.global_attrs['TRUELAT1'] = np.float32(self.projection_data['standard_parallel'])
                self.global_attrs['TRUELAT2'] = np.float32(0)
                self.global_attrs['MOAD_CEN_LAT'] = np.float32(self.projection_data['standard_parallel'])
                self.global_attrs['STAND_LON'] = np.float32(self.projection_data['longitude_of_projection_origin'])
                self.global_attrs['CEN_LAT'] = np.float32(self.projection_data['standard_parallel'])
                self.global_attrs['CEN_LON'] = np.float32(self.projection_data['longitude_of_projection_origin'])

        return None

    def create_dimensions(self, netcdf):
        """
        Create 'time', 'time_bnds', 'lev', 'lon' and 'lat' dimensions.

        Parameters
        ----------
        self : nes.Nes
            Nes Object.
        netcdf : Dataset
            netcdf4-python open dataset.
        """

        netcdf.createDimension('Time', len(self._time))
        netcdf.createDimension('DateStrLen', 19)
        netcdf.createDimension('emissions_zdim', len(self._lev['data']))
        if isinstance(self, nes.LCCNes):
            netcdf.createDimension('west_east', len(self._x['data']))
            netcdf.createDimension('south_north', len(self._y['data']))

        return None

    def create_dimension_variables(self, netcdf):
        """
        Create the 'y' and 'x' variables.

        Parameters
        ----------
        self : nes.Nes
            Nes Object.
        netcdf : Dataset
            NetCDF object.
        """

        times = netcdf.createVariable('Times', 'S1', ('Time', 'DateStrLen', ))
        times[:] = create_times_var(self)

        return None

    def create_variables(self, netcdf):
        """
        Create the netCDF file variables.

        Parameters
        ----------
        self : nes.Nes
            Nes Object.
        netcdf : Dataset
            netcdf4-python open dataset.
        """

        for var_name, var_info in self.variables.items():
            var = netcdf.createVariable(var_name, 'f', ('Time', 'emissions_zdim', 'south_north', 'west_east',),
                                        zlib=self.zip_lvl > 0, complevel=self.zip_lvl)
            var.FieldType = var_info['FieldType']
            var.MemoryOrder = var_info['MemoryOrder']
            var.description = var_info['description']
            var.units = var_info['units']
            var.stagger = var_info['stagger']
            var.coordinates = var_info['coordinates']

            if var_info['data'] is not None:
                if self.info:
                    print("Rank {0:03d}: Filling {1})".format(self.rank, var_name))

                if isinstance(var_info['data'], int) and var_info['data'] == 0:
                    var[self.write_axis_limits['t_min']:self.write_axis_limits['t_max'],
                        self.write_axis_limits['z_min']:self.write_axis_limits['z_max'],
                        self.write_axis_limits['y_min']:self.write_axis_limits['y_max'],
                        self.write_axis_limits['x_min']:self.write_axis_limits['x_max']] = 0

                elif len(var_info['data'].shape) == 4:
                    var[self.write_axis_limits['t_min']:self.write_axis_limits['t_max'],
                        self.write_axis_limits['z_min']:self.write_axis_limits['z_max'],
                        self.write_axis_limits['y_min']:self.write_axis_limits['y_max'],
                        self.write_axis_limits['x_min']:self.write_axis_limits['x_max']] = var_info['data']

        return None