#!/usr/bin/env python

import sys
import warnings
import numpy as np
from numpy.ma.core import MaskError
from copy import deepcopy
from nes.nc_projections.default_nes import Nes
from nes.nc_projections.points_nes import PointsNes


class PointsNesGHOST(PointsNes):
    """

    Attributes
    ----------
    _var_dim : tuple
        Tuple with the name of the Y and X dimensions for the variables.
        ('lat', 'lon') for a regular latitude-longitude projection.
    _lat_dim : tuple
        Tuple with the name of the dimensions of the Latitude values.
        ('lat',) for a regular latitude-longitude projection.
    _lon_dim : tuple
        Tuple with the name of the dimensions of the Longitude values.
        ('lon',) for a regular latitude-longitude projection.
    """
    def __init__(self, comm=None, path=None, info=False, dataset=None, xarray=False, parallel_method='X',
                 avoid_first_hours=0, avoid_last_hours=0, first_level=0, last_level=None, create_nes=False, 
                 **kwargs):
        """
        Initialize the PointsNes 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
        parallel_method : str
            Indicates the parallelization method that you want. Default over Y axis
            accepted values: ['Y', 'T']
        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.
        """

        super(PointsNesGHOST, self).__init__(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_nes=create_nes,
                                             **kwargs)

        self._flag = self._get_coordinate_dimension(['flag'])
        self.flag = self._get_coordinate_values(self._flag, 'X')

        self._qa = self._get_coordinate_dimension(['qa'])
        self.qa = self._get_coordinate_values(self._qa, 'X')

    @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 = PointsNesGHOST(comm=comm, path=path, info=info, dataset=dataset, xarray=xarray, balanced=balanced,
                             parallel_method=parallel_method, avoid_first_hours=avoid_first_hours,
                             avoid_last_hours=avoid_last_hours, first_level=first_level, last_level=last_level)
        return new

    def _create_dimensions(self, netcdf):
        """
        Create the 'N_flag_codes' and 'N_qa_codes' dimensions and the super dimensions ('time', 'station').

        Parameters
        ----------
        netcdf : Dataset
            NetCDF object.
        """

        super(PointsNesGHOST, self)._create_dimensions(netcdf)

        netcdf.createDimension('N_flag_codes', self._flag['data'].shape[2])
        netcdf.createDimension('N_qa_codes', self._qa['data'].shape[2])

        return None

    def _create_dimension_variables(self, netcdf):
        """
        Create the 'station' variables.

        Parameters
        ----------
        netcdf : Dataset
            NetCDF object.
        """

        super(PointsNesGHOST, self)._create_dimension_variables(netcdf)

        # N FLAG CODES
        flag = netcdf.createVariable('flag', np.int64, ('station', 'time', 'N_flag_codes',), 
                                     zlib=self.zip_lvl > 0, complevel=self.zip_lvl)
        flag.units = ""
        flag.axis = ""
        flag.long_name = ""
        flag.standard_name = "flag"
        if self.size > 1:
            flag.set_collective(True)
        flag[:] = self._flag['data']

        # N QA CODES
        qa = netcdf.createVariable('qa', np.int64, ('station', 'time', 'N_qa_codes',), 
                                   zlib=self.zip_lvl > 0, complevel=self.zip_lvl)
        qa.units = ""
        qa.axis = ""
        qa.long_name = ""
        qa.standard_name = "N_qa_codes"
        if self.size > 1:
            qa.set_collective(True)
        qa[:] = self._qa['data']

        self.free_vars(('flag', 'qa'))

    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: ['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 == '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['x_min']:self.read_axis_limits['x_max'], self.read_axis_limits['t_min']:self.read_axis_limits['t_max']]
            elif coordinate_len == 3:
                values['data'] = values['data'][self.read_axis_limits['x_min']:self.read_axis_limits['x_max'], self.read_axis_limits['t_min']:self.read_axis_limits['t_max'], :]
            else:
                raise NotImplementedError("The coordinate has wrong dimensions: {dim}".format(
                    dim=values['data'].shape))

        return values

    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

        # Remove strlen (maximum number of characters that a string can have) from dimensions and join characters
        if 'strlen' in nc_var.dimensions:
            nc_var = np.array([''.join(i) for i in np.char.decode(nc_var[:].data)])
            var_dims = tuple([', '.join(dim for dim in var_dims if dim != 'strlen')])

        # Read data in 1 or 2 dimensions
        # TODO: Ask Dene why x, t instead of t, x
        if len(var_dims) < 2:
            data = nc_var[self.read_axis_limits['x_min']:self.read_axis_limits['x_max']]
        elif len(var_dims) == 2:
            data = nc_var[self.read_axis_limits['x_min']:self.read_axis_limits['x_max'],
                          self.read_axis_limits['t_min']:self.read_axis_limits['t_max']]
        elif len(var_dims) == 3:
            data = nc_var[self.read_axis_limits['x_min']:self.read_axis_limits['x_max'],
                          self.read_axis_limits['t_min']:self.read_axis_limits['t_max'],
                          :]
        else:
            raise NotImplementedError('Error with {0}. Only can be read netCDF with 3 dimensions or less'.format(
                var_name))
                
        # Missing to nan
        try:
            data[data.mask == True] = np.nan
        except (AttributeError, MaskError, ValueError):
            pass

        return data

    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:

                    # Define dimensions depending on the type of variable
                    if len(var_dict['data'].shape) == 1:
                        # Metadata
                        var_dims = self._var_dim
                    elif len(var_dict['data'].shape) == 2:
                        # Different from metadata (e.g. concentrations of pm10)
                        var_dims = self._var_dim + ('time',)
                    else:
                        # Flags and qa variables
                        if var_name == 'flag':
                            var_dims = self._var_dim + ('time', 'N_flag_codes',)
                        elif var_name == 'qa':
                            var_dims = self._var_dim + ('time', 'N_qa_codes',)

                    # ESDAC iwahashi landform and other vars are given as objects, transform to strings
                    if var_dict['data'].dtype == np.dtype(object):
                        var_dtype = np.dtype(str)
                    else:
                        var_dtype = var_dict['data'].dtype

                    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_dtype, var_dims,
                                                        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_dtype, var_dims,
                                                        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 len(var_dict['data'].shape) == 1:
                                    try:
                                        var[self.write_axis_limits['x_min']:self.write_axis_limits['x_max']] = att_value
                                    except IndexError:
                                        raise IndexError("Different shapes. out_shape={0}, data_shp={1}".format(
                                            var[self.write_axis_limits['x_min']:self.write_axis_limits['x_max']].shape,
                                            att_value.shape))
                                elif len(var_dict['data'].shape) == 2:
                                    try:
                                        var[self.write_axis_limits['x_min']:self.write_axis_limits['x_max'],
                                            self.write_axis_limits['t_min']:self.write_axis_limits['t_max']] = att_value
                                    except IndexError:
                                        raise IndexError("Different shapes. out_shape={0}, data_shp={1}".format(
                                            var[self.write_axis_limits['x_min']:self.write_axis_limits['x_max'],
                                                self.write_axis_limits['t_min']:self.write_axis_limits['t_max']].shape,
                                            att_value.shape))  
                                elif len(var_dict['data'].shape) == 3:
                                    try:
                                        var[self.write_axis_limits['x_min']:self.write_axis_limits['x_max'],
                                            self.write_axis_limits['t_min']:self.write_axis_limits['t_max'],
                                            :] = att_value
                                    except IndexError:
                                        raise IndexError("Different shapes. out_shape={0}, data_shp={1}".format(
                                            var[self.write_axis_limits['x_min']:self.write_axis_limits['x_max'],
                                                self.write_axis_limits['t_min']:self.write_axis_limits['t_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 _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:
                    if self.parallel_method == 'X':
                        axis = 0
                    else:
                        raise NotImplementedError(
                            "Parallel method '{meth}' is not implemented. Use one of these: {accept}".format(
                                meth=self.parallel_method, accept=['X']))
                    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