#!/usr/bin/env python

import sys
import warnings
import numpy as np
import pandas as pd
from copy import deepcopy
from netCDF4 import Dataset, date2num, stringtochar
from numpy.ma.core import MaskError
from nes.nc_projections.default_nes import Nes


class PointsNes(Nes):
    """

    Attributes
    ----------
    _var_dim : tuple
        Tuple with the name of the Y and X dimensions for the variables.
        ('lat', 'lon') for a points grid.
    _lat_dim : tuple
        Tuple with the name of the dimensions of the Latitude values.
        ('lat',) for a points grid.
    _lon_dim : tuple
        Tuple with the name of the dimensions of the Longitude values.
        ('lon',) for a points grid.
    _station : tuple
        Tuple with the name of the dimensions of the station values.
        ('station',) for a points grid.
    """
    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, 
                 times=None, **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: ['X']
        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(PointsNes, 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,
                                        times=times, **kwargs)

        if create_nes:
            # Complete dimensions
            self._station = {'data': np.arange(len(self._lon['data']))}
            
            # Dimensions screening
            self.lat = self._get_coordinate_values(self._lat, 'X')
            self.lon = self._get_coordinate_values(self._lon, 'X')
            self.station = deepcopy(self._station)
        else:
            self._station = self._get_coordinate_dimension(['station'])
            self.station = self._get_coordinate_values(self._station, 'X')
        
        self.strlen = self._get_strlen(create_nes)

        # Set axis limits for parallel writing
        self.write_axis_limits = self.get_write_axis_limits()

        self._var_dim = ('station',)
        self._lat_dim = ('station',)
        self._lon_dim = ('station',)

    @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 = PointsNes(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 'lev', 'time_nv', 'station' dimensions.

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

        netcdf.createDimension('time', None)
        if self._time_bnds is not None:
            netcdf.createDimension('time_nv', 2)

        # The number of longitudes is equal to the number of stations
        netcdf.createDimension('station', len(self._lon['data']))

        if hasattr(self, 'strlen'):
            if self.strlen is not None:
                netcdf.createDimension('strlen', self.strlen)

        return None

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

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

        # 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, time_var.calendar)

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

        # STATIONS
        stations = netcdf.createVariable('station', np.float64, ('station',), zlib=self.zip_lvl > 0, complevel=self.zip_lvl)
        stations.units = ""
        stations.axis = "X"
        stations.long_name = ""
        stations.standard_name = "station"
        if self.size > 1:
            stations.set_collective(True)
        stations[:] = self._station['data']

        return None

    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.
        """

        nc_var = super(PointsNes, self)._get_coordinate_dimension(possible_names)

        if isinstance(possible_names, str):
            possible_names = [possible_names]

        if 'station' in possible_names:
            nc_var['data'] = np.arange(len(self._lon['data']))

        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: ['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['t_min']:self.read_axis_limits['t_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))

        return values

    def _get_lazy_variables_not_used(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
                    # Remove strlen as a dimension
                    if 'strlen' in var_info.dimensions:
                        variables[var_name]['dimensions'] = tuple([dim for dim in var_info.dimensions
                                                                   if dim != 'strlen'])
                    else:
                        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 _get_strlen(self, create_nes):
        """
        Read the string length dimension of some variables.

        Returns
        -------
        strlen: np.int
            String length.
        """

        if not create_nes:
            if 'strlen' in self.netcdf.dimensions:
                strlen = self.netcdf.dimensions['strlen'].size
            else:
                strlen = None
        else:
            # TODO: Check if strings are always divided into 75 in the real networks
            strlen = 75

        return strlen

    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
        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:
            if 'strlen' in nc_var.dimensions:
                data = nc_var[self.read_axis_limits['x_min']:self.read_axis_limits['x_max'], :]
            else:
                data = nc_var[self.read_axis_limits['t_min']:self.read_axis_limits['t_max'],
                              self.read_axis_limits['x_min']:self.read_axis_limits['x_max']]
        else:
            raise NotImplementedError('Error with {0}. Only can be read netCDF with 2 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:

                    # Set dimensions
                    if 'dimensions' in var_dict.keys(): 
                        var_dims = var_dict['dimensions']
                    else:
                        if len(var_dict['data'].shape) == 1:
                            # For data that depends only on station (e.g. station_code)
                            var_dims = self._var_dim
                        else:
                            # For data that is dependent on time and station (e.g. PM10)
                            var_dims = ('time',) + self._var_dim
                        
                        # Add strlen as a dimension if needed
                        if var_dict['dtype'] == np.str:
                            var_dims += ('strlen',)

                    # Ensure data is of type numpy array (to create NES)
                    if not isinstance(var_dict['data'], (np.ndarray, np.generic)):
                        try:
                            var_dict['data'] = var_dict['data'].to_numpy()
                        except AttributeError:
                            raise AttributeError("Data for variable {0} must be a numpy array (np.ndarray or np.generic)".var_name)

                    # Set dtype
                    if 'dtype' in var_dict.keys():
                        var_dtype = var_dict['dtype']
                        if var_dtype != var_dict['data'].dtype:
                            msg = 'WARNING!!! '
                            msg += 'Different data types for variable {0}'.format(var_name)
                            msg += 'Input dtype={0}, data dtype={1}'.format(var_dtype, 
                                                                            var_dict['data'].dtype)
                            warnings.warn(msg)
                            try:
                                var_dict['data'] = var_dict['data'].astype(var_dtype)
                            except Exception as e: # TODO: Detect exception
                                raise e('It was not possible to cast the data to the input dtype.')
                    else:
                        var_dtype = var_dict['data'].dtype

                    # Transform objects into strings (e.g. for ESDAC iwahashi landform in GHOST)
                    if var_dtype == np.dtype(object):
                        var_dtype = np.dtype(str)

                    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
                                # TODO: Change chunk size (add strlen) as tuple
                            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:
                                    if 'strlen' in var_dict['dimensions']:
                                        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))
                                    else:
                                        try:
                                            var[self.write_axis_limits['t_min']:self.write_axis_limits['t_max'],
                                                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['t_min']:self.write_axis_limits['t_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', 'dtype']:
                                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, var_info in data_list.items():
            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)
                    if self.parallel_method == 'X':
                        if shp_len == 1:
                            # concatenate over first axis (station) when dims are (station)
                            axis = 0
                        elif shp_len == 2:
                            if 'strlen' in var_info['dimensions']:
                                # concatenate over first axis (station) when dims are (station, strlen)
                                axis = 0
                            else:
                                # concatenate over second axis (station) when dims are (time, station)
                                axis = 1     
                    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

    def _create_centroids(self, **kwargs):
        """
        Calculate center latitudes and longitudes from points.

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

        # Calculate center latitudes
        self.center_lats = kwargs['lat']

        # Calculate center longitudes
        self.center_lons = kwargs['lon']

        return {'data': self.center_lats}, {'data': self.center_lons}

    def _create_metadata(self, netcdf):
        """
        Create metadata variables

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

        # 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 set_read_axis_limits(self):
        """
        Calculate the 4D reading axis limits

        Returns
        -------
        dict
            Dictionary with the 2D limits of the rank data to read.
            t_min, t_max, x_min and x_max
        """

        axis_limits = {'x_min': None, 'x_max': None,
                       't_min': None, 't_max': None}

        if self.parallel_method == 'X':
            x_len = self._lon['data'].shape[-1]
            if x_len < self.size:
                raise IndexError('More processors (size={0}) selected than X elements (size={1})'.format(self.size, x_len))
            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)
        else:
            raise NotImplementedError("Parallel method '{meth}' is not implemented. Use one of these: {accept}".format(
                meth=self.parallel_method, accept=['X']))

        return axis_limits
