#!/usr/bin/env python

import numpy as np
import nes
from netCDF4 import Dataset
from mpi4py import MPI


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

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

    self.to_dtype(np.float32)

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

    # Create dimensions
    self._create_dimensions(netcdf)

    # Create dimension variables
    self._lev['data'] = np.array(self._lev['data'], dtype=np.float32)
    self._lat['data'] = np.array(self._lat['data'], dtype=np.float32)
    self._lat_bnds['data'] = np.array(self._lat_bnds['data'], dtype=np.float32)
    self._lon['data'] = np.array(self._lon['data'], dtype=np.float32)
    self._lon_bnds['data'] = np.array(self._lon_bnds['data'], dtype=np.float32)
    if isinstance(self, nes.RotatedNes):
        self._rlat['data'] = np.array(self._rlat['data'], dtype=np.float32)
        self._rlon['data'] = np.array(self._rlon['data'], dtype=np.float32)
    if isinstance(self, nes.LCCNes) or isinstance(self, nes.MercatorNes):
        self._y['data'] = np.array(self._y['data'], dtype=np.float32)
        self._x['data'] = np.array(self._x['data'], dtype=np.float32)

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

    # Create cell measures
    if 'cell_area' in self.cell_measures.keys():
        self.cell_measures['cell_area']['data'] = np.array(self.cell_measures['cell_area']['data'], dtype=np.float32)
    self._create_cell_measures(netcdf)

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

    if keep_open:
        self.netcdf = netcdf
    else:
        netcdf.close()

    return None


def to_monarch_units(self):
    """
    Change the data values according to the MONARCH conventions.

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

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

    for var_name in self.variables.keys():
        if isinstance(self.variables[var_name]['data'], np.ndarray):
            if self.variables[var_name]['units'] == 'mol.s-1.m-2':
                # Kmol to mol
                self.variables[var_name]['data'] = np.array(self.variables[var_name]['data'] * 1000, dtype=np.float32)
            elif self.variables[var_name]['units'] == 'kg.s-1.m-2':
                # No unit change needed
                self.variables[var_name]['data'] = np.array(self.variables[var_name]['data'], dtype=np.float32)

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