#!/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_cams_ra(self, path):
    """
    Horizontal interpolation from one grid to another one.

    Parameters
    ----------
    self : nes.Nes
        Source projection Nes Object.
    path : str
            Path to the output netCDF file.
    """
    
    if not isinstance(self, nes.LatLonNes):
        raise TypeError("CAMS Re-Analysis format must have Regular Lat-Lon projection")
    if '<level>' not in path:
        raise ValueError("AMS Re-Analysis path must contain '<level>' as pattern; current: '{0}'".format(path))
    
    orig_path = copy(path)
    
    for i_lev, level in enumerate(self.lev['data']):
        path = orig_path.replace('<level>', 'l{0}'.format(i_lev))
        # 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))
        self.to_dtype(data_type=np.float32)

        # Create dimensions
        create_dimensions(self, netcdf)

        # Create variables
        create_variables(self, netcdf, i_lev)

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

        # 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.close()

    return None


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

    Parameters
    ----------
    self : nes.Nes
        Source projection Nes Object.
    netcdf : Dataset
        netcdf4-python open dataset.
    """

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

    # Create lev, lon and lat dimensions
    netcdf.createDimension('lat', len(self._lat['data']))
    netcdf.createDimension('lon', len(self._lon['data']))

    return None


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

    Parameters
    ----------
    self : nes.Nes
        Source projection Nes Object.
    netcdf : Dataset
        netcdf4-python open dataset.
    """

    # LATITUDES
    lat = netcdf.createVariable('lat', np.float64, ('lat',))
    lat.standard_name = 'latitude'
    lat.long_name = 'latitude'
    lat.units = 'degrees_north'
    lat.axis = 'Y'

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

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

    # TIMES
    time_var = netcdf.createVariable('time', np.float64, ('time',))
    time_var.standard_name = 'time'
    time_var.units = 'day as %Y%m%d.%f'
    time_var.calendar = 'proleptic_gregorian'
    time_var.axis = 'T'
    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)

    return None


def create_variables(self, netcdf, i_lev):
    """
    Create the netCDF file variables.

    Parameters
    ----------
    self : nes.Nes
        Source projection Nes Object.
    netcdf : Dataset
        netcdf4-python open dataset.
    """

    for i, (var_name, var_dict) in enumerate(self.variables.items()):
        if var_dict['data'] is not None:
            if self.info:
                print("Rank {0:03d}: Writing {1} var ({2}/{3})".format(self.rank, var_name, i + 1, len(self.variables)))
            try:
                var = netcdf.createVariable(var_name, np.float32, ('time', 'lat', 'lon',),
                                            zlib=True, complevel=7, least_significant_digit=3)

                if self.info:
                    print("Rank {0:03d}: Var {1} created ({2}/{3})".format(
                        self.rank, var_name, i + 1, len(self.variables)))
                if self.size > 1:
                    var.set_collective(True)
                    if self.info:
                        print("Rank {0:03d}: Var {1} collective ({2}/{3})".format(
                            self.rank, var_name, i + 1, len(self.variables)))

                if self.info:
                    print("Rank {0:03d}: Filling {1})".format(self.rank, var_name))
                var[self.write_axis_limits['t_min']:self.write_axis_limits['t_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']] = var_dict['data'][:, i_lev, :, :]

                if self.info:
                    print("Rank {0:03d}: Var {1} data ({2}/{3})".format(
                        self.rank, var_name, i + 1, len(self.variables)))
                var.long_name = var_dict['long_name']
                var.units = var_dict['units']
                var.number_of_significant_digits = np.int32(3)

                if self.info:
                    print("Rank {0:03d}: Var {1} completed ({2}/{3})".format(self.rank, var_name, i + 1,
                                                                             len(self.variables)))
            except Exception as e:
                print("**ERROR** an error has occurred while writing the '{0}' variable".format(var_name))
                # print("**ERROR** an error has occurredred while writing the '{0}' variable".format(var_name),
                #       file=sys.stderr)
                raise e
        else:
            msg = 'WARNING!!! '
            msg += 'Variable {0} was not loaded. It will not be written.'.format(var_name)
            warnings.warn(msg)

    return None


def date2num(time_array, time_units=None, time_calendar=None):
    
    time_res = []
    for aux_time in time_array:
        time_res.append(float(aux_time.strftime("%Y%m%d")) + (float(aux_time.strftime("%H")) / 24))
    time_res = np.array(time_res, dtype=np.float64)

    return time_res
