#!/usr/bin/env python

import sys
import os
import warnings
import numpy as np
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
from copy import deepcopy
import datetime
from ..interpolation 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
    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, **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
        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']
        """

        # 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.print_info = info
        self.is_xarray = xarray
        self.__ini_path = path

        # 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.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 = self.__get_time(create_nes, **kwargs)
            self._time_bnds = self.__get_time_bnds(create_nes)
            self._lev = {'data': np.array([0]),
                         'units': '',
                         'positive': 'up'}
            self._lat, self._lon = self._create_centroids(**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)

            # 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'])

            # 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')

            # 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=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: ['Y', 'T']
        """
        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, balanced=balanced)
        return new

    def __del__(self):
        self.close()
        for var_name, var_info in self.variables.items():
            del var_info['data']
        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
        except AttributeError:
            pass

        del self

        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 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)

        Returns
        -------
        nessy : Nes
            Copy of the Nes object
        """
        nessy = deepcopy(self)
        nessy.netcdf = None
        if copy_vars:
            nessy.variables = nessy.__get_lazy_variables()
        else:
            nessy.variables = {}

        return nessy

    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 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:
            for var_name in var_list:
                if self.variables is not None:
                    if var_name in self.variables:
                        del self.variables[var_name]

        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

    # ==================================================================================================================
    #                                                 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

    # ==================================================================================================================
    #                                                 Reading
    # ==================================================================================================================
    def get_read_axis_limits(self):
        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}

        if self.parallel_method == 'Y':
            y_len = self._lat['data'].shape[0]
            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
            if self.rank + 1 < self.size:
                axis_limits['y_max'] = (y_len // self.size) * (self.rank + 1)
            # Spin up
            axis_limits['t_min'] = self.get_time_id(self.hours_start, first=True)
            axis_limits['t_max'] = self.get_time_id(self.hours_end, first=False)
        elif self.parallel_method == 'X':
            x_len = self._lon['data'].shape[-1]
            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
            if self.rank + 1 < self.size:
                axis_limits['x_max'] = (x_len // self.size) * (self.rank + 1)
            # Spin up
            axis_limits['t_min'] = self.get_time_id(self.hours_start, first=True)
            axis_limits['t_max'] = self.get_time_id(self.hours_end, first=False)
        elif self.parallel_method == 'T':
            first_time_idx = self.get_time_id(self.hours_start, first=True)
            last_time_idx = self.get_time_id(self.hours_end, first=False)
            t_len = last_time_idx - first_time_idx
            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) + first_time_idx
            if self.rank + 1 < self.size:
                axis_limits['t_max'] = ((t_len // self.size) * (self.rank + 1)) + first_time_idx

        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
        """
        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}

        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':
            first_time_idx = self.get_time_id(self.hours_start, first=True)
            last_time_idx = self.get_time_id(self.hours_end, first=False)
            len_to_split = last_time_idx - first_time_idx
            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]

        # 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 youy want to avoid from the first hours (True) or from the last (False)
            Default: True

        Returns
        -------
        int
            Possition 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 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
            Opened 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
            Opened 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 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

        Parameter
        ---------
        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

        Parameter
        ---------
        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, create_nes=False, **kwargs):
        """
        Get the NetCDF time values

        Returns
        -------
        time : list
            List of times (datetime) of the NetCDF data
        """
        if self.is_xarray:
            time = self.variables['time']
        else:
            if self.master:
                if create_nes:
                    if 'time' in kwargs:
                        time = kwargs['time']
                    else:
                        units = 'days since 1996-12-31 00:00:00'
                        calendar = 'standard'
                        time = num2date([0.], units=units, calendar=calendar)
                else:
                    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

        Returns
        -------
        time : 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_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):
        """
        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']
        Returns
        -------
        values : dict
            Dictionary with the portion of data corresponding to the rank
        """
        values = deepcopy(coordinate_info)
        if isinstance(coordinate_info, list):
            values = {'data': deepcopy(coordinate_info)}
        coordinate_len = len(values['data'].shape)

        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 = {}
                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

                    for attrname in var_info.ncattrs():
                        # Avoiding some attributes
                        if attrname not in ['missing_value', '_FillValue']:
                            value = getattr(var_info, attrname)
                            if value in ['unitless', '-']:
                                value = ''
                            variables[var_name][attrname] = value
            else:
                variables = None
            variables = self.comm.bcast(variables, root=0)

        return variables

    def _read_variable(self, var_name):
        """
        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.
        """
        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.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.print_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)
            if self.print_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 concatenate(self, aux_nessy):
        """
        Concatenate different variables into the same nes object

        Parameters
        ----------
        self : Nes

        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
        aux_nessy.load()

        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

        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):
        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 the 'lev' and 'time' dimension.

        Parameters
        ----------
        netcdf : Dataset
            netcdf4-python opened Dataset
        """
        netcdf.createDimension('time', None)
        if self._time_bnds is not None:
            netcdf.createDimension('time_nv', 2)
        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 'lev' and 'time' variables.

        Parameters
        ----------
        netcdf : Dataset
            netcdf4-python opened 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)
        lev.units = Units(self._lev['units'], formatted=True).units
        lev.positive = 'up'
        if self.size > 1:
            lev.set_collective(True)
        lev[:] = self._lev['data']

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

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

        return None

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

        Parameters
        ----------
        netcdf : Dataset
            netcdf4-python opened 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.print_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.print_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.print_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.print_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.print_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.print_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 hase occurred 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)

    def _create_centroids(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.print_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.print_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.print_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(self, path, compression_level=0, serial=False, info=False, chunking=False):
        """
        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.print_info
        self.print_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
                    new_nc.__to_netcdf_py(path)

            else:
                self.__to_netcdf_py(path, chunking=chunking)

        self.print_info = old_info

        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 hase occur while gathering the '{0}' variable.\n".format(var_name))
                sys.stderr.write("**ERROR** an error hase occur 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.print_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 hase occur while gathering the '{0}' variable.\n".format(var_name))
                sys.stderr.write("**ERROR** an error hase occur 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


# Methods
Nes.interpolate_vertical = interpolate_vertical
Nes.add_4d_vertical_info = add_4d_vertical_info
