
import xarray as xr
import numpy as np
import matplotlib.pyplot as plt
import cartopy.crs as ccrs

def _plot_significance(da, ax, **kwargs):
    print(da)
    
    da.plot.contourf(
        ax=ax, transform=ccrs.PlateCarree(),
        levels=[0,0.5],
        hatches=[None, "xx"], add_colorbar=False, alpha=0
    )

def plot_array_to_map(
    da: xr.DataArray, 
    ax=None, title=None, 
    projection=None, 
    mask=None, 
    coastlines_color='black',
    contour=False, 
    contour_color = 'r',
    significance: xr.DataArray = None,
    **kwargs # used either in contourf or plot
    ):

    # if kwargs add_colorbar is true, create a new variable _add_colorbar and set kwargs['add_colorbar'] = False
    if 'add_colorbar' in kwargs:
        _add_colorbar = kwargs['add_colorbar']
        _cbar_kwargs = kwargs['cbar_kwargs'] if 'cbar_kwargs' in kwargs else {}
        kwargs['add_colorbar'] = False
        kwargs.pop('cbar_kwargs')
    else:
        _add_colorbar = False
        _cbar_kwargs = {}
    
    if mask is not None:
        da = da.where(mask)
        
        if significance is not None:
            significance = significance.where(mask)

    if projection is None:

        if ax is None:
            ax = plt.axes()

        if contour:
            #im = da.plot.contourf(ax=ax, **kwargs)
            da.plot.contour(ax=ax, colors=contour_color)

        else:
            im = da.plot(ax=ax, **kwargs)
            
        if significance is not None: _plot_significance(significance, ax)

    else:

        if ax is None:
            ax = plt.axes(projection=projection)
        if contour:
            #im = da.plot.contourf(
            #    ax=ax, transform=ccrs.PlateCarree(), **kwargs
            #    )
            im = da.plot.contour(
                ax=ax, colors=contour_color, transform=ccrs.PlateCarree()
                )
            
            if significance is not None: _plot_significance(significance, ax)

        else:
            im = da.plot(
                ax=ax,transform=ccrs.PlateCarree(), **kwargs
                )

            if significance is not None: _plot_significance(significance, ax)


        ax.coastlines(color=coastlines_color, linewidth=4)
        
    if title is not None:
        ax.set_title(title, fontsize=10, fontweight="bold")
        
    if _add_colorbar:

        cb = plt.colorbar(im, ax=ax, **_cbar_kwargs)
        cb.ax.tick_params(labelsize=30)
        # add kwargs to the colorbar

    return ax, im
