from dask import delayed
from dask import array
import xarray as xr
import numpy as np
from dask.distributed import Client
from warnings import warn
import copy
import scipy as sp

def setup_dask_clinet(workers=2, threads=2, memory_limit_per_worker='2GB'):
    Client(n_workers=workers, threads_per_worker=threads, memory_limit=memory_limit_per_worker)


class COAsT:
    def __init__(self, file = None, chunks: dict=None, multiple=False,
                 workers=2, threads=2, memory_limit_per_worker='2GB'):
        # self.client = Client(n_workers=workers, threads_per_worker=threads, memory_limit=memory_limit_per_worker)
        self.dataset = None
        self.earth_raids = 6371.007176 # Radius of the earth in km
        self.set_dimension_mapping()
        self.set_variable_mapping()
        #self.set_grid_ref_attr()
        if file is None:
            print("Object created but no file or directory specified: " + 
                  str(self) + " \n" +
                  "Use COAsT.load() to load a netCDF file from file path or " +
                  "directory into this object.")
        else:
            self.load(file, chunks, multiple)

    def load(self, file_or_dir, chunks:dict=None, multiple=False):
        """
        Loads a file into a COAsT object's dataset variable using xarray
        
        Args:
            file (str)     : file name or directory to multiple files.
            chunks (dict)  : Chunks to use in Dask [default None]
            multiple (bool): If true, load in multiple files from directory.
                             If false load a single file [default False]
        """
        if multiple:
            self.load_multiple(file_or_dir, chunks)
        else:
            self.load_single(file_or_dir, chunks)    

    def __getitem__(self, name: str):
        return self.dataset[name]

    def load_single(self, file, chunks: dict = None):
        """ Loads a single file into COAsT object's dataset variable. """
        self.dataset = xr.open_dataset(file, chunks=chunks)

    def load_multiple(self, directory_to_files, chunks: dict = None):
        """ Loads multiple files from directory into dataset variable. """
        self.dataset = xr.open_mfdataset(
            directory_to_files, chunks=chunks, parallel=True, 
            combine="by_coords") #, compat='override')

    def load_dataset(self, dataset):
        """        
        :param dataset: The dataset to use
        :type dataset: xarray.Dataset
        """
        
        self.dataset = dataset
    
    def set_dimension_mapping(self):
        self.dim_mapping = None
        
    def set_variable_mapping(self):
        self.var_mapping = None

    def set_grid_ref_attr(self):
        self.grid_ref_attr_mapping = None

    def set_dimension_names(self, dim_mapping: dict):
        """ 
        Relabel dimensions in COAsT object xarray.dataset to ensure
        consistent naming throughout the COAsT package.
        
        Args:
            dim_mapping (dict): keys are dimension names to change and values
                                new dimension names
        """
        if dim_mapping is None: return
        for key, value in dim_mapping.items():
            try:
                self.dataset = self.dataset.rename_dims({ key : value })
            except:
                pass
                #print(str(self) + ': Problem renaming dimension: ' + 
                #      key + ' -> ' + value)
                
    def set_variable_names(self, var_mapping: dict):
        """ 
        Relabel variables in COAsT object xarray.dataset to ensure
        consistent naming throughout the COAsT package.
        
        Args:
            var_mapping (dict): keys are variable names to change and values
                                are new variable names
        """
        if var_mapping is None: return
        for key, value in var_mapping.items():
            try:
                self.dataset = self.dataset.rename_vars({ key : value })
            except:
                pass
                #print(str(self) + ': Problem renaming variable: ' + 
                #      key + ' -> ' + value)

    def set_variable_grid_ref_attribute(self, grid_ref_attr_mapping: dict):
        """
        Set attributes for variables to access within package.
        Set grid attributes to identify with grid variable is associated with.
        """
        if grid_ref_attr_mapping is None: return
        for key, value in grid_ref_attr_mapping.items():
            try:
                self.dataset[key].attrs['grid_ref'] = value
            except:
                pass
                #print(str(self) + ': Problem assigning ' + key +
                #        'grid_ref attribute ' + value)

    def copy(self):
        return copy.copy(self)
        
    def isel(self, indexers: dict = None, drop: bool = False,
             **kwargs):
        '''
        Indexes COAsT object along specified dimensions using xarray isel. 
        Input is of same form as xarray.isel. Basic use, hand in either:
            1. Dictionary with keys = dimensions, values = indices
            2. **kwargs of form dimension = indices
        '''
        obj_copy = self.copy()
        obj_copy.dataset = obj_copy.dataset.isel(indexers, drop, **kwargs)
        return obj_copy
    
    def sel(self, indexers: dict = None, drop: bool = False,
             **kwargs):
        '''
        Indexes COAsT object along specified dimensions using xarray sel. 
        Input is of same form as xarray.sel. Basic use, hand in either:
            1. Dictionary with keys = dimensions, values = indices
            2. **kwargs of form dimension = indices
        '''
        obj_copy = self.copy()
        obj_copy.dataset = obj_copy.dataset.sel(indexers, drop, **kwargs)
        return obj_copy
    
    def rename(self, rename_dict, inplace: bool=None, **kwargs):
        self.dataset = self.dataset.rename(rename_dict, inplace, **kwargs)
        return

    def subset(self, **kwargs):
        '''
        Subsets all variables within the dataset inside self (a COAsT object).
        Input is a set of keyword argument pairs of the form:
            dimension_name = indices
        The entire object is then subsetted along this dimension at indices
        '''
        self.dataset = self.dataset.isel(kwargs)
        
    def subset_as_copy(self, **kwargs):
        '''
        Similar to COAsT.subset() however applies the subsetting to a copy of
        the original COAsT object. This subsetted copy is then returned.
        Useful for preserving the original object whilst creating smaller
        subsetted object copies.
        '''
        obj_copy = self.copy()
        obj_copy.subset(**kwargs)
        return obj_copy

    def distance_between_two_points(self):
        raise NotImplementedError
        
    def subset_indices_by_distance(self, centre_lon: float, centre_lat: float, 
                                   radius: float):
        """
        This method returns a `tuple` of indices within the `radius` of the lon/lat point given by the user.

        Distance is calculated as haversine - see `self.calculate_haversine_distance`

        :param centre_lon: The longitude of the users central point
        :param centre_lat: The latitude of the users central point
        :param radius: The haversine distance (in km) from the central point
        :return: All indices in a `tuple` with the haversine distance of the central point
        """

        # Flatten NEMO domain stuff.
        lon = self.dataset.longitude
        lat = self.dataset.latitude

        # Calculate the distances between every model point and the specified
        # centre. Calls another routine dist_haversine.

        dist = self.calculate_haversine_distance(centre_lon, centre_lat, 
                                                 lon, lat)
        indices_bool = dist < radius
        indices = np.where(indices_bool.compute())

        return xr.DataArray(indices[0]), xr.DataArray(indices[1])
    
    def subset_indices_lonlat_box(self, lonbounds, latbounds):
        """Generates array indices for data which lies in a given lon/lat box.

        Keyword arguments:
        lon       -- Longitudes, 1D or 2D.
        lat       -- Latitudes, 1D or 2D
        lonbounds -- Array of form [min_longitude=-180, max_longitude=180]
        latbounds -- Array of form [min_latitude, max_latitude]
        
        return: Indices corresponding to datapoints inside specified box
        """
        lon_str = 'longitude'
        lat_str = 'latitude'
        lon = self.dataset[lon_str].copy()
        lat = self.dataset[lat_str]
        ff = ( lon > lonbounds[0] ).astype(int)
        ff = ff*( lon < lonbounds[1] ).astype(int)
        ff = ff*( lat > latbounds[0] ).astype(int)
        ff = ff*( lat < latbounds[1] ).astype(int)
        return np.where(ff)

    def calculate_haversine_distance(self, lon1, lat1, lon2, lat2):
        '''
        # Estimation of geographical distance using the Haversine function.
        # Input can be single values or 1D arrays of locations. This
        # does NOT create a distance matrix but outputs another 1D array.
        # This works for either location vectors of equal length OR a single loc
        # and an arbitrary length location vector.
        #
        # lon1, lat1 :: Location(s) 1.
        # lon2, lat2 :: Location(s) 2.
        '''

        # Convert to radians for calculations
        lon1 = xr.ufuncs.deg2rad(lon1)
        lat1 = xr.ufuncs.deg2rad(lat1)
        lon2 = xr.ufuncs.deg2rad(lon2)
        lat2 = xr.ufuncs.deg2rad(lat2)

        # Latitude and longitude differences
        dlat = (lat2 - lat1) / 2
        dlon = (lon2 - lon1) / 2

        # Haversine function.
        distance = xr.ufuncs.sin(dlat) ** 2 + xr.ufuncs.cos(lat1) * xr.ufuncs.cos(lat2) * xr.ufuncs.sin(dlon) ** 2
        distance = 2 * 6371.007176 * xr.ufuncs.arcsin(xr.ufuncs.sqrt(distance))

        return distance

    def get_subset_as_xarray(self, var: str, points_x: slice, points_y: slice, line_length: int = None,
                             time_counter: int = 0):
        """
        This method gets a subset of the data across the x/y indices given for the chosen variable.

        Setting time_counter to None will treat `var` as only having 3 dimensions depth, y, x

        there is a check on `var` to see the size of the time_counter, if 1 then time_counter is fixed to index 0.

        :param var: the name of the variable to get data from
        :param points_x: a list/array of indices for the x dimension
        :param points_y: a list/array of indices for the y dimension
        :param line_length: (Optional) the length of your subset (assuming simple line transect)
        :param time_counter: (Optional) which time slice to get data from, if None and the variable only has one a time
                             channel of length 1 then time_counter is fixed too an index of 0
        :return: data across all depths for the chosen variable along the given indices
        """

        try:
            [time_size, _, _, _] = self.dataset[var].shape
            if time_size == 1:
                time_counter == 0

        except ValueError:
            time_counter = None

        dx = xr.DataArray(points_x)
        dy = xr.DataArray(points_y)

        if time_counter is None:
            smaller = self.dataset[var].isel(x_dim=dx, y_dim=dy)
        else:
            smaller = self.dataset[var].isel(t_dim=time_counter, x_dim=dx, y_dim=dy)

        return smaller

    def get_2d_subset_as_xarray(self, var: str, points_x: slice, points_y: slice, line_length: int = None,
                                time_counter: int = 0):
        """

        :param var:
        :param points_x:
        :param points_y:
        :param line_length:
        :param time_counter:
        :return:
        """

        try:
            [time_size, _, _, _] = self.dataset[var].shape
            if time_size == 1:
                time_counter == 0
        except ValueError:
            time_counter = None

        if time_counter is None:
            smaller = self.dataset[var].isel(x=points_x, y=points_y)
        else:
            smaller = self.dataset[var].isel(time_counter=time_counter, x=points_x, y=points_y)

        return smaller

    def plot_simple_2d(self, x, y, data: xr.DataArray, cmap, plot_info: dict):
        """
        This is a simple method that will plot data in a 2d. It is a wrapper for matplotlib's 'pcolormesh' method.

        `cmap` and `plot_info` are required to run this method, `cmap` is passed directly to `pcolormesh`.

        `plot_info` contains all the required information for setting the figure;
         - ylim
         - xlim
         - clim
         - title
         - fig_size
         - ylabel

        :param x: The variable contain the x axis information
        :param y: The variable contain the y axis information
        :param data: the DataArray a user wishes to plot
        :param cmap:
        :param plot_info:
        :return:
        """
        import matplotlib.pyplot as plt

        plt.close('all')

        fig = plt.figure()
        plt.rcParams['figure.figsize'] = plot_info['fig_size']

        ax = fig.add_subplot(411)
        plt.pcolormesh(x, y, data, cmap=cmap)

        plt.ylim(plot_info['ylim'])
        plt.xlim(plot_info['xlim'])
        plt.title(plot_info['title'])
        plt.ylabel(plot_info['ylabel'])
        plt.clim(plot_info['clim'])
        plt.colorbar()

        return plt

    def plot_cartopy(self, var: str, plot_var: array, params, time_counter: int = 0):
        try:
            import cartopy.crs as ccrs  # mapping plots
            import cartopy.feature  # add rivers, regional boundaries etc
            from cartopy.mpl.gridliner import LONGITUDE_FORMATTER, LATITUDE_FORMATTER  # deg symb
            from cartopy.feature import NaturalEarthFeature  # fine resolution coastline
        except ImportError:
            import sys
            warn("No cartopy found - please run\nconda install -c conda-forge cartopy")
            sys.exit(-1)

        import matplotlib.pyplot as plt
        plt.close('all')
        fig = plt.figure(figsize=(10, 10))
        ax = fig.gca()
        ax = plt.subplot(1, 1, 1, projection=ccrs.PlateCarree())

        cset = self.dataset[var].isel(time_counter=time_counter, deptht=0).plot.pcolormesh \
            (np.ma.masked_where(plot_var == np.NaN, plot_var), transform=ccrs.PlateCarree(), cmap=params.cmap)

        cset.set_clim([params.levs[0], params.levs[-1]])

        ax.add_feature(cartopy.feature.OCEAN)
        ax.add_feature(cartopy.feature.BORDERS, linestyle=':')
        ax.add_feature(cartopy.feature.RIVERS)
        coast = NaturalEarthFeature(category='physical', scale='10m', facecolor='none', name='coastline')
        ax.add_feature(coast, edgecolor='gray')

        gl = ax.gridlines(crs=ccrs.PlateCarree(), draw_labels=True,
                          linewidth=0.5, color='gray', alpha=0.5, linestyle='-')

        gl.xlabels_top = False
        gl.xlabels_bottom = True
        gl.ylabels_right = False
        gl.ylabels_left = True
        gl.xformatter = LONGITUDE_FORMATTER
        gl.yformatter = LATITUDE_FORMATTER

        plt.colorbar(cset, shrink=params.colorbar_shrink, pad=.05)

        # tmp = self.dataset.votemper
        # tmp.attrs = self.dataset.votemper.attrs
        # tmp.isel(time_counter=time_counter, deptht=0).plot.contourf(ax=ax, transform=ccrs.PlateCarree())
        # ax.set_global()
        # ax.coastlines()
        plt.show()

    def plot_movie(self):
        raise NotImplementedError
