#!/usr/bin/env python

import sys
import warnings
import numpy as np
import pandas as pd
from copy import deepcopy
import geopandas as gpd
from netCDF4 import date2num, stringtochar
from numpy.ma.core import MaskError
from .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, strlen=75, **kwargs):
        """
        Initialize the PointsNes class.

        Parameters
        ----------
        comm: MPI.Communicator
            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, None
            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: 'X'.
            accepted values: ['X', 'T'].
        strlen: int
            Maximum length of strings in NetCDF. Default: 75.
        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:
            # Dimensions screening
            self.lat = self._get_coordinate_values(self._lat, 'X')
            self.lon = self._get_coordinate_values(self._lon, 'X')
            self.strlen = strlen
        else:
            # Dimensions screening
            self.strlen = self._get_strlen()

        # Complete dimensions
        self._station = {'data': np.arange(len(self._lon['data']))}

        # Dimensions screening
        self.station = self._get_coordinate_values(self._station, 'X')

        # 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='X', 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: 'X'.
            accepted values: ['X', '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 'lev', 'time_nv', 'station', 'spatial_nv' and 'strlen' dimensions.

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

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

        # Create time_nv (number of vertices) dimension
        if self._time_bnds is not None:
            netcdf.createDimension('time_nv', 2)

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

        # Create string length dimension
        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', 'station', 'lat', 'lat_bnds', 'lon' and 'lon_bnds' 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']

        # LATITUDES
        lat = netcdf.createVariable('lat', np.float64, self._lat_dim, 
                                     zlib=self.zip_lvl > 0, complevel=self.zip_lvl)
        lat.units = 'degrees_north'
        lat.axis = 'Y'
        lat.long_name = 'latitude coordinate'
        lat.standard_name = 'latitude'
        if self._lat_bnds is not None:
            lat.bounds = 'lat_bnds'
        if self.size > 1:
            lat.set_collective(True)
        lat[:] = self._lat['data']

        # LONGITUDES
        lon = netcdf.createVariable('lon', np.float64, self._lon_dim, 
                                     zlib=self.zip_lvl > 0, complevel=self.zip_lvl)
        lon.units = 'degrees_east'
        lon.axis = 'X'
        lon.long_name = 'longitude coordinate'
        lon.standard_name = 'longitude'
        if self._lon_bnds is not None:
            lon.bounds = 'lon_bnds'
        if self.size > 1:
            lon.set_collective(True)
        lon[:] = self._lon['data']

        return None

    def _get_coordinate_values(self, coordinate_info, coordinate_axis, bounds=False):
        """
        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'].
        bounds : bool
            Boolean variable to know if there are coordinate bounds.
        Returns
        -------
        values : dict
            Dictionary with the portion of data corresponding to the rank.
        """

        if coordinate_info is None:
            return None

        if not isinstance(coordinate_info, dict):
            values = {'data': deepcopy(coordinate_info)}
        else:
            values = deepcopy(coordinate_info)

        coordinate_len = len(values['data'].shape)
        if bounds:
            coordinate_len -= 1

        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_strlen(self):
        """
        Read the string length dimension of some variables.

        Returns
        -------
        int, None
            String length. None means no string data.
        """

        if 'strlen' in self.netcdf.dimensions:
            strlen = self.netcdf.dimensions['strlen'].size
        else:
            strlen = None

        return strlen

    def _read_variable(self, var_name, nc_var=None, var_dims=None):
        """
        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.
        """

        if nc_var is None and var_dims is None:
            nc_var = self.netcdf.variables[var_name]
            var_dims = nc_var.dimensions

        # 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 var_dims:
                data = nc_var[self.read_axis_limits['x_min']:self.read_axis_limits['x_max'], :]
                data = np.array([''.join(i.tostring().decode('ascii').replace('\x00', '')) for i in data])
            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 open 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:

                    # Get data type
                    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_dict['data'] = var_dict['data'].astype(str)
                        var_dtype = var_dict['data'].dtype

                    # Get dimensions when reading datasets
                    if 'dimensions' in var_dict.keys():
                        var_dims = var_dict['dimensions']
                    # Get dimensions when creating new datasets
                    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

                    # Convert list of strings to chars for parallelization
                    try:
                        unicode_type = len(max(var_dict['data'], key=len))
                        if ((var_dict['data'].dtype == np.dtype('<U' + str(unicode_type))) or 
                            (var_dict['data'].dtype == np.str)):
                            # Add strlen as a dimension if needed
                            if self.strlen is not None:
                                if 'strlen' not in var_dims:
                                    var_dims += ('strlen',)
                                var_dict['data'] = stringtochar(np.array([v.encode('ascii', 'ignore') 
                                                                for v in var_dict['data']]).astype('S' + str(self.strlen)))
                                var_dtype = 'S' + str(self.strlen)
                    except:
                        pass
                    
                    # 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.".format(var_name))

                    if self.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.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)))
                        
                        for att_name, att_value in var_dict.items():
                            if att_name == 'data':
                                if len(att_value.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))
                                    except ValueError:
                                            raise ValueError("Axis limits cannot be accessed. 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(att_value.shape) == 2:
                                    if 'strlen' in var_dims:
                                        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))
                                        except ValueError:
                                            raise ValueError("Axis limits cannot be accessed. 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))
                                        except ValueError:
                                            raise ValueError("Axis limits cannot be accessed. 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.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.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 _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':
                        # concatenate over station
                        if shp_len == 1:
                            # dimensions = (station)
                            axis = 0
                        elif shp_len == 2:
                            if 'strlen' in var_info['dimensions']:
                                # dimensions = (station, strlen)
                                axis = 0
                            else:
                                # dimensions = (time, station)
                                axis = 1  
                        else:
                            msg = "The points NetCDF must have "
                            msg += "surface values (without levels)."
                            raise NotImplementedError(msg)
                    elif self.parallel_method == 'T':
                        # concatenate over time
                        if shp_len == 1: 
                            # dimensions = (station)
                            axis = None
                            continue
                        elif shp_len == 2:
                            if 'strlen' in var_info['dimensions']:
                                # dimensions = (station, strlen)
                                axis = None
                                continue
                            else:
                                # dimensions = (time, station)
                                axis = 0
                        else:
                            raise NotImplementedError("The points NetCDF must only have surface values (without levels).")
                    else:
                         raise NotImplementedError(
                            "Parallel method '{meth}' is not implemented. Use one of these: {accept}".format(
                                meth=self.parallel_method, accept=['X', 'T']))
                    data_list[var_name]['data'] = np.concatenate(data_aux, axis=axis)
            except Exception as e:
                print("**ERROR** an error has occurred while gathering the '{0}' variable.\n".format(var_name))
                sys.stderr.write("**ERROR** an error has occurred 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_centre_coordinates(self, **kwargs):
        """
        Calculate centre latitudes and longitudes from points.

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

        # Calculate centre latitudes
        centre_lat_data = kwargs['lat']
        centre_lat = {'data': centre_lat_data}

        # Calculate centre longitudes
        centre_lon_data = kwargs['lon']
        centre_lon = {'data': centre_lon_data}

        return centre_lat, centre_lon

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

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

        return None

    def create_spatial_bounds(self):
        """
        Calculate longitude and latitude bounds and set them.
        """

        raise NotImplementedError("Spatial bounds cannot be created for points datasets.")

    def to_providentia(self, model_centre_lon, model_centre_lat, grid_edge_lon, grid_edge_lat):
        """
        Transform a PointsNes into a PointsNesProvidentia object
        
        Returns
        ----------
        points_nes_providentia : nes.Nes
            Points Nes Providentia Object
        """

        from .points_nes_providentia import PointsNesProvidentia

        points_nes_providentia = PointsNesProvidentia(comm=self.comm, 
                                                      info=self.info, 
                                                      balanced=self.balanced, 
                                                      parallel_method=self.parallel_method, 
                                                      avoid_first_hours=self.hours_start, 
                                                      avoid_last_hours=self.hours_end, 
                                                      first_level=self.first_level, 
                                                      last_level=self.last_level, 
                                                      create_nes=True,
                                                      times=self.time,
                                                      model_centre_lon=model_centre_lon,
                                                      model_centre_lat=model_centre_lat,
                                                      grid_edge_lon=grid_edge_lon,
                                                      grid_edge_lat=grid_edge_lat,
                                                      lat=self.lat['data'], 
                                                      lon=self.lon['data']
                                                      )
        
        # Convert dimensions (time, lev, lat, lon) to (station, time) for interpolated variables and reshape data
        variables = {}
        interpolated_variables = deepcopy(self.variables)
        for var_name, var_info in interpolated_variables.items():
            variables[var_name] = {}
            # ('time', 'lev', 'lat', 'lon') or ('time', 'lat', 'lon') to ('station', 'time')
            if len(var_info['dimensions']) != len(var_info['data'].shape):
                variables[var_name]['data'] = var_info['data'].T
                variables[var_name]['dimensions'] = ('station', 'time')
            else:
                variables[var_name]['data'] = var_info['data']
                variables[var_name]['dimensions'] = var_info['dimensions']

        # Set variables
        points_nes_providentia.variables = variables

        return points_nes_providentia

    def to_grib2(self, path, grib_keys, grib_template_path, lat_flip=False, info=False):
        """
        Write output file with grib2 format.

        Parameters
        ----------
        path : str
            Path to the output file.
        grib_keys : dict
            Dictionary with the grib2 keys.
        grib_template_path : str
            Path to the grib2 file to use as template.
        info : bool
            Indicates if you want to print extra information during the process.
        """

        raise NotImplementedError("Grib2 format cannot be written with point data.")

    def create_shapefile(self):
        """
        Create spatial geodataframe (shapefile).
        """

        # Create dataframe cointaining all points
        geometry = gpd.points_from_xy(self.lon['data'], self.lat['data'])
        fids = np.arange(len(self._lon['data']))
        fids = fids[self.read_axis_limits['x_min']:self.read_axis_limits['x_max']]
        gdf = gpd.GeoDataFrame(index=pd.Index(name='FID', data=fids), 
                               geometry=geometry, 
                               crs="EPSG:4326")
        self.shapefile = gdf

        return gdf

    @staticmethod
    def _get_axis_index_(axis):
        if axis == 'T':
            value = 0
        elif axis == 'X':
            value = 1
        else:
            raise ValueError("Unknown axis: {0}".format(axis))
        return value

    @staticmethod
    def _set_var_crs(var):
        """
        Set the grid_mapping

        Parameters
        ----------
        var : Variable
            netCDF4-python variable object.
        """
        var.coordinates = "lat lon"

        return None
