#!/usr/bin/env python

import sys
import os
from urllib.parse import _NetlocResultMixinStr
import warnings
import numpy as np
from pandas import DataFrame, read_csv
import timeit
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

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 choosen 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, 
                 **kwargs):
        """
        Initialize the Nes class

        Parameters
        ----------
        comm: MPI.COMM
            Path to the CSV file that contains all the information.
        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']
        """
        # 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

        # NetCDF object
        if create_nes:
            self.variables = None

            self._time = self.__get_time(create_nes)
            self._time_bnds = self.__get_time_bnds(create_nes)

            self._lev = {'data': np.array([0]), 
                         'units': '',
                         'positive': 'up'}
            self.lev = deepcopy(self._lev)

            self._lat, self._lon = self._create_centroids(**kwargs)

            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 dimension
            self._time = self.__get_time()
            self._lev = self._get_coordinate_dimension(['lev', 'level', 'lm'])
            self._lat = self._get_coordinate_dimension(['lat', 'latitude'])
            self._lon = self._get_coordinate_dimension(['lon', 'longitude'])
            self._time_bnds = self.__get_time_bnds()

            # Axis limits
            self.parallel_method = parallel_method   
            self.read_axis_limits = self.set_read_axis_limits()

            # Dimensions screening
            self.time = self._time[self.read_axis_limits['t_min']:self.read_axis_limits['t_max']]
            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.time_bnds = self._time_bnds

            self.write_axis_limits = self.set_write_axis_limits()

            # 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

    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.set_read_axis_limits()
        self.write_axis_limits = self.set_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(f'{time_bnd} is not a datetime object')
                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. Time bounds will not be set.'
                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 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
            Position 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 set_read_axis_limits(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]
            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]
            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
            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=['Y']))

        # Verical 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 set_write_axis_limits(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 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

    @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):
        """
        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:
                    units = "days since 1996-12-31 00:00:00"
                    calendar = "gregorian"
                    time = num2date([0.], units=units, calendar=calendar)
                else:
                    nc_var = self.netcdf.variables['time']
                    time = num2date(nc_var[:], self.__parse_time_unit(nc_var.units), calendar=nc_var.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]
        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 nc_var['units'] in ['unitless', '-']:
                nc_var['units'] = ''

        self.free_vars(dimension_name)

        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
        if 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('Only can be read netCDF with 4 dimensions or less')

        # Missing to nan
        try:
            data[data.mask == True] = np.nan
        except (AttributeError, MaskError):
            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 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 __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

    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 _create_dimensions(self, netcdf):
        """
        Create the 'lev' and 'time' dimension.

        Parameters
        ----------
        netcdf : Dataset
            netcdf4-python opened Dataset
        """
        netcdf.createDimension('time', None)
        netcdf.createDimension('lev', len(self.lev['data']))
        if self._time_bnds is not None:
            netcdf.createDimension('time_nv', 2)

        return None

    def _create_dimension_variables(self, netcdf):
        """
        Create the 'lev', 'time', 'lat' and 'lon' 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, calendar='standard')

        # 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
        """
        if self.variables is not None:
            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.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':

                                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 _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._set_crs(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')

        self.comm.Barrier()
        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(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
