#!/usr/bin/env python

import sys
import warnings
import numpy as np
import os
import nes
from netCDF4 import Dataset
from mpi4py import MPI
from copy import copy


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._create_dimension_variables(netcdf)
    if self.info:
        print("Rank {0:03d}: Dimensions done".format(self.rank))

    # Create cell measures
    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(variables):
    """
    Change the data values according to the MONARCH conventions

    Parameters
    ----------
    variables : dict
        Variables data

    Returns
    -------
    dict
        Variable in the MONARCH units
    """
    for var_name in variables.keys():
        if variables['units'] == 'mol.s-1.m-2':
            # Kmol to mol
            variables['data'] *= 1000
        elif variables['units'] == 'kg.s-1.m-2':
            # No action needed
            pass
        else:
            raise TypeError("The unit '{0}' of specie {1} is not defined correctly. ".format(
                variables['units'], var_name) +
                            "Should be 'mol.s-1.m-2' or 'kg.s-1.m-2'")
    return variables

