#!/usr/bin/env python

import sys
import gc
import os
import warnings
import numpy as np
import pandas as pd
import datetime
from xarray import open_dataset
from netCDF4 import Dataset, num2date, date2num
from mpi4py import MPI
from cfunits import Units
from numpy.ma.core import MaskError
import geopandas as gpd
from shapely.geometry import Polygon
from copy import deepcopy, copy
import datetime
from ..interpolation import vertical_interpolation
from ..interpolation import horizontal_interpolation
from ..nes_formats import to_netcdf_cams_ra
from ..nc_projections import *

class Nes(object):
    """

    Attributes
    ----------
    comm : MPI.Communicator.
    rank : int
        MPI rank.
    master : bool
        True when rank == 0.
    size : int
        Size of the communicator.
    print_info : bool
        Indicates if you want to print reading/writing info.
    is_xarray : bool
        (Not working) Indicates if you want to use xarray as default.
    __ini_path : str
        Path to the original file to read when open_netcdf is called.
    hours_start : int
        Number of hours to avoid from the first original values.
    hours_end : int
        Number of hours to avoid from the last original values.
    dataset : xr.Dataset
        (not working) xArray Dataset.
    netcdf : Dataset
        netcdf4-python Dataset.
    variables : dict
        Variables information.
        The variables are stored in a dictionary with the var_name as key and another dictionary with the information.
        The information dictionary contains the 'data' key with None (if the variable is not loaded) or the array values
        and the other keys are the variable attributes or description.
    _time : list
        Complete list of original time step values.
    _lev : dict
        Vertical level dictionary with the complete 'data' key for all the values and the rest of the attributes.
    _lat : dict
        Latitudes dictionary with the complete 'data' key for all the values and the rest of the attributes.
    _lon _ dict
        Longitudes dictionary with the complete 'data' key for all the values and the rest of the attributes.
    parallel_method : str
        Parallel method to read/write.
        Can be chosen any of the following axis to parallelize: 'T', 'Y' or 'X'.
    read_axis_limits : dict
        Dictionary with the 4D limits of the rank data to read.
        t_min, t_max, z_min, z_max, y_min, y_max, x_min and x_max.
    write_axis_limits : dict
        Dictionary with the 4D limits of the rank data to write.
        t_min, t_max, z_min, z_max, y_min, y_max, x_min and x_max.
    time : list
        List of time steps of the rank data.
    lev : dict
        Vertical levels dictionary with the portion of 'data' corresponding to the rank values.
    lat : dict
        Latitudes dictionary with the portion of 'data' corresponding to the rank values.
    lon : dict
        Longitudes dictionary with the portion of 'data' corresponding to the rank values.
    global_attrs : dict
        Global attributes with the attribute name as key and data as values.
    _var_dim : None, tuple
        Tuple with the name of the Y and X dimensions for the variables.
    _lat_dim : None, tuple
        Tuple with the name of the dimensions of the Latitude values.
    _lon_dim : None, tuple
        Tuple with the name of the dimensions of the Longitude values.
    """
    def __init__(self, comm=None, path=None, info=False, dataset=None, xarray=False, parallel_method='Y',
                 avoid_first_hours=0, avoid_last_hours=0, first_level=0, last_level=None, create_nes=False,
                 balanced=False, times=None, **kwargs):
        """
        Initialize the Nes class

        Parameters
        ----------
        comm: MPI.COMM
            MPI Communicator.
        path: str
            Path to the NetCDF to initialize the object.
        info: bool
            Indicates if you want to get reading/writing info.
        dataset: Dataset
            NetCDF4-python Dataset to initialize the class.
        xarray: bool:
            (Not working) Indicates if you want to use xarray as default.
        parallel_method : str
            Indicates the parallelization method that you want. Default over Y axis
            accepted values: ['X', 'Y', 'T'].
        balanced : bool
            Indicates if you want a balanced parallelization or not. 
            Balanced dataset cannot be written in chunking mode.
        avoid_first_hours : int
            Number of hours to remove from first time steps.
        avoid_last_hours : int
            Number of hours to remove from last time steps.
        first_level : int
            Index of the first level to use.
        last_level : int, None
            Index of the last level to use. None if it is the last.
        create_nes : bool
            Indicates if you want to create the object from scratch (True) or through an existing file.
        times : list, None
            List of times to substitute the current ones while creation.
        kwargs :
            Projection dependent parameters to create it from scratch
        """

        # MPI Initialization
        if comm is None:
            self.comm = MPI.COMM_WORLD
        else:
            self.comm = comm
        self.rank = self.comm.Get_rank()
        self.master = self.rank == 0
        self.size = self.comm.Get_size()

        # General info
        self.info = info
        self.is_xarray = xarray
        self.__ini_path = path
        self.shapefile = None

        # Selecting info
        self.hours_start = avoid_first_hours
        self.hours_end = avoid_last_hours
        self.first_level = first_level
        self.last_level = last_level
        self.lat_min = None
        self.lat_max = None
        self.lon_min = None
        self.lon_max = None
        self.balanced = balanced

        # Define parallel method
        self.parallel_method = parallel_method

        # NetCDF object
        if create_nes:

            self.netcdf = None
            self.dataset = None

            # Initialize variables
            self.variables = {}

            # Complete dimensions
            self._time = times
            self._time_bnds = self.__get_time_bnds(create_nes)
            self._lat_bnds, self._lon_bnds = self.__get_coordinates_bnds(create_nes)
            self._lev = {'data': np.array([0]),
                         'units': '',
                         'positive': 'up'}
            self._lat, self._lon = self._create_centre_coordinates(**kwargs)

            # Set axis limits for parallel reading
            self.read_axis_limits = self.get_read_axis_limits()

            # Dimensions screening
            self.time = self._time[self.read_axis_limits['t_min']:self.read_axis_limits['t_max']]
            self.time_bnds = self._time_bnds
            self.lev = deepcopy(self._lev)
            self.lat_bnds, self.lon_bnds = self._lat_bnds, self._lon_bnds

            # Set NetCDF attributes
            self.global_attrs = self.__get_global_attributes(create_nes)

        else:

            if dataset is not None:
                if self.is_xarray:
                    self.dataset = dataset
                    self.netcdf = None
                else:
                    self.dataset = None
                    self.netcdf = dataset
            elif self.__ini_path is not None:
                if self.is_xarray:
                    self.dataset = self.__open_dataset()
                    self.netcdf = None
                else:
                    self.dataset = None
                    self.netcdf = self.__open_netcdf4()

            # Lazy variables
            self.variables = self._get_lazy_variables()

            # Complete dimensions
            self._time = self.__get_time()
            self._time_bnds = self.__get_time_bnds()
            self._lev = self._get_coordinate_dimension(['lev', 'level', 'lm', 'plev'])
            self._lat = self._get_coordinate_dimension(['lat', 'latitude'])
            self._lon = self._get_coordinate_dimension(['lon', 'longitude'])
            self._lat_bnds, self._lon_bnds = self.__get_coordinates_bnds()

            # Set axis limits for parallel reading
            self.read_axis_limits = self.get_read_axis_limits()

            # Dimensions screening
            self.time = self._time[self.read_axis_limits['t_min']:self.read_axis_limits['t_max']]
            self.time_bnds = self._time_bnds
            self.lev = self._get_coordinate_values(self._lev, 'Z')
            self.lat = self._get_coordinate_values(self._lat, 'Y')
            self.lon = self._get_coordinate_values(self._lon, 'X')
            self.lat_bnds = self._get_coordinate_values(self._lat_bnds, 'Y', bounds=True)
            self.lon_bnds = self._get_coordinate_values(self._lon_bnds, 'X', bounds=True)

            # Set axis limits for parallel writing
            self.write_axis_limits = self.get_write_axis_limits()

            # Set NetCDF attributes
            self.global_attrs = self.__get_global_attributes()

        # Writing options
        self.zip_lvl = 0

        # Dimensions information
        self._var_dim = None
        self._lat_dim = None
        self._lon_dim = None

        self.vertical_var_name = None

    @staticmethod
    def new(comm=None, path=None, info=False, dataset=None, xarray=False, create_nes=False, balanced=False,
            parallel_method='Y', avoid_first_hours=0, avoid_last_hours=0, first_level=0, last_level=None):
        """
        Initialize the Nes class.

        Parameters
        ----------
        comm: MPI.COMM
            MPI Communicator.
        path: str
            Path to the NetCDF to initialize the object.
        info: bool
            Indicates if you want to get reading/writing info.
        dataset: Dataset
            NetCDF4-python Dataset to initialize the class.
        xarray: bool:
            (Not working) Indicates if you want to use xarray as default.
        avoid_first_hours : int
            Number of hours to remove from first time steps.
        avoid_last_hours : int
            Number of hours to remove from last time steps.
        parallel_method : str
            Indicates the parallelization method that you want. Default over Y axis
            accepted values: ['X', 'Y', 'T'].
        balanced : bool
            Indicates if you want a balanced parallelization or not. 
            Balanced dataset cannot be written in chunking mode.
        first_level : int
            Index of the first level to use.
        last_level : int, None
            Index of the last level to use. None if it is the last.
        create_nes : bool
            Indicates if you want to create the object from scratch (True) or through an existing file.
        """

        new = Nes(comm=comm, path=path, info=info, dataset=dataset, xarray=xarray, parallel_method=parallel_method,
                  avoid_first_hours=avoid_first_hours, avoid_last_hours=avoid_last_hours, first_level=first_level,
                  last_level=last_level, create=create_nes, balanced=balanced)

        return new

    def __del__(self):
        """
        To delete the Nes object and close all the open datasets.
        """

        self.close()
        self.free_vars(list(self.variables.keys()))
        del self.variables
        try:
            del self.time
            del self._time
            del self.time_bnds
            del self._time_bnds
            del self.lev
            del self._lev
            del self.lat
            del self._lat
            del self.lon
            del self._lon
            del self._lat_bnds
            del self.lat_bnds
            del self._lon_bnds
            del self.lon_bnds
        except AttributeError:
            pass

        del self
        gc.collect()

        return None

    def __getstate__(self):
        """
        Read the CSV file that contains all the Reduce variable specifications.

        Returns
        -------
        state : dict
            Dictionary with the class parameters.
        """

        d = self.__dict__
        state = {k: d[k] for k in d if k not in ['comm', 'variables', 'netcdf']}

        return state

    def __setstate__(self, state):
        """
        Set the state of the class.

        Parameters
        ----------
        state: dict
            Dictionary with the class parameters.
        """

        self.__dict__ = state

        return None

    def copy(self, copy_vars=False, copy_comm=False):
        """
        Copy the Nes object.
        The copy will avoid to copy the communicator, dataset and variables by default.

        Parameters
        ----------
        copy_vars: bool
            Indicates if you want to copy the variables (in lazy mode).
        copy_comm: bool
            Indicates if you want to copy the communicator.

        Returns
        -------
        nessy : Nes
            Copy of the Nes object.
        """

        nessy = deepcopy(self)

        if copy_comm:
            nessy.comm = self.comm

        if copy_vars:
            if not hasattr(nessy, 'netcdf'):
                nessy.netcdf = None
                nessy.variables = self._get_lazy_variables()
            else:
                nessy.variables = nessy._get_lazy_variables()
        else:
            nessy.variables = {}

        return nessy

    def get_full_times(self):
        return self._time

    def get_full_levels(self):
        return self._lev

    def set_level_direction(self, new_direction):
        if new_direction not in ['up', 'down']:
            raise ValueError("Level direction mus be up or down. '{0}' is not a valid option".format(new_direction))
        self._lev['positive'] = new_direction
        self.lev['positive'] = new_direction

        return True

    def reverse_level_direction(self):
        if 'positive' in self._lev.keys():
            if self._lev['positive'] == 'up':
                self._lev['positive'] = 'down'
                self.lev['positive'] = 'down'
            else:
                self._lev['positive'] = 'up'
                self.lev['positive'] = 'up'
        return True

    def clear_communicator(self):
        """
        Erase the communicator and the parallelization indexes.
        """

        self.comm = None
        self.rank = 0
        self.master = 0
        self.size = 0

        return None

    def set_communicator(self, comm):
        """
        Set a new communicator and the correspondent parallelization indexes.

        Parameters
        ----------
        comm: MPI.COMM
            Communicator to be set.
        """

        self.comm = comm
        self.rank = self.comm.Get_rank()
        self.master = self.rank == 0
        self.size = self.comm.Get_size()

        self.read_axis_limits = self.get_read_axis_limits()
        self.write_axis_limits = self.get_write_axis_limits()
        return None

    def set_levels(self, levels):
        """
        Modify the original level values with new ones.

        Parameters
        ----------
        levels : dict
            Dictionary with the new level information to be set.
        """

        self._lev = deepcopy(levels)
        self.lev = deepcopy(levels)

        return None

    def set_time_bnds(self, time_bnds):
        """
        Modify the original time bounds values with new ones.

        Parameters
        ----------
        time_bnds : list
            List with the new time bounds information to be set.
        """

        correct_format = True
        for time_bnd in np.array(time_bnds).flatten():
            if not isinstance(time_bnd, datetime.datetime):
                print("{0} is not a datetime object".format(time_bnd))
                correct_format = False
        if correct_format:
            if len(self._time) == len(time_bnds):
                self._time_bnds = deepcopy(time_bnds)
                self.time_bnds = deepcopy(time_bnds)
            else:
                msg = "WARNING!!! "
                msg += "The given time bounds list has a different length than the time array. "
                msg += "(time:{0}, bnds:{1}). Time bounds will not be set.".format(len(self._time), len(time_bnds))
                warnings.warn(msg)
        else:
            msg = 'WARNING!!! '
            msg += 'There is at least one element in the time bounds to be set that is not a datetime object. '
            msg += 'Time bounds will not be set.'
            warnings.warn(msg)

        return None

    def create_single_spatial_bounds(self, coordinates, inc, spatial_nv=2, inverse=False):
        """
        Calculate the vertices coordinates.

        Parameters
        ----------
        coordinates : np.array
            Coordinates in degrees (latitude or longitude).
        inc : float 
            Increment between centre values.
        spatial_nv : int 
            Non mandatory parameter that informs the number of vertices that must have the
            boundaries. Default: 2.
        inverse : bool 
            For some grid latitudes.

        Returns
        ----------
        bounds : np.array
            Array with as many elements as vertices for each value of coords.
        """
        
        # Create new arrays moving the centres half increment less and more.
        coords_left = coordinates - inc / 2
        coords_right = coordinates + inc / 2

        # Defining the number of corners needed. 2 to regular grids and 4 for irregular ones.
        if spatial_nv == 2:
            # Create an array of N arrays of 2 elements to store the floor and the ceil values for each cell
            bounds = np.dstack((coords_left, coords_right))
            bounds = bounds.reshape((len(coordinates), spatial_nv))
        elif spatial_nv == 4:
            # Create an array of N arrays of 4 elements to store the corner values for each cell
            # It can be stored in clockwise starting form the left-top element, or in inverse mode.
            if inverse:
                bounds = np.dstack((coords_left, coords_left, coords_right, coords_right))
            else:
                bounds = np.dstack((coords_left, coords_right, coords_right, coords_left))
        else:
            raise ValueError('The number of vertices of the boundaries must be 2 or 4.')

        return bounds

    def create_spatial_bounds(self):
        """
        Calculate longitude and latitude bounds and set them.
        """

        inc_lat = np.abs(np.mean(np.diff(self._lat['data'])))
        lat_bnds = self.create_single_spatial_bounds(self._lat['data'], inc_lat, spatial_nv=2)

        self._lat_bnds = deepcopy(lat_bnds)
        self.lat_bnds = lat_bnds[self.write_axis_limits['y_min']:self.write_axis_limits['y_max'],
                                 :]

        inc_lon = np.abs(np.mean(np.diff(self._lon['data'])))
        lon_bnds = self.create_single_spatial_bounds(self._lon['data'], inc_lon, spatial_nv=2)

        self._lon_bnds = deepcopy(lon_bnds)
        self.lon_bnds = lon_bnds[self.write_axis_limits['x_min']:self.write_axis_limits['x_max'], 
                                 :]

        return None

    def free_vars(self, var_list):
        """
        Erase the selected variables from the variables information.

        Parameters
        ----------
        var_list : list, str
            List (or single string) of the variables to be loaded.
        """

        if isinstance(var_list, str):
            var_list = [var_list]

        if self.is_xarray:
            self.dataset = self.dataset.drop_vars(var_list)
            self.variables = self._get_lazy_variables()
        else:
            if self.variables is not None:
                for var_name in var_list:
                    if var_name in self.variables:
                        if 'data' in self.variables[var_name].keys():
                            del self.variables[var_name]['data']
                        del self.variables[var_name]
        gc.collect()

        return None

    def keep_vars(self, var_list):
        """
        Keep the selected variables and erases the rest.

        Parameters
        ----------
        var_list : list, str
            List (or single string) of the variables to be loaded.
        """

        if isinstance(var_list, str):
            var_list = [var_list]

        to_remove = list(set(self.variables.keys()).difference(set(var_list)))

        self.free_vars(to_remove)

        return None

    def get_time_interval(self):
        """
        Calculate the interrval of hours between time steps.

        Returns
        -------
        int
            Number of hours between time steps.
        """

        time_interval = self._time[1] - self._time[0]
        time_interval = int(time_interval.seconds // 3600)

        return time_interval

    def sel_time(self, time, copy=False):
        """
        To select only one time step.

        Parameters
        ----------
        time : datetime.datetime
            Time stamp to select.
        copy : bool
            Indicates if you want a copy with the selected time step (True) or to modify te existing one (False).

        Returns
        -------
        Nes
            Nes object with the data (and metadata) of the selected time step.
        """

        if copy:
            aux_nessy = self.copy(copy_vars=False)
            aux_nessy.comm = self.comm
        else:
            aux_nessy = self

        aux_nessy.hours_start = 0
        aux_nessy.hours_end = 0

        idx_time = aux_nessy.time.index(time)

        aux_nessy.time = [self.time[idx_time]]
        aux_nessy._time = aux_nessy.time
        for var_name, var_info in self.variables.items():
            if copy:
                aux_nessy.variables[var_name] = {}
                for att_name, att_value in var_info.items():
                    if att_name == 'data':
                        if att_value is None:
                            raise ValueError("{} data not loaded".format(var_name))
                        aux_nessy.variables[var_name][att_name] = att_value[[idx_time]]
                    else:
                        aux_nessy.variables[var_name][att_name] = att_value
            else:
                aux_nessy.variables[var_name]['data'] = aux_nessy.variables[var_name]['data'][[idx_time]]
        
        return aux_nessy

    def sel(self, hours_start=None, time_min=None, hours_end=None, time_max=None, lev_min=None, lev_max=None,
            lat_min=None, lat_max=None, lon_min=None, lon_max=None):
            
        loaded_vars = False
        for var_info in self.variables.values():
            if var_info['data'] is not None:
                loaded_vars = True
                # var_info['data'] = None
        if loaded_vars:
            raise ValueError("Some variables have been loaded. Use select function before load.")

        # First time filter
        if hours_start is not None:
            if time_min is not None:
                raise ValueError("Choose to select by hours_start or time_min but not both")
            self.hours_start = hours_start
        elif time_min is not None:
            if time_min <= self._time[0]:
                self.hours_start = 0
            else:
                self.hours_start = int((time_min - self._time[0]).total_seconds() // 3600)

        # Last time Filter
        if hours_end is not None:
            if time_max is not None:
                raise ValueError("Choose to select by hours_end or time_max but not both")
            self.hours_end = hours_end
        elif time_max is not None:
            if time_max >= self._time[-1]:
                self.hours_end = 0
            else:
                self.hours_end = int((self._time[-1] - time_max).total_seconds() // 3600)

        # Level filter
        self.first_level = lev_min
        self.last_level = lev_max

        # Coordinate filter
        self.lat_min = lat_min
        self.lat_max = lat_max
        self.lon_min = lon_min
        self.lon_max = lon_max

        # New Axis limits
        self.read_axis_limits = self.get_read_axis_limits()

        # Dimensions screening
        self.time = self._time[self.read_axis_limits['t_min']:self.read_axis_limits['t_max']]
        self.time_bnds = self._time_bnds
        self.lev = self._get_coordinate_values(self._lev, 'Z')
        self.lat = self._get_coordinate_values(self._lat, 'Y')
        self.lon = self._get_coordinate_values(self._lon, 'X')

        self.lat_bnds = self._get_coordinate_values(self._lat_bnds, 'Y', bounds=True)
        self.lon_bnds = self._get_coordinate_values(self._lon_bnds, 'X', bounds=True)

        # Removing complete coordinates
        self.write_axis_limits = self.get_write_axis_limits()

        self.filter_coordinates_selection()

        return None

    def filter_coordinates_selection(self):
        idx = self.get_idx_intervals()

        self._time = self._time[idx['idx_t_min']:idx['idx_t_max']]

        if len(self._lat['data'].shape) == 1:
            # Regular projection
            self._lat['data'] = self._lat['data'][idx['idx_y_min']:idx['idx_y_max']]
            self._lon['data'] = self._lon['data'][idx['idx_x_min']:idx['idx_x_max']]

            if self._lat_bnds is not None:
                self._lat_bnds = self._lat_bnds[idx['idx_y_min']:idx['idx_y_max'], :]
            if self._lon_bnds is not None:
                self._lon_bnds = self._lon_bnds[idx['idx_x_min']:idx['idx_x_max'], :]
        else:
            # Irregular projections
            self._lat['data'] = self._lat['data'][idx['idx_y_min']:idx['idx_y_max'], idx['idx_x_min']:idx['idx_x_max']]
            self._lon['data'] = self._lon['data'][idx['idx_y_min']:idx['idx_y_max'], idx['idx_x_min']:idx['idx_x_max']]

            if self._lat_bnds is not None:
                self._lat_bnds = self._lat_bnds[idx['idx_y_min']:idx['idx_y_max'], idx['idx_x_min']:idx['idx_x_max'], :]
            if self._lon_bnds is not None:
                self._lon_bnds = self._lon_bnds[idx['idx_y_min']:idx['idx_y_max'], idx['idx_x_min']:idx['idx_x_max'], :]

        return None

    def get_idx_intervals(self):
        idx = {'idx_t_min': self.get_time_id(self.hours_start, first=True),
               'idx_t_max': self.get_time_id(self.hours_end, first=False)}

        # Axis Y
        if self.lat_min is None:
            idx['idx_y_min'] = 0
        else:
            idx['idx_y_min'] = self.get_coordinate_id(self._lat['data'], self.lat_min, axis=0)
        if self.lat_max is None:
            idx['idx_y_max'] = self._lat['data'].shape[0] + 1
        else:
            idx['idx_y_max'] = self.get_coordinate_id(self._lat['data'], self.lat_max, axis=0) + 1

        if idx['idx_y_min'] > idx['idx_y_max']:
            idx_aux = copy(idx['idx_y_min'])
            idx['idx_y_min'] = idx['idx_y_max']
            idx['idx_y_max'] = idx_aux

        # Axis X

        if self.lon_min is None:
            idx['idx_x_min'] = 0
        else:
            if len(self._lon['data'].shape) == 1:
                axis = 0
            else:
                axis = 1
            idx['idx_x_min'] = self.get_coordinate_id(self._lon['data'], self.lon_min, axis=axis)
        if self.lon_max is None:
            idx['idx_x_max'] = self._lon['data'].shape[-1] + 1
        else:
            if len(self._lon['data'].shape) == 1:
                axis = 0
            else:
                axis = 1
            idx['idx_x_max'] = self.get_coordinate_id(self._lon['data'], self.lon_max, axis=axis) + 1

        if idx['idx_x_min'] > idx['idx_x_max']:
            idx_aux = copy(idx['idx_x_min'])
            idx['idx_x_min'] = idx['idx_x_max']
            idx['idx_x_max'] = idx_aux
        return idx

    # ==================================================================================================================
    #                                                 Statistics
    # ==================================================================================================================

    def last_time_step(self):
        """
        Modify variables to keep only the last time step.
        """

        if self.parallel_method == 'T':
            raise NotImplementedError("Statistics are not implemented on time axis paralelitation method.")
        aux_time = self._time[0].replace(hour=0, minute=0, second=0, microsecond=0)
        self._time = [aux_time]
        self.time = [aux_time]

        for var_name, var_info in self.variables.items():
            if var_info['data'] is None:
                self.load(var_name)
            aux_data = var_info['data'][-1, :]
            if len(aux_data.shape) == 3:
                aux_data = aux_data.reshape((1, aux_data.shape[0], aux_data.shape[1], aux_data.shape[2]))
            self.variables[var_name]['data'] = aux_data
        self.hours_start = 0
        self.hours_end = 0

        return None

    def daily_statistic(self, op, type_op='calendar'):
        """
        Calculate daily statistic.

        Parameters
        ----------
        op : str
            Statistic to perform. Accepted values: "max", "mean" and "min".
        type_op : str
            Type of statistic to perform. Accepted values: "calendar", "alltsteps", and "withoutt0".
            - "calendar": Calculate the statistic using the time metadata.
                It will avoid single time step by day calculations
            - "alltsteps": Calculate a single time statistic with all the time steps.
            - "withoutt0": Calculate a single time statistic with all the time steps avoiding the first one.
        """

        if self.parallel_method == 'T':
            raise NotImplementedError("Statistics are not implemented on time axis paralelitation method.")
        time_interval = self.get_time_interval()
        if type_op == 'calendar':
            aux_time_bounds = []
            aux_time = []
            day_list = [date_aux.day for date_aux in self.time]
            for var_name, var_info in self.variables.items():
                if var_info['data'] is None:
                    self.load(var_name)
                stat_data = None
                for day in np.unique(day_list):
                    idx_first = next(i for i, val in enumerate(day_list, 0) if val == day)
                    idx_last = len(day_list) - next(i for i, val in enumerate(reversed(day_list), 1) if val == day)
                    if idx_first != idx_last:  # To avoid single time step statistic
                        if idx_last != len(day_list):
                            if op == 'mean':
                                data_aux = var_info['data'][idx_first:idx_last + 1, :, :, :].mean(axis=0)
                            elif op == 'max':
                                data_aux = var_info['data'][idx_first:idx_last + 1, :, :, :].max(axis=0)
                            elif op == 'min':
                                data_aux = var_info['data'][idx_first:idx_last + 1, :, :, :].min(axis=0)
                            else:
                                raise NotImplementedError("Statistic operation '{0}' is not implemented.".format(op))
                            aux_time_bounds.append([self.time[idx_first], self.time[idx_last]])
                        else:
                            if op == 'mean':
                                data_aux = var_info['data'][idx_first:, :, :, :].mean(axis=0)
                            elif op == 'max':
                                data_aux = var_info['data'][idx_first:, :, :, :].max(axis=0)
                            elif op == 'min':
                                data_aux = var_info['data'][idx_first:, :, :, :].min(axis=0)
                            else:
                                raise NotImplementedError("Statistic operation '{0}' is not implemented.".format(op))
                            aux_time_bounds.append([self.time[idx_first], self.time[-1]])

                        data_aux = data_aux.reshape((1, data_aux.shape[0], data_aux.shape[1], data_aux.shape[2]))
                        aux_time.append(self.time[idx_first].replace(hour=0, minute=0, second=0))
                        # Append over time dimension
                        if stat_data is None:
                            stat_data = data_aux.copy()
                        else:
                            stat_data = np.vstack([stat_data, data_aux])
                self.variables[var_name]['data'] = stat_data
                self.variables[var_name]['cell_methods'] = "time: {0} (interval: {1}hr)".format(op, time_interval)
            self.time = aux_time
            self._time = self.time

            self.set_time_bnds(aux_time_bounds)

        elif type_op == 'alltsteps':
            for var_name, var_info in self.variables.items():
                if var_info['data'] is None:
                    self.load(var_name)
                if op == 'mean':
                    aux_data = var_info['data'].mean(axis=0)
                elif op == 'max':
                    aux_data = var_info['data'].max(axis=0)
                elif op == 'min':
                    aux_data = var_info['data'].min(axis=0)
                else:
                    raise NotImplementedError("Statistic operation '{0}' is not implemented.".format(op))
                if len(aux_data.shape) == 3:
                    aux_data = aux_data.reshape((1, aux_data.shape[0], aux_data.shape[1], aux_data.shape[2]))
                self.variables[var_name]['data'] = aux_data
                self.variables[var_name]['cell_methods'] = "time: {0} (interval: {1}hr)".format(op, time_interval)

            aux_time = self.time[0].replace(hour=0, minute=0, second=0, microsecond=0)
            aux_time_bounds = [[self.time[0], self.time[-1]]]
            self.time = [aux_time]
            self._time = self.time

            self.set_time_bnds(aux_time_bounds)

        elif type_op == 'withoutt0':
            for var_name, var_info in self.variables.items ():
                if var_info['data'] is None:
                    self.load(var_name)
                if op == 'mean':
                    aux_data = var_info['data'][1:, :].mean(axis=0)
                elif op == 'max':
                    aux_data = var_info['data'][1:, :].max(axis=0)
                elif op == 'min':
                    aux_data = var_info['data'][1:, :].min(axis=0)
                else:
                    raise NotImplementedError("Statistic operation '{0}' is not implemented.".format(op))
                if len(aux_data.shape) == 3:
                    aux_data = aux_data.reshape((1, aux_data.shape[0], aux_data.shape[1], aux_data.shape[2]))
                self.variables[var_name]['data'] = aux_data
                self.variables[var_name]['cell_methods'] = "time: {0} (interval: {1}hr)".format(op, time_interval)
            aux_time = self._time[1].replace(hour=0, minute=0, second=0, microsecond=0)
            aux_time_bounds = [[self._time[1], self._time[-1]]]
            self.time = [aux_time]
            self._time = self.time

            self.set_time_bnds(aux_time_bounds)
        else:
            raise NotImplementedError("Statistic operation type '{0}' is not implemented.".format(type_op))
        self.hours_start = 0
        self.hours_end = 0

        return None

    @staticmethod
    def _get_axis_index_(axis):
        if axis == 'T':
            value = 0
        elif axis == 'Z':
            value = 1
        elif axis == 'Y':
            value = 2
        elif axis == 'X':
            value = 3
        else:
            raise ValueError("Unknown axis: {0}".format(axis))
        return value

    def sum_axis(self, axis='Z'):
        if self.parallel_method == axis:
            raise NotImplementedError("It is not possible to sum the axis with it is parallelized '{0}'".format(
                self.parallel_method))

        for var_name, var_info in self.variables.items():
            if var_info['data'] is not None:
                self.variables[var_name]['data'] = self.variables[var_name]['data'].sum(
                    axis=self._get_axis_index_(axis), keepdims=True)
                if axis == 'T':
                    self.variables[var_name]['cell_methods'] = "time: sum (interval: {0}hr)".format(
                        (self.time[-1] - self.time[0]).total_seconds() // 3600)

        if axis == 'T':
            self.set_time_bnds([self.time[0], self.time[-1]])
            self.time = [self.time[0]]
            self._time = [self._time[0]]
        if axis == 'Z':
            self.lev['data'] = [self.lev['data'][0]]
            self._lev['data'] = [self._lev['data'][0]]
        return None

    # ==================================================================================================================
    #                                                 Reading
    # ==================================================================================================================
    
    def get_read_axis_limits(self):
        """
        Calculate the 4D reading axis limits depending on if them have to balanced or not.

        Returns
        -------
        dict
            Dictionary with the 4D limits of the rank data to read.
            t_min, t_max, z_min, z_max, y_min, y_max, x_min and x_max.
        """

        if self.balanced:
            return self.get_read_axis_limits_balanced()
        else:
            return self.get_read_axis_limits_unbalanced()

    def get_read_axis_limits_unbalanced(self):
        """
        Calculate the 4D reading axis limits.

        Returns
        -------
        dict
            Dictionary with the 4D limits of the rank data to read.
            t_min, t_max, z_min, z_max, y_min, y_max, x_min and x_max.
        """

        axis_limits = {'x_min': None, 'x_max': None,
                       'y_min': None, 'y_max': None,
                       'z_min': None, 'z_max': None,
                       't_min': None, 't_max': None}
        
        idx = self.get_idx_intervals()
        
        if self.parallel_method == 'Y':
            y_len = idx['idx_y_max'] - idx['idx_y_min']
            if y_len < self.size:
                raise IndexError('More processors (size={0}) selected than Y elements (size={1})'.format(
                    self.size, y_len))
            axis_limits['y_min'] = ((y_len // self.size) * self.rank) + idx['idx_y_min']
            if self.rank + 1 < self.size:
                axis_limits['y_max'] = ((y_len // self.size) * (self.rank + 1)) + idx['idx_y_max']
            else:
                axis_limits['y_max'] = idx['idx_y_max']

            # Non parallel filters
            axis_limits['x_min'] = idx['idx_x_min']
            axis_limits['x_max'] = idx['idx_x_max']

            axis_limits['t_min'] = idx['idx_t_min']
            axis_limits['t_max'] = idx['idx_t_max']
            
        elif self.parallel_method == 'X':
            x_len = idx['idx_x_max'] - idx['idx_x_min']
            if x_len < self.size:
                raise IndexError('More processors (size={0}) selected than X elements (size={1})'.format(
                    self.size, x_len))
            axis_limits['x_min'] = ((x_len // self.size) * self.rank) + idx['idx_x_min']
            if self.rank + 1 < self.size:
                axis_limits['x_max'] = ((x_len // self.size) * (self.rank + 1)) + idx['idx_x_max']
            else:
                axis_limits['x_max'] = idx['idx_x_max']

            # Non parallel filters
            axis_limits['y_min'] = idx['idx_y_min']
            axis_limits['y_max'] = idx['idx_y_max']

            axis_limits['t_min'] = idx['idx_t_min']
            axis_limits['t_max'] = idx['idx_t_max']
            
        elif self.parallel_method == 'T':
            t_len = idx['idx_t_min'] - idx['idx_t_max']
            if t_len < self.size:
                raise IndexError('More processors (size={0}) selected than T elements (size={1})'.format(
                    self.size, t_len))
            axis_limits['t_min'] = ((t_len // self.size) * self.rank) + idx['idx_t_min']
            if self.rank + 1 < self.size:
                axis_limits['t_max'] = ((t_len // self.size) * (self.rank + 1)) + idx['idx_t_max']

            # Non parallel filters
            axis_limits['y_min'] = idx['idx_y_min']
            axis_limits['y_max'] = idx['idx_y_max']

            axis_limits['x_min'] = idx['idx_x_min']
            axis_limits['x_max'] = idx['idx_x_max']

        else:
            raise NotImplementedError("Parallel method '{meth}' is not implemented. Use one of these: {accept}".format(
                meth=self.parallel_method, accept=['X', 'Y', 'T']))

        # Vertical levels selection:
        axis_limits['z_min'] = self.first_level
        if self.last_level == -1 or self.last_level is None:
            self.last_level = None
        elif self.last_level + 1 == len(self._lev['data']):
            self.last_level = None
        else:
            self.last_level += 1
        axis_limits['z_max'] = self.last_level

        return axis_limits

    def get_read_axis_limits_balanced(self):
        """
        Calculate the 4D reading balanced axis limits.

        Returns
        -------
        dict
            Dictionary with the 4D limits of the rank data to read.
            t_min, t_max, z_min, z_max, y_min, y_max, x_min and x_max.
        """
        idx = self.get_idx_intervals()

        fid_dist = {}
        if self.parallel_method == 'Y':
            len_to_split = idx['idx_y_max'] - idx['idx_y_min']
            if len_to_split < self.size:
                raise IndexError('More processors (size={0}) selected than Y elements (size={1})'.format(
                    self.size, len_to_split))
            min_axis = 'y_min'
            max_axis = 'y_max'
            to_add = idx['idx_y_min']
        elif self.parallel_method == 'X':
            len_to_split = idx['idx_x_max'] - idx['idx_x_min']
            if len_to_split < self.size:
                raise IndexError('More processors (size={0}) selected than X elements (size={1})'.format(
                    self.size, len_to_split))
            min_axis = 'x_min'
            max_axis = 'x_max'
            to_add = idx['idx_x_min']
        elif self.parallel_method == 'T':
            len_to_split = idx['idx_t_max'] - idx['idx_t_min']
            if len_to_split < self.size:
                raise IndexError('More processors (size={0}) selected than T elements (size={1})'.format(
                    self.size, len_to_split))
            min_axis = 't_min'
            max_axis = 't_max'
            to_add = idx['idx_t_min']
        else:
            raise NotImplementedError("Parallel method '{meth}' is not implemented. Use one of these: {accept}".format(
                meth=self.parallel_method, accept=['X', 'Y', 'T']))

        procs_len = len_to_split // self.size
        procs_rows_extended = len_to_split - (procs_len * self.size)

        rows_sum = 0
        for proc in range(self.size):
            fid_dist[proc] = {'x_min': None, 'x_max': None,
                              'y_min': None, 'y_max': None,
                              'z_min': None, 'z_max': None,
                              't_min': None, 't_max': None}
            if proc < procs_rows_extended:
                aux_rows = procs_len + 1
            else:
                aux_rows = procs_len

            len_to_split -= aux_rows
            if len_to_split < 0:
                rows = len_to_split + aux_rows
            else:
                rows = aux_rows

            fid_dist[proc][min_axis] = rows_sum
            fid_dist[proc][max_axis] = rows_sum + rows

            if to_add is not None:
                fid_dist[proc][min_axis] += to_add
                fid_dist[proc][max_axis] += to_add

            # # Last element
            # if len_to_split == 0 and to_add == 0:
            #     fid_dist[proc][max_axis] = None

            rows_sum += rows

        axis_limits = fid_dist[self.rank]

        # Non parallel filters
        if self.parallel_method != 'T':
            axis_limits['t_min'] = idx['idx_t_min']
            axis_limits['t_max'] = idx['idx_t_max']
        if self.parallel_method != 'X':
            axis_limits['x_min'] = idx['idx_x_min']
            axis_limits['x_max'] = idx['idx_x_max']
        if self.parallel_method != 'Y':
            axis_limits['y_min'] = idx['idx_y_min']
            axis_limits['y_max'] = idx['idx_y_max']

        # Vertical levels selection:
        axis_limits['z_min'] = self.first_level
        if self.last_level == -1 or self.last_level is None:
            self.last_level = None
        elif self.last_level + 1 == len(self._lev['data']):
            self.last_level = None
        else:
            self.last_level += 1
        axis_limits['z_max'] = self.last_level

        return axis_limits

    def get_time_id(self, hours, first=True):
        """
        Get the index of the corresponding time value.

        Parameters
        ----------
        hours : int
            Number of hours to avoid.
        first : bool
            Indicates if you want to avoid from the first hours (True) or from the last (False).
            Default: True.

        Returns
        -------
        int
            Index of the time array.
        """

        from datetime import timedelta

        if first:
            idx = self._time.index(self._time[0] + timedelta(hours=hours))
        else:
            idx = self._time.index(self._time[-1] - timedelta(hours=hours)) + 1

        return idx

    def get_coordinate_id(self, array, value, axis=0):
        """
        Get the index of the corresponding coordinate value.

        Parameters
        ----------
        value : float
            Coordinate value to search.
        min : bool
            Indicates if you want to apply the filter as minimum (True) or maximum (False) values
            Default: True.

        Returns
        -------
        int
            Index of the coordinate array.
        """
        idx = (np.abs(array - value)).argmin(axis=axis).min()

        return idx

    def open(self):
        """
        Open the NetCDF.
        """

        if self.is_xarray:
            self.dataset = self.__open_dataset()
            self.netcdf = None
        else:
            self.dataset = None
            self.netcdf = self.__open_netcdf4()

        return None

    def __open_dataset(self):
        """
        Open the NetCDF with xarray.

        Returns
        -------
        dataset : xr.Dataset
            Open dataset.
        """

        if self.master:
            warnings.filterwarnings('ignore')  # Disabling warnings while reading MONARCH original file
            dataset = open_dataset(self.__ini_path, decode_coords='all')
            warnings.filterwarnings('default')  # Re-activating warnings
        else:
            dataset = None

        dataset = self.comm.bcast(dataset, root=0)
        self.dataset = dataset

        return dataset

    def __open_netcdf4(self, mode='r'):
        """
        Open the NetCDF with netcdf4-python.

        Parameters
        ----------
        mode : str
            Inheritance from mode parameter from https://unidata.github.io/netcdf4-python/#Dataset.__init__
            Default: 'r' (read-only).
        Returns
        -------
        netcdf : Dataset
            Open dataset.
        """

        if self.size == 1:
            netcdf = Dataset(self.__ini_path, format="NETCDF4", mode=mode, parallel=False)
        else:
            netcdf = Dataset(self.__ini_path, format="NETCDF4", mode=mode, parallel=True, comm=self.comm,
                             info=MPI.Info())
        self.netcdf = netcdf

        return netcdf

    def close(self):
        """
        Close the NetCDF with netcdf4-python.
        """

        if hasattr(self, 'netcdf'):
            if self.netcdf is not None:
                self.netcdf.close()
                self.netcdf = None

        return None

    def __get_dates_from_months(self, time, units, calendar):
        """
        Calculates the number of days since the first date
        in the 'time' list and store in new list:
        This is useful when the units are 'months since',
        which cannot be transformed to dates using num2date.

        Parameter
        ---------
        time: list
            Original time.
        units: str
            CF compliant time units.
        calendar: str
            Original calendar.

        Returns
        -------
        time: list
            CF compliant time.
        """

        start_date_str = time.units.split('since')[1].lstrip()
        start_date = datetime.datetime(int(start_date_str[0:4]),
                                       int(start_date_str[5:7]),
                                       int(start_date_str[8:10]))
        new_time = []

        for current_date in time:

            # Transform current_date into number of days since base date
            current_date = num2date(current_date, self.__parse_time_unit(units), calendar=calendar)

            # Calculate number of days between base date and the other dates
            n_days = (current_date - start_date).days

            # Store in list
            new_time.append(n_days)

        return new_time

    def __parse_time(self, time):
        """
        Parses the time to be CF compliant.

        Parameters
        ----------
        time: str
            Original time.

        Returns
        -------
        time : str
            CF compliant time.
        """

        units = time.units
        if not hasattr(time, 'calendar'):
            calendar = 'standard'
        else:
            calendar = time.calendar

        if 'months since' in time.units:
            units = 'days since ' + time.units.split('since')[1].lstrip()
            time = self.__get_dates_from_months(time, units, calendar)

        return time, units, calendar

    @staticmethod
    def __parse_time_unit(t_units):
        """
        Parses the time units to be CF compliant.

        Parameters
        ----------
        t_units : str
            Original time units.

        Returns
        -------
        t_units : str
            CF compliant time units.
        """

        if 'h @' in t_units:
            t_units = 'hour since {0}-{1}-{2} {3}:{4}:{5} UTC'.format(
                t_units[4:8], t_units[8:10], t_units[10:12], t_units[13:15], t_units[15:17], t_units[17:-4])

        return t_units

    def __get_time(self):
        """
        Get the NetCDF file time values.

        Returns
        -------
        time : list
            List of times (datetime.datetime) of the NetCDF data.
        """

        if self.is_xarray:
            time = self.variables['time']
        else:
            if self.master:
                nc_var = self.netcdf.variables['time']
                nc_var, units, calendar = self.__parse_time(nc_var)
                time = num2date(nc_var[:], self.__parse_time_unit(units), calendar=calendar)
                time = [aux.replace(second=0, microsecond=0) for aux in time]
            else:
                time = None
            time = self.comm.bcast(time, root=0)
        self.free_vars('time')

        return time

    def __get_time_bnds(self, create_nes=False):
        """
        Get the NetCDF time bounds values.

        Parameters
        ----------
        create_nes : bool
            Indicates if you want to create the object from scratch (True) or through an existing file.

        Returns
        -------
        time_bnds : list
            List of time bounds (datetime) of the NetCDF data.
        """

        if self.is_xarray:
            time_bnds = self.variables['time_bnds']
        else:
            if self.master:
                if not create_nes:
                    if 'time_bnds' in self.netcdf.variables.keys():
                        time = self.netcdf.variables['time']
                        nc_var = self.netcdf.variables['time_bnds']
                        time_bnds = num2date(nc_var[:], self.__parse_time_unit(time.units), 
                                             calendar=time.calendar).tolist()
                    else:
                        time_bnds = None
                else:
                    time_bnds = None
            else:
                time_bnds = None
            time_bnds = self.comm.bcast(time_bnds, root=0)

        self.free_vars('time_bnds')

        return time_bnds

    def __get_coordinates_bnds(self, create_nes=False):
        """
        Get the NetCDF coordinates bounds values.

        Parameters
        ----------
        create_nes : bool
            Indicates if you want to create the object from scratch (True) or through an existing file.

        Returns
        -------
        lat_bnds : list
            List of latitude bounds of the NetCDF data.
        lon_bnds : list
            List of longitude bounds of the NetCDF data.
        """

        if self.is_xarray:
            lat_bnds = self.variables['lat_bnds']
            lon_bnds = self.variables['lon_bnds']
        else:
            if self.master:
                if not create_nes:
                    if 'lat_bnds' in self.netcdf.variables.keys():
                        lat_bnds = self.netcdf.variables['lat_bnds'][:]
                    else:
                        lat_bnds = None
                    if 'lon_bnds' in self.netcdf.variables.keys():
                        lon_bnds = self.netcdf.variables['lon_bnds'][:]
                    else:
                        lon_bnds = None
                else:
                    lat_bnds = None
                    lon_bnds = None
            else:
                lat_bnds = None
                lon_bnds = None
            lat_bnds = self.comm.bcast(lat_bnds, root=0)
            lon_bnds = self.comm.bcast(lon_bnds, root=0)

        self.free_vars(['lat_bnds', 'lon_bnds'])

        return lat_bnds, lon_bnds

    def _get_coordinate_dimension(self, possible_names):
        """
        Read the coordinate dimension data.

        This will read the complete data of the coordinate.

        Parameters
        ----------
        possible_names: list, str
            List (or single string) of the possible names of the coordinate (e.g. ['lat', 'latitude']).

        Returns
        -------
        nc_var : dict
            Dictionary with the 'data' key with the coordinate variable values. and the attributes as other keys.
        """

        if isinstance(possible_names, str):
            possible_names = [possible_names]

        try:
            dimension_name = set(possible_names).intersection(set(self.variables.keys())).pop()
            if self.is_xarray:
                nc_var = self.dataset[dimension_name]
            else:
                nc_var = self.variables[dimension_name].copy()
                nc_var['data'] = self.netcdf.variables[dimension_name][:]
                if hasattr(nc_var, 'units'):
                    if nc_var['units'] in ['unitless', '-']:
                        nc_var['units'] = ''
            self.free_vars(dimension_name)
        except KeyError:
            nc_var = {'data': np.array([0]),
                      'units': ''
                     }

        return nc_var

    def _get_coordinate_values(self, coordinate_info, coordinate_axis, bounds=False):
        """
        Get the coordinate data of the current portion.

        Parameters
        ----------
        coordinate_info : dict, list
            Dictionary with the 'data' key with the coordinate variable values. and the attributes as other keys.
        coordinate_axis : str
            Name of the coordinate to extract. Accepted values: ['Z', 'Y', 'X'].
        bounds : bool
            Boolean variable to know if there are coordinate bounds.
        Returns
        -------
        values : dict
            Dictionary with the portion of data corresponding to the rank.
        """

        if coordinate_info is None:
            return None

        if not isinstance(coordinate_info, dict):
            values = {'data': deepcopy(coordinate_info)}
        else:
            values = deepcopy(coordinate_info)

        coordinate_len = len(values['data'].shape)
        if bounds:
            coordinate_len -= 1

        if coordinate_axis == 'Y':
            if coordinate_len == 1:
                values['data'] = values['data'][self.read_axis_limits['y_min']:self.read_axis_limits['y_max']]
            elif coordinate_len == 2:
                values['data'] = values['data'][self.read_axis_limits['y_min']:self.read_axis_limits['y_max'], self.read_axis_limits['x_min']:self.read_axis_limits['x_max']]
            else:
                raise NotImplementedError("The coordinate has wrong dimensions: {dim}".format(
                    dim=values['data'].shape))
        elif coordinate_axis == 'X':
            if coordinate_len == 1:
                values['data'] = values['data'][self.read_axis_limits['x_min']:self.read_axis_limits['x_max']]
            elif coordinate_len == 2:
                values['data'] = values['data'][self.read_axis_limits['y_min']:self.read_axis_limits['y_max'], self.read_axis_limits['x_min']:self.read_axis_limits['x_max']]
            else:
                raise NotImplementedError("The coordinate has wrong dimensions: {dim}".format(
                    dim=values['data'].shape))
        elif coordinate_axis == 'Z':
            if coordinate_len == 1:
                values['data'] = values['data'][self.read_axis_limits['z_min']:self.read_axis_limits['z_max']]
            else:
                raise NotImplementedError("The coordinate has wrong dimensions: {dim}".format(
                    dim=values['data'].shape))

        return values

    def _get_lazy_variables(self):
        """
        Get all the variables information.

        Returns
        -------
        variables : dict
            Dictionary with the variable name as key and another dictionary as value.
            De value dictionary will have the 'data' key with None as value and all the variable attributes as the
            other keys.
            e.g.
            {'var_name_1': {'data': None, 'attr_1': value_1_1, 'attr_2': value_1_2, ...},
             'var_name_2': {'data': None, 'attr_1': value_2_1, 'attr_2': value_2_2, ...},
             ...}
        """

        if self.is_xarray:
            variables = self.dataset.variables
        else:
            if self.master:
                variables = {}
                if self.netcdf is not None:
                    # Initialise data
                    for var_name, var_info in self.netcdf.variables.items():
                        variables[var_name] = {}
                        variables[var_name]['data'] = None
                        variables[var_name]['dimensions'] = var_info.dimensions

                        # Avoid some attributes
                        for attrname in var_info.ncattrs():
                            if attrname not in ['missing_value', '_FillValue']:
                                value = getattr(var_info, attrname)
                                if value in ['unitless', '-']:
                                    value = ''
                                variables[var_name][attrname] = value
                else:
                    # Unload loaded data (used in to_shapefile with data from interpolated or loaded datasets)
                    for var_name, var_info in self.variables.items():
                        variables[var_name] = {}
                        variables[var_name]['data'] = None
                        variables[var_name]['dimensions'] = var_info['dimensions']
            else:
                variables = None
            variables = self.comm.bcast(variables, root=0)

        return variables

    def _read_variable(self, var_name, nc_var=None, var_dims=None):
        """
        Read the corresponding variable data according to the current rank.

        Parameters
        ----------
        var_name : str
            Name of the variable to read.

        Returns
        -------
        data: np.array
            Portion of the variable data corresponding to the rank.
        """

        if nc_var is None and var_dims is None:
            nc_var = self.netcdf.variables[var_name]
            var_dims = nc_var.dimensions

        # Read data in 4 dimensions
        if len(var_dims) < 2:
            data = nc_var[:]
        elif len(var_dims) == 2:
            data = nc_var[self.read_axis_limits['y_min']:self.read_axis_limits['y_max'],
                          self.read_axis_limits['x_min']:self.read_axis_limits['x_max']]
            data = data.reshape(1, 1, data.shape[-2], data.shape[-1])
        elif len(var_dims) == 3:
            data = nc_var[self.read_axis_limits['t_min']:self.read_axis_limits['t_max'],
                          self.read_axis_limits['y_min']:self.read_axis_limits['y_max'],
                          self.read_axis_limits['x_min']:self.read_axis_limits['x_max']]
            data = data.reshape(data.shape[-3], 1, data.shape[-2], data.shape[-1])
        elif len(var_dims) == 4:
            data = nc_var[self.read_axis_limits['t_min']:self.read_axis_limits['t_max'],
                          self.read_axis_limits['z_min']:self.read_axis_limits['z_max'],
                          self.read_axis_limits['y_min']:self.read_axis_limits['y_max'],
                          self.read_axis_limits['x_min']:self.read_axis_limits['x_max']]
        # elif len(var_dims) == 5:
        #     data = nc_var[self.read_axis_limits['t_min']:self.read_axis_limits['t_max'],
        #                   :,
        #                   self.read_axis_limits['z_min']:self.read_axis_limits['z_max'],
        #                   self.read_axis_limits['y_min']:self.read_axis_limits['y_max'],
        #                   self.read_axis_limits['x_min']:self.read_axis_limits['x_max']]
        else:
            raise NotImplementedError('Error with {0}. Only can be read netCDF with 4 dimensions or less'.format(
                var_name))
        # Missing to nan
        try:
            data[data.mask == True] = np.nan
        except (AttributeError, MaskError, ValueError):
            pass

        return data

    def load(self, var_list=None):
        """
        Load of the selected variables.

        That function will fill the variable 'data' key with the corresponding values.

        Parameters
        ----------
        var_list : list, str
            List (or single string) of the variables to be loaded.
        """
        
        if (self.__ini_path is None) and (self.dataset is None) and (self.netcdf is None):
            raise RuntimeError('Only data from existent files can be loaded.')

        if self.netcdf is None:
            self.__open_dataset()
            close = True
        else:
            close = False

        if isinstance(var_list, str):
            var_list = [var_list]
        elif var_list is None:
            var_list = list(self.variables.keys())

        for i, var_name in enumerate(var_list):
            if self.info:
                print("Rank {0:03d}: Loading {1} var ({2}/{3})".format(self.rank, var_name, i + 1, len(var_list)))
            if self.variables[var_name]['data'] is None:
                self.variables[var_name]['data'] = self._read_variable(var_name)
            else:
                if self.master:
                    print("Data for {0} was previously loaded. Skipping variable.".format(var_name))
            if self.info:
                print("Rank {0:03d}: Loaded {1} var ({2})".format(
                    self.rank, var_name, self.variables[var_name]['data'].shape))

        if close:
            self.close()

        return None

    def to_dtype(self, data_type='float32'):
        
        for var_name, var_info in self.variables.items():
            if var_info['data'] is not None:
                self.variables[var_name]['data'] = self.variables[var_name]['data'].astype(data_type)
        
        return None

    def concatenate(self, aux_nessy):
        """
        Concatenate different variables into the same nes object.

        Parameters
        ----------
        aux_nessy : Nes, str
            Nes object or str with the path to the NetCDF file that contains the variables to add.

        Returns
        -------
        list
            List of var names added.
        """

        if isinstance(aux_nessy, str):
            aux_nessy = self.new(path=aux_nessy, comm=self.comm, parallel_method=self.parallel_method,
                                 xarray=self.is_xarray,
                                 avoid_first_hours=self.hours_start, avoid_last_hours=self.hours_end,
                                 first_level=self.first_level, last_level=self.last_level)
            new = True
        else:
            new = False
        for var_name, var_info in aux_nessy.variables.items():
            if var_info['data'] is None:
                aux_nessy.load(var_name)

        new_vars_added = []
        for new_var_name, new_var_data in aux_nessy.variables.items():
            if new_var_name not in self.variables.keys():
                self.variables[new_var_name] = deepcopy(new_var_data)
            new_vars_added.append(new_var_name)

        if new:
            del aux_nessy

        return new_vars_added

    def __get_global_attributes(self, create_nes=False):
        """
        Read the netcdf global attributes.

        Parameters
        ----------
        create_nes : bool
            Indicates if you want to create the object from scratch (True) or through an existing file.

        Returns
        -------
        gl_attrs : dict
            Dictionary with the netCDF global attributes.
        """

        gl_attrs = {}
        if self.is_xarray:
            gl_attrs = self.dataset.attrs
        else:
            if not create_nes:
                for attrname in self.netcdf.ncattrs():
                    gl_attrs[attrname] = getattr(self.netcdf, attrname)

        return gl_attrs

    # ==================================================================================================================
    #                                                 Writing
    # ==================================================================================================================
    
    def get_write_axis_limits(self):
        """
        Calculate the 4D writing axis limits depending on if them have to balanced or not.

        Returns
        -------
        dict
            Dictionary with the 4D limits of the rank data to write.
            t_min, t_max, z_min, z_max, y_min, y_max, x_min and x_max.
        """

        if self.balanced:
            return self.get_write_axis_limits_balanced()
        else:
            return self.get_write_axis_limits_unbalanced()

    def get_write_axis_limits_unbalanced(self):
        """
        Calculate the 4D writing axis limits.

        Returns
        -------
        dict
            Dictionary with the 4D limits of the rank data to write.
            t_min, t_max, z_min, z_max, y_min, y_max, x_min and x_max.
        """

        axis_limits = {'x_min': None, 'x_max': None,
                       'y_min': None, 'y_max': None,
                       'z_min': None, 'z_max': None,
                       't_min': None, 't_max': None}

        if self.parallel_method == 'Y':
            y_len = self._lat['data'].shape[0]
            axis_limits['y_min'] = (y_len // self.size) * self.rank
            if self.rank + 1 < self.size:
                axis_limits['y_max'] = (y_len // self.size) * (self.rank + 1)
        elif self.parallel_method == 'X':
            x_len = self._lon['data'].shape[-1]
            axis_limits['x_min'] = (x_len // self.size) * self.rank
            if self.rank + 1 < self.size:
                axis_limits['x_max'] = (x_len // self.size) * (self.rank + 1)
        elif self.parallel_method == 'T':
            t_len = len(self._time)
            axis_limits['t_min'] = ((t_len // self.size) * self.rank)
            if self.rank + 1 < self.size:
                axis_limits['t_max'] = (t_len // self.size) * (self.rank + 1)
        else:
            raise NotImplementedError("Parallel method '{meth}' is not implemented. Use one of these: {accept}".format(
                meth=self.parallel_method, accept=['X', 'Y', 'T']))

        return axis_limits

    def get_write_axis_limits_balanced(self):
        """
        Calculate the 4D reading balanced axis limits.

        Returns
        -------
        dict
            Dictionary with the 4D limits of the rank data to read.
            t_min, t_max, z_min, z_max, y_min, y_max, x_min and x_max.
        """

        fid_dist = {}
        if self.parallel_method == 'Y':
            len_to_split = self._lat['data'].shape[0]
            min_axis = 'y_min'
            max_axis = 'y_max'
        elif self.parallel_method == 'X':
            len_to_split = self._lon['data'].shape[-1]
            min_axis = 'x_min'
            max_axis = 'x_max'
        elif self.parallel_method == 'T':
            len_to_split = len(self._time)
            min_axis = 't_min'
            max_axis = 't_max'
        else:
            raise NotImplementedError("Parallel method '{meth}' is not implemented. Use one of these: {accept}".format(
                meth=self.parallel_method, accept=['X', 'Y', 'T']))

        procs_len = len_to_split // self.size
        procs_rows_extended = len_to_split - (procs_len * self.size)

        rows_sum = 0
        for proc in range(self.size):
            fid_dist[proc] = {'x_min': None, 'x_max': None,
                              'y_min': None, 'y_max': None,
                              'z_min': None, 'z_max': None,
                              't_min': None, 't_max': None}
            if proc < procs_rows_extended:
                aux_rows = procs_len + 1
            else:
                aux_rows = procs_len

            len_to_split -= aux_rows
            if len_to_split < 0:
                rows = len_to_split + aux_rows
            else:
                rows = aux_rows

            fid_dist[proc][min_axis] = rows_sum
            fid_dist[proc][max_axis] = rows_sum + rows

            # Last element
            if len_to_split == 0:
                fid_dist[proc][max_axis] = None

            rows_sum += rows

        axis_limits = fid_dist[self.rank]

        return axis_limits

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

        Parameters
        ----------
        netcdf : Dataset
            netcdf4-python open dataset.
        """

        # Create time dimension
        netcdf.createDimension('time', None)

        # Create time_nv (number of vertices) dimension
        if self._time_bnds is not None:
            netcdf.createDimension('time_nv', 2)

        # Create spatial_nv (number of vertices) dimension
        if (self._lat_bnds is not None) and (self._lon_bnds is not None):
            netcdf.createDimension('spatial_nv', 2)
        
        # Create lev, lon and lat dimensions
        netcdf.createDimension('lev', len(self.lev['data']))
        netcdf.createDimension('lon', len(self._lon['data']))
        netcdf.createDimension('lat', len(self._lat['data']))

        return None

    def _create_dimension_variables(self, netcdf):
        """
        Create the 'time', 'time_bnds', 'lev', 'lat', 'lat_bnds', 'lon' and 'lon_bnds' variables.

        Parameters
        ----------
        netcdf : Dataset
            netcdf4-python open dataset.
        """

        # TIMES
        time_var = netcdf.createVariable('time', np.float64, ('time',), zlib=self.zip_lvl > 0, complevel=self.zip_lvl)
        time_var.units = 'hours since {0}'.format(
            self._time[self.get_time_id(self.hours_start, first=True)].strftime('%Y-%m-%d %H:%M:%S'))
        time_var.standard_name = 'time'
        time_var.calendar = 'standard'
        time_var.long_name = 'time'
        if self._time_bnds is not None:
            time_var.bounds = 'time_bnds'
        if self.size > 1:
            time_var.set_collective(True)
        time_var[:] = date2num(self._time[self.get_time_id(self.hours_start, first=True):
                                          self.get_time_id(self.hours_end, first=False)],
                               time_var.units, time_var.calendar)

        # TIME BOUNDS
        if self._time_bnds is not None:
            time_bnds_var = netcdf.createVariable('time_bnds', np.float64, ('time', 'time_nv',), zlib=self.zip_lvl,
                                                  complevel=self.zip_lvl)
            if self.size > 1:
                time_bnds_var.set_collective(True)
            time_bnds_var[:] = date2num(self._time_bnds, time_var.units, calendar='standard')

        # LEVELS
        lev = netcdf.createVariable('lev', np.float64, ('lev',), zlib=self.zip_lvl > 0, complevel=self.zip_lvl)
        if 'units' in self._lev.keys():
            lev.units = Units(self._lev['units'], formatted=True).units
        else:
            lev.units = ''
        if 'positive' in self._lev.keys():
            lev.positive = self._lev['positive']

        if self.size > 1:
            lev.set_collective(True)
        lev[:] = self._lev['data']

        # LATITUDES
        lat = netcdf.createVariable('lat', np.float64, self._lat_dim, zlib=self.zip_lvl > 0, complevel=self.zip_lvl)
        lat.units = 'degrees_north'
        lat.axis = 'Y'
        lat.long_name = 'latitude coordinate'
        lat.standard_name = 'latitude'
        if self._lat_bnds is not None:
            lat.bounds = 'lat_bnds'
        if self.size > 1:
            lat.set_collective(True)
        lat[:] = self._lat['data']

        # LATITUDES BOUNDS
        if self._lat_bnds is not None:
            lat_bnds_var = netcdf.createVariable('lat_bnds', np.float64, self._lat_dim + ('spatial_nv',), 
                                                 zlib=self.zip_lvl > 0, complevel=self.zip_lvl)
            if self.size > 1:
                lat_bnds_var.set_collective(True)
            lat_bnds_var[:] = self._lat_bnds[:]

        # LONGITUDES
        lon = netcdf.createVariable('lon', np.float64, self._lon_dim, zlib=self.zip_lvl > 0, complevel=self.zip_lvl)
        lon.units = 'degrees_east'
        lon.axis = 'X'
        lon.long_name = 'longitude coordinate'
        lon.standard_name = 'longitude'
        if self._lon_bnds is not None:
            lon.bounds = 'lon_bnds'
        if self.size > 1:
            lon.set_collective(True)
        lon[:] = self._lon['data']

        # LONGITUDES BOUNDS
        if self._lon_bnds is not None:
            lon_bnds_var = netcdf.createVariable('lon_bnds', np.float64, self._lon_dim + ('spatial_nv',), 
                                                 zlib=self.zip_lvl > 0, complevel=self.zip_lvl)
            if self.size > 1:
                lon_bnds_var.set_collective(True)
            lon_bnds_var[:] = self._lon_bnds[:]

        return None

    def _create_variables(self, netcdf, chunking=False):
        """
        Create the netCDF file variables.

        Parameters
        ----------
        netcdf : Dataset
            netcdf4-python open dataset.
        chunking : bool
            Indicates if you want to chunk the output netCDF.
        """

        for i, (var_name, var_dict) in enumerate(self.variables.items()):
            if var_dict['data'] is not None:
                if self.info:
                    print("Rank {0:03d}: Writing {1} var ({2}/{3})".format(self.rank, var_name, i + 1, len(self.variables)))
                try:
                    if not chunking:
                        var = netcdf.createVariable(var_name, var_dict['data'].dtype, ('time', 'lev',) + self._var_dim,
                                                    zlib=self.zip_lvl > 0, complevel=self.zip_lvl)
                    else:
                        if self.balanced:
                            raise NotImplementedError("A balanced data cannot be chunked.")
                        if self.master:
                            chunk_size = var_dict['data'].shape
                        else:
                            chunk_size = None
                        chunk_size = self.comm.bcast(chunk_size, root=0)
                        var = netcdf.createVariable(var_name, var_dict['data'].dtype, ('time', 'lev',) + self._var_dim,
                                                    zlib=self.zip_lvl > 0, complevel=self.zip_lvl, chunksizes=chunk_size)
                    if self.info:
                        print("Rank {0:03d}: Var {1} created ({2}/{3})".format(
                            self.rank, var_name, i + 1, len(self.variables)))
                    if self.size > 1:
                        var.set_collective(True)
                        if self.info:
                            print("Rank {0:03d}: Var {1} collective ({2}/{3})".format(
                                self.rank, var_name, i + 1, len(self.variables)))

                    for att_name, att_value in var_dict.items():
                        if att_name == 'data':
                            if self.info:
                                print("Rank {0:03d}: Filling {1})".format(self.rank, var_name))
                            try:
                                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']] = att_value
                            except ValueError:
                                var[self.write_axis_limits['t_min']:self.write_axis_limits['t_max'],
                                    0,
                                    self.write_axis_limits['y_min']:self.write_axis_limits['y_max'],
                                    self.write_axis_limits['x_min']:self.write_axis_limits['x_max']] = att_value
                                # msg = "*WARNING* '{0}' variable is a 3D field. Setting it on first (0) layer.".format(
                                #     var_name)
                                # warn(msg)
                            except IndexError:
                                raise IndexError("Different shapes. out_shape={0}, data_shp={1}".format(
                                    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']].shape,
                                    att_value.shape))
                            if self.info:
                                print("Rank {0:03d}: Var {1} data ({2}/{3})".format(self.rank, var_name, i + 1,
                                                                                    len(self.variables)))
                        elif att_name not in ['chunk_size', 'var_dims', 'dimensions']:
                            var.setncattr(att_name, att_value)
                    self._set_var_crs(var)
                    if self.info:
                        print("Rank {0:03d}: Var {1} completed ({2}/{3})".format(self.rank, var_name, i + 1,
                                                                                len(self.variables)))
                except Exception as e:
                    print("**ERROR** an error has occurred while writing the '{0}' variable".format(var_name))
                    # print("**ERROR** an error has occurredred while writing the '{0}' variable".format(var_name),
                    #       file=sys.stderr)
                    raise e
            else:
                msg = 'WARNING!!! '
                msg += 'Variable {0} was not loaded. It will not be written.'.format(var_name)
                warnings.warn(msg)

        return None

    def _create_centre_coordinates(self):
        """
        Must be implemented on inner class.
        """

        return None

    def _create_metadata(self, netcdf):
        """
        Must be implemented on inner class.
        """

        return None

    def _set_crs(self, netcdf):
        """
        Must be implemented on inner class.

        Parameters
        ----------
        netcdf : Dataset
            netcdf4-python Dataset.
        """

        return None

    @staticmethod
    def _set_var_crs(var):
        """
        Must be implemented on inner class.

        Parameters
        ----------
        var : Variable
            netCDF4-python variable object.
        """

        return None

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

        Parameters
        ----------
        path : str
            Path to the output netCDF file.
        chunking: bool
            Indicates if you want to chunk the output netCDF.
        """

        # 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
        self._create_dimensions(netcdf)

        # Create dimension variables
        self._create_dimension_variables(netcdf)
        if self.info:
            print("Rank {0:03d}: Dimensions done".format(self.rank))

        # Create variables
        self._create_variables(netcdf, chunking=chunking)

        # Create metadata
        self._create_metadata(netcdf)

        # Close NetCDF
        if self.global_attrs is not None:
            for att_name, att_value in self.global_attrs.items():
                netcdf.setncattr(att_name, att_value)
        netcdf.setncattr('Conventions', 'CF-1.7')

        netcdf.close()

        return None

    def __to_netcdf_cams_ra(self, path):
        return to_netcdf_cams_ra(self, path)

    def to_netcdf(self, path, compression_level=0, serial=False, info=False,
                  chunking=False, type='NES'):
        """
        Write the netCDF output file.

        Parameters
        ----------
        path : str
            Path to the output netCDF file.
        compression_level : int
            Level of compression (0 to 9) Default: 0 (no compression).
        serial : bool
            Indicates if you want to write in serial or not. Default: False.
        info : bool
            Indicates if you want to print the information of each writing step by stdout Default: False.
        chunking : bool
            Indicates if you want a chunked netCDF output. Only available with non serial writes. Default: False.
        """

        old_info = self.info
        self.info = info

        self.zip_lvl = compression_level
        if self.is_xarray:
            raise NotImplementedError("Writing with xarray not implemented")
        else:
            # if serial:
            if serial and self.size > 1:
                data = self._gather_data()
                if self.master:
                    new_nc = self.copy(copy_vars=False)
                    new_nc.set_communicator(MPI.COMM_SELF)
                    new_nc.variables = data
                    if type == 'NES':
                        new_nc.__to_netcdf_py(path)
                    elif type == 'CAMS_RA':
                        new_nc.__to_netcdf_cams_ra(path)
                    else:
                        raise ValueError(
                            "Unknown NetCDF type '{0}'. Use 'CAMS_RA' or 'NES'; default='NES'".format(type))
            else:
                if type == 'NES':
                    self.__to_netcdf_py(path, chunking=chunking)
                elif type == 'CAMS_RA':
                    self.__to_netcdf_cams_ra(path)
                else:
                    raise ValueError("Unknown NetCDF type '{0}'. Use 'CAMS_RA' or 'NES'; default='NES'".format(type))

        self.info = old_info

        return None

    def __to_grib2(self, path, grib_keys, grib_template_path, lat_flip=True, info=False):
        """
        Private method to write output file with grib2 format.

        Parameters
        ----------
        path : str
            Path to the output file.
        grib_keys : dict
            Dictionary with the grib2 keys.
        grib_template_path : str
            Path to the grib2 file to use as template.
        info : bool
            Indicates if you want to print extra information during the process.
        """

        from eccodes import codes_grib_new_from_file
        from eccodes import codes_keys_iterator_new
        from eccodes import codes_keys_iterator_next
        from eccodes import codes_keys_iterator_get_name
        from eccodes import codes_get_string
        from eccodes import codes_keys_iterator_delete
        from eccodes import codes_clone
        from eccodes import codes_set
        from eccodes import codes_set_values
        from eccodes import codes_write
        from eccodes import codes_release

        fout = open(path, 'wb')

        # read template
        fin = open(grib_template_path, 'rb')

        gid = codes_grib_new_from_file(fin)
        if gid is None:
            sys.exit(1)

        iterid = codes_keys_iterator_new(gid, 'ls')
        while codes_keys_iterator_next(iterid):
            keyname = codes_keys_iterator_get_name(iterid)
            keyval = codes_get_string(gid, keyname)
            if info:
                print("%s = %s" % (keyname, keyval))

        codes_keys_iterator_delete(iterid)
        for var_name, var_info in self.variables.items():
            for i_time, time in enumerate(self.time):
                for i_lev, lev in enumerate(self.lev['data']):
                    clone_id = codes_clone(gid)

                    # Adding grib2 keys to file
                    for key, value in grib_keys.items():
                        if value not in ['', 'None', None]:
                            try:
                                codes_set(clone_id, key, value)
                            except Exception as e:
                                print("Something went wrong while writing the Grib key '{0}': {1}".format(key, value))
                                raise e

                    # Time dependent keys
                    if 'dataTime' in grib_keys.keys() and grib_keys['dataTime'] in ['', 'None', None]:
                        codes_set(clone_id, 'dataTime', int(i_time * 100))
                    if 'stepRange' in grib_keys.keys() and grib_keys['stepRange'] in ['', 'None', None]:
                        n_secs = (time - self._time[0]).total_seconds()
                        codes_set(clone_id, 'stepRange', int(n_secs // 3600))
                    if 'forecastTime' in grib_keys.keys() and grib_keys['forecastTime'] in ['', 'None', None]:
                        n_secs = (time - self._time[0]).total_seconds()
                        codes_set(clone_id, 'forecastTime', int(n_secs))

                    # Level dependent keys
                    if 'typeOfFirstFixedSurface' in grib_keys.keys() and \
                            grib_keys['typeOfFirstFixedSurface'] in ['', 'None', None]:
                        if float(lev) == 0:
                            codes_set(clone_id, 'typeOfFirstFixedSurface', 1)
                            # grib_keys['typeOfFirstFixedSurface'] = 1
                        else:
                            codes_set(clone_id, 'typeOfFirstFixedSurface', 103)
                            # grib_keys['typeOfFirstFixedSurface'] = 103
                    if 'level' in grib_keys.keys() and grib_keys['level'] in ['', 'None', None]:
                        codes_set(clone_id, 'level', float(lev))

                    newval = var_info['data'][i_time, i_lev, :, :]
                    if lat_flip:
                        newval = np.flipud(newval)

                    # TODO Check default NaN Value
                    newval[np.isnan(newval)] = 0.

                    codes_set_values(clone_id, np.array(newval.ravel(), dtype='float64'))
                    # codes_set_values(clone_id, newval.ravel())
                    codes_write(clone_id, fout)
                    del newval
        codes_release(gid)
        fout.close()
        fin.close()

        return None

    def to_grib2(self, path, grib_keys, grib_template_path, lat_flip=True, info=False):
        """
        Write output file with grib2 format.

        Parameters
        ----------
        path : str
            Path to the output file.
        grib_keys : dict
            Dictionary with the grib2 keys.
        grib_template_path : str
            Path to the grib2 file to use as template.
        info : bool
            Indicates if you want to print extra information during the process.
        """

        # if serial:
        if self.parallel_method in ['X', 'Y'] and self.size > 1:
            data = self._gather_data()
            if self.master:
                new_nc = self.copy(copy_vars=False)
                new_nc.set_communicator(MPI.COMM_SELF)
                new_nc.variables = data
                new_nc.__to_grib2(path, grib_keys, grib_template_path, lat_flip=lat_flip, info=info)
        else:
            self.__to_grib2(path, grib_keys, grib_template_path, lat_flip=lat_flip, info=info)

        return None

    def create_shapefile(self):
        """
        Create spatial geodataframe (shapefile).
        """

        if self._lat_bnds is None or self._lon_bnds is None:
            self.create_spatial_bounds()

        # Reshape arrays to create geometry
        aux_shape = (self.lat_bnds.shape[0], self.lon_bnds.shape[0], 4)
        lon_bnds_aux = np.empty(aux_shape)
        lon_bnds_aux[:, :, 0] = self.lon_bnds[np.newaxis, :, 0]
        lon_bnds_aux[:, :, 1] = self.lon_bnds[np.newaxis, :, 1]
        lon_bnds_aux[:, :, 2] = self.lon_bnds[np.newaxis, :, 1]
        lon_bnds_aux[:, :, 3] = self.lon_bnds[np.newaxis, :, 0]

        lon_bnds = lon_bnds_aux
        del lon_bnds_aux

        lat_bnds_aux = np.empty(aux_shape)
        lat_bnds_aux[:, :, 0] = self.lat_bnds[:, np.newaxis, 0]
        lat_bnds_aux[:, :, 1] = self.lat_bnds[:, np.newaxis, 0]
        lat_bnds_aux[:, :, 2] = self.lat_bnds[:, np.newaxis, 1]
        lat_bnds_aux[:, :, 3] = self.lat_bnds[:, np.newaxis, 1]

        lat_bnds = lat_bnds_aux
        del lat_bnds_aux

        aux_b_lats = lat_bnds.reshape((lat_bnds.shape[0] * lat_bnds.shape[1], lat_bnds.shape[2]))
        aux_b_lons = lon_bnds.reshape((lon_bnds.shape[0] * lon_bnds.shape[1], lon_bnds.shape[2]))

        # Create dataframe cointaining all polygons
        geometry = []
        for i in range(aux_b_lons.shape[0]):
            geometry.append(Polygon([(aux_b_lons[i, 0], aux_b_lats[i, 0]),
                                     (aux_b_lons[i, 1], aux_b_lats[i, 1]),
                                     (aux_b_lons[i, 2], aux_b_lats[i, 2]),
                                     (aux_b_lons[i, 3], aux_b_lats[i, 3]),
                                     (aux_b_lons[i, 0], aux_b_lats[i, 0])]))
        fids = np.arange(len(self._lat['data']) * len(self._lon['data']))
        fids = fids.reshape((len(self._lat['data']), len(self._lon['data'])))
        fids = fids[self.read_axis_limits['y_min']:self.read_axis_limits['y_max'],
                    self.read_axis_limits['x_min']:self.read_axis_limits['x_max']]
        gdf = gpd.GeoDataFrame(index=pd.Index(name='FID', data=fids.ravel()), 
                               geometry=geometry, 
                               crs="EPSG:4326")
        self.shapefile = gdf

        return gdf

    def write_shapefile(self, path):
        """ 
        Save spatial geodataframe (shapefile).

        Parameters
        ----------
        path : str
            Path to the output file.
        """
        
        if self.shapefile is None:
            raise ValueError('Shapefile was not created.')
        
        if self.size == 1:
            # In serial, avoid gather
            self.shapefile.to_file(path)
        else:
            # In parallel
            data = self.comm.gather(self.shapefile, root=0)
            if self.master:
                data = pd.concat(data)
                data.to_file(path)
        
        return None

    def to_shapefile(self, path, time=None, lev=None, var_list=None):
        """
        Create shapefile from NES data.
        
        1. Check 2D data is selected.
        2. Create grid shapefile.
        3. Add variables to shapefile (as independent function).
        4. Write shapefile.

        Parameters
        ----------
        path : str
            Path to the output file.
        time : datetime.datetime
            Time stamp to select.
        lev : int
            Vertical level to select.
        var_list : list, str
            List (or single string) of the variables to be loaded and saved in the shapefile.
        """

        # Make a copy, so we do not lose data of  original dataset
        aux_nessy = self.copy(copy_vars=True, copy_comm=True)

        if var_list is None:
            var_list = deepcopy(aux_nessy.variables).keys()
        else:
            # Transform to list if there is only one element
            if isinstance(var_list, str):
                var_list = [var_list]
            # Check if variable exists
            for var_name in var_list:
                if var_name not in aux_nessy.variables.keys():
                    raise ValueError('{} does not exist in dataset.'.format(var_name))
            # Skip variables (if desired)
            aux_nessy.keep_vars(var_list)

        # Check if time and level have to be selected
        sel_time, sel_lev = False, False
        data_list = deepcopy(aux_nessy.variables)
        for var_name in data_list.keys():
            dim = data_list[var_name]['dimensions']
            if 'time' in dim:
                sel_time = True
            elif 'lev' in dim:
                sel_lev = True
            break

        # Select first vertical level (if needed)
        if sel_lev:
            if lev is None:
                msg = 'No vertical level has been specified. The first one will be selected.'
                warnings.warn(msg)
                aux_nessy.sel(lev_min=0, lev_max=0)
            else:
                if lev not in aux_nessy.lev['data']:
                    raise ValueError('Level {} is not available. Choose from {}'.format(lev, aux_nessy.lev['data']))
                aux_nessy.sel(lev_min=lev, lev_max=lev)

        # Select first time (if needed)
        if sel_time:
            if time is None:
                msg = 'No time has been specified. The first one will be selected.'
                warnings.warn(msg)
                aux_nessy.sel(time_min=aux_nessy.time[0], time_max=aux_nessy.time[0])
            else:
                if time not in aux_nessy.time:
                    raise ValueError('Time {} is not available. Choose from {}'.format(time, aux_nessy.time))
                aux_nessy.sel(time_min=time, time_max=time)
        
        # Create shapefile
        aux_nessy.create_shapefile()

        # Load variables from original file and get data for selected time / level
        self.add_variables_to_shapefile(aux_nessy, var_list)

        # Add data to shapefile
        self.shapefile = aux_nessy.shapefile

        # Write shapefile
        aux_nessy.write_shapefile(path)

        return None

    def add_variables_to_shapefile(self, aux_nessy, var_list):
        """
        Add variables data to shapefile.

        var_list : list, str
            List (or single string) of the variables to be loaded and saved in the shapefile.
        """

        if var_list:
               
            # Load data
            # We cannot load it directly with the function load() because aux_nessy.netcdf is None
            for i, var_name in enumerate(var_list):
                if self.info:
                    print("Rank {0:03d}: Loading {1} var ({2}/{3})".format(self.rank, var_name, i + 1, 
                                                                           len(var_list)))
                
                # Add data
                if aux_nessy.variables[var_name]['data'] is None:

                    # Get data from variables in original file
                    nc_var = self.variables[var_name]['data']
                    var_dims = self.variables[var_name]['dimensions']

                    # Do not add data if there is no data loaded in original
                    if self.variables[var_name]['data'] is None:
                        msg = 'Data for {} was not loaded, it will not be saved in the shapefile. '.format(var_name)
                        msg += 'Load the data using load() if you want to include it.'
                        print(msg)
                        return
                    
                    # Add data
                    aux_nessy.variables[var_name]['data'] = aux_nessy._read_variable(var_name, nc_var, var_dims)

                else:
                    if self.master:
                        print("Data for {0} was previously loaded. Skipping variable.".format(var_name))
                
                if self.info:
                    print("Rank {0:03d}: Loaded {1} var ({2})".format(self.rank, var_name, 
                                                                      aux_nessy.variables[var_name]['data'].shape))

            # Add data
            for var_name in var_list:
                aux_nessy.shapefile[var_name] = aux_nessy.variables[var_name]['data'][:].ravel()

        return None

    def spatial_join(self, mask, method=None):
        """
        Compute overlay intersection of two GeoPandasDataFrames.

        Parameters
        ----------
        mask : GeoPandasDataFrame
            File from where the data will be obtained on the intersection.
        method : str
            Overlay method. Accepted values: ['nearest', 'intersection', None].
        """

        # Nearest centroids to the mask polygons
        if method == 'nearest':
            
            # Make copy
            shapefile_aux = deepcopy(self.shapefile)

            # Get centroids of shapefile to mask
            shapefile_aux.geometry = shapefile_aux.centroid

            # Calculate spatial joint by distance
            shapefile_aux = gpd.sjoin_nearest(shapefile_aux, mask.to_crs(shapefile_aux.crs), distance_col='distance')
            
            # Get data from closest shapes to centroids
            del shapefile_aux['geometry'], shapefile_aux['index_right']
            self.shapefile.loc[shapefile_aux.index, shapefile_aux.columns] = shapefile_aux

        # Intersect the areas of the mask polygons, outside of the mask there will be NaN
        elif method == 'intersection':
        
            # Get intersected areas
            inp, res = mask.sindex.query_bulk(self.shapefile.geometry, predicate='intersects')
            print('Rank {0:03d}: {1} intersected areas found'.format(self.rank, len(inp)))
            
            # Calculate intersected areas and fractions
            intersection = pd.concat([self.shapefile.geometry[inp].reset_index(), mask.geometry[res].reset_index()], 
                                     axis=1, ignore_index=True)
            intersection.columns = (list(self.shapefile.geometry[inp].reset_index().columns) + 
                                    list(mask.geometry[res].reset_index().rename(columns={'geometry': 'geometry_mask', 
                                                                                          'index': 'index_mask'}).columns))
            intersection['area'] = intersection.apply(lambda x: x['geometry'].intersection(x['geometry_mask']).buffer(0).area, 
                                                      axis=1)
            intersection['fraction'] = intersection.apply(lambda x: x['area'] / x['geometry'].area, axis=1)
            
            # Choose biggest area from intersected areas with multiple options
            intersection.sort_values('fraction', ascending=False, inplace=True)
            intersection = intersection.drop_duplicates(subset='FID', keep="first")
            intersection = intersection.sort_values('FID').set_index('FID')

            # Get data from mask
            del mask['geometry']
            self.shapefile.loc[intersection.index, mask.columns] = np.array(mask.loc[intersection.index_mask, :])

        # Centroids that fall on the mask polygons, outside of the mask there will be NaN
        elif method is None:

            # Make copy
            shapefile_aux = deepcopy(self.shapefile)

            # Get centroids of shapefile to mask
            shapefile_aux.geometry = shapefile_aux.centroid

            # Calculate spatial joint
            shapefile_aux = gpd.sjoin(shapefile_aux, mask.to_crs(shapefile_aux.crs))

            # Get data from shapes where there are centroids, rest will be NaN
            del shapefile_aux['geometry'], shapefile_aux['index_right']
            self.shapefile.loc[shapefile_aux.index, shapefile_aux.columns] = shapefile_aux
            
        return None

    def __gather_data_py_object(self):
        """
        Gather all the variable data into the MPI rank 0 to perform a serial write.

        Returns
        -------
        data_list: dict
            Variables dictionary with all the data from all the ranks.
        """

        data_list = deepcopy(self.variables)
        for var_name in data_list.keys():
            try:
                # noinspection PyArgumentList
                data_aux = self.comm.gather(data_list[var_name]['data'], root=0)
                if self.rank == 0:
                    shp_len = len(data_list[var_name]['data'].shape)
                    add_dimension = False  # to Add a dimension
                    if self.parallel_method == 'Y':
                        if shp_len == 2:
                            # if is a 2D concatenate over first axis
                            axis = 0
                        elif shp_len == 3:
                            # if is a 3D concatenate over second axis
                            axis = 1
                        else:
                            # if is a 4D concatenate over third axis
                            axis = 2
                    elif self.parallel_method == 'X':
                        if shp_len == 2:
                            # if is a 2D concatenate over second axis
                            axis = 1
                        elif shp_len == 3:
                            # if is a 3D concatenate over third axis
                            axis = 2
                        else:
                            # if is a 4D concatenate over forth axis
                            axis = 3
                    elif self.parallel_method == 'T':
                        if shp_len == 2:
                            # if is a 2D add dimension
                            add_dimension = True
                            axis = None  # Not used
                        elif shp_len == 3:
                            # if is a 3D concatenate over first axis
                            axis = 0
                        else:
                            # if is a 4D concatenate over second axis
                            axis = 0
                    else:
                        raise NotImplementedError(
                            "Parallel method '{meth}' is not implemented. Use one of these: {accept}".format(
                                meth=self.parallel_method, accept=['X', 'Y', 'T']))
                    if add_dimension:
                        data_list[var_name]['data'] = np.stack(data_aux)
                    else:
                        data_list[var_name]['data'] = np.concatenate(data_aux, axis=axis)
            except Exception as e:
                print("**ERROR** an error has occurred while gathering the '{0}' variable.\n".format(var_name))
                sys.stderr.write("**ERROR** an error has occurred while gathering the '{0}' variable.\n".format(var_name))
                print(e)
                sys.stderr.write(str(e))
                # print(e, file=sys.stderr)
                sys.stderr.flush()
                self.comm.Abort(1)
                raise e

        return data_list

    def _gather_data(self):
        """
        Gather all the variable data into the MPI rank 0 to perform a serial write.

        Returns
        -------
        data_list: dict
            Variables dictionary with all the data from all the ranks.
        """

        data_list = deepcopy(self.variables)
        for var_name in data_list.keys():
            if self.info and self.master:
                print("Gathering {0}".format(var_name))
            shp_len = len(data_list[var_name]['data'].shape)
            try:
                # Collect local array sizes using the high-level mpi4py gather
                rank_shapes = np.array(self.comm.gather(data_list[var_name]['data'].shape, root=0))
                sendbuf = data_list[var_name]['data'].flatten()
                sendcounts = np.array(self.comm.gather(len(sendbuf), root=0))
                if self.master:
                    recvbuf = np.empty(sum(sendcounts), dtype=type(sendbuf[0]))
                else:
                    recvbuf = None
                self.comm.Gatherv(sendbuf=sendbuf, recvbuf=(recvbuf, sendcounts), root=0)
                if self.master:
                    recvbuf = np.split(recvbuf, np.cumsum(sendcounts))
                    # TODO ask
                    # I don't understand why it is giving one more split
                    if len(recvbuf) > len(sendcounts):
                        recvbuf = recvbuf[:-1]
                    for i, shape in enumerate(rank_shapes):
                        recvbuf[i] = recvbuf[i].reshape(shape)
                    add_dimension = False  # to Add a dimension
                    if self.parallel_method == 'Y':
                        if shp_len == 2:
                            # if is a 2D concatenate over first axis
                            axis = 0
                        elif shp_len == 3:
                            # if is a 3D concatenate over second axis
                            axis = 1
                        else:
                            # if is a 4D concatenate over third axis
                            axis = 2
                    elif self.parallel_method == 'X':
                        if shp_len == 2:
                            # if is a 2D concatenate over second axis
                            axis = 1
                        elif shp_len == 3:
                            # if is a 3D concatenate over third axis
                            axis = 2
                        else:
                            # if is a 4D concatenate over forth axis
                            axis = 3
                    elif self.parallel_method == 'T':
                        if shp_len == 2:
                            # if is a 2D add dimension
                            add_dimension = True
                            axis = None  # Not used
                        elif shp_len == 3:
                            # if is a 3D concatenate over first axis
                            axis = 0
                        else:
                            # if is a 4D concatenate over second axis
                            axis = 0
                    else:
                        raise NotImplementedError(
                            "Parallel method '{meth}' is not implemented. Use one of these: {accept}".format(
                                meth=self.parallel_method, accept=['X', 'Y', 'T']))
                    if add_dimension:
                        data_list[var_name]['data'] = np.stack(recvbuf)
                    else:
                        data_list[var_name]['data'] = np.concatenate(recvbuf, axis=axis)
            except Exception as e:
                print("**ERROR** an error has occurred while gathering the '{0}' variable.\n".format(var_name))
                sys.stderr.write("**ERROR** an error has occurred while gathering the '{0}' variable.\n".format(var_name))
                print(e)
                sys.stderr.write(str(e))
                # print(e, file=sys.stderr)
                sys.stderr.flush()
                self.comm.Abort(1)
                raise e

        return data_list

    # ==================================================================================================================
    #                                            Extra Methods
    # ==================================================================================================================
    
    def add_4d_vertical_info(self, info_to_add):
        """
        To add the vertical information from other source.

        Parameters
        ----------
        info_to_add : nes.Nes, str
            Nes object with the vertical information as variable or str with the path to the NetCDF file that contains
            the vertical data.
        """

        return vertical_interpolation.add_4d_vertical_info(self, info_to_add)

    def interpolate_vertical(self, new_levels, new_src_vertical=None, kind='linear', extrapolate=None, info=None):
        """
        Vertical interpolation method.

        Parameters
        ----------
        self : Nes
            Source Nes object.
        new_levels : list
            List of new vertical levels.
        new_src_vertical
        kind : str
            Vertical interpolation type.
        extrapolate : None, tuple, str
            Extrapolate method (for non linear operations).
        info: None, bool
            Indicates if you want to print extra information.
        """

        return vertical_interpolation.interpolate_vertical(
            self, new_levels, new_src_vertical=new_src_vertical, kind=kind, extrapolate=extrapolate, info=info)

    def interpolate_horizontal(self, dst_grid, weight_matrix_path=None, kind='NearestNeighbour', n_neighbours=4,
                               info=False, to_providentia=False):
        """
        Horizontal interpolation from the current grid to another one.

        Parameters
        ----------
        dst_grid : nes.Nes
            Final projection Nes object.
        weight_matrix_path : str, None
            Path to the weight matrix to read/create.
        kind : str
            Kind of horizontal interpolation. choices = ['NearestNeighbour'].
        n_neighbours: int
            Used if kind == NearestNeighbour. Number of nearest neighbours to interpolate. Default: 4.
        info: bool
            Indicates if you want to print extra info during the interpolation process.
        to_providentia : bool
            Indicates if we want the interpolated grid in Providentia format.
        """
        
        return horizontal_interpolation.interpolate_horizontal(
            self, dst_grid, weight_matrix_path=weight_matrix_path, kind=kind, n_neighbours=n_neighbours, info=info, 
            to_providentia=to_providentia)
