import os
from PIL import Image
import cv2
from os import listdir
from os.path import join

import matplotlib.pyplot as plt
import matplotlib
from matplotlib.colors import LogNorm
from io_utils.io_common import create_folder
from viz_utils.constants import PlotMode, BackgroundType
import pylab
import numpy as np
import cmocean
import shapely

import cartopy.crs as ccrs
import cartopy.feature as cfeature
import cartopy

def select_colormap(field_name):
    '''
    Based on the name if the field it chooses a colormap from cmocean
    Args:
        field_name:

    Returns:

    '''
    if np.any([field_name.find(x) != -1 for x in ('ssh', 'srfhgt', 'adt','surf_el')]):
        # cmaps_fields.append(cmocean.cm.deep_r)
        return cmocean.cm.curl
    elif np.any([field_name.find(x) != -1 for x in ('temp', 'sst', 'temperature')]):
        return cmocean.cm.thermal
    elif np.any([field_name.find(x) != -1 for x in ('vorticity', 'vort')]):
        return cmocean.cm.curl
    elif np.any([field_name.find(x) != -1 for x in ('salin', 'sss', 'sal')]):
        return cmocean.cm.haline
    elif field_name.find('error') != -1:
        return cmocean.cm.diff
    elif field_name.find('binary') != -1:
        return cmocean.cm.oxy
    elif np.any([field_name.find(x) != -1 for x in ('u_', 'v_', 'u-vel.', 'v-vel.','velocity')]):
        return cmocean.cm.speed


class EOAImageVisualizer:
    """This class makes plenty of plots assuming we are plotting Geospatial data (maps).
    It is made to read xarrays, numpy arrays, and numpy arrays in dictionaries
    vizobj = new EOAImageVisualizer(disp_images=True, output_folder='output',
                                    lats=[lats],lons=[lons])
    """
    _COLORS = ['y', 'r', 'c', 'b', 'g', 'w', 'k', 'y', 'r', 'c', 'b', 'g', 'w', 'k']
    _figsize = 8
    _font_size = 30
    _units = ''
    _max_imgs_per_row = 4
    _mincbar = np.nan  # User can set a min and max colorbar values to 'force' same color bar to all plots
    _maxcbar  = np.nan
    _flip_data = True
    _eoas_pyutils_path =  './eoas_pyutils'# This is the path where the eoas_utils folder is stored with respect to the main project
    _contourf = False  # When plotting non-regular grids and need precision
    _background = BackgroundType.BLUE_MARBLE_LR  # Select the background to use
    _auto_colormap = True  # Selects the colormap based on the name of the field
    _show_var_names = False  # Includes the name of the field name in the titles
    _additional_polygons = []  # MUST BE SHAPELY GEOMETRIES In case we want to include additional polygons in the plots (all of them)
    # If you want to add a streamplot of a vector field. It must be a dictionary with keys x,y,u,v
    # and optional density, color, cmap, arrowsize, arrowstyle, minlength
    _vector_field = None
    _norm = None  # Use to normalize the colormap. For example with LogNorm

    # vizobj = EOAImageVisualizer(disp_images=True, output_folder='output',
    #                                 lats=[lats],lons=[lons])
    def __init__(self, disp_images=True, output_folder='output',
                 lats=[-90,90], lons =[-180,180],
                 projection=ccrs.PlateCarree(), **kwargs):
        # All the arguments that are passed to the constructor of the class MUST have its name on it.
        self._disp_images = disp_images
        self._output_folder = output_folder
        self._projection = projection
        bbox = self.getExtent(lats, lons)
        self._extent = bbox
        self._lats = lats
        self._lons = lons
        self._fig_prop = (bbox[1]-bbox[0])/(bbox[3]-bbox[2])
        self._contour_labels = False
        for arg_name, arg_value in kwargs.items():
            self.__dict__["_" + arg_name] = arg_value
            print(self.__dict__["_" + arg_name])

    def __getattr__(self, attr):
        '''Generic getter for all the properties of the class'''
        return self.__dict__["_" + attr]

    def __setattr__(self, attr, value):
        '''Generic setter for all the properties of the class'''
        self.__dict__["_" + attr] = value

    def add_colorbar(self, fig, im, ax, show_color_bar, label=""):
        # https://matplotlib.org/api/_as_gen/matplotlib.pyplot.colorbar.html
        if show_color_bar:
            font_size_cbar = self._font_size * .5
            # TODO how to make this automatic and works always
            cbar = fig.colorbar(im, ax=ax, shrink=.7)
            cbar.ax.tick_params(labelsize=font_size_cbar)
            if label != "":
                cbar.set_label(label, fontsize=font_size_cbar*1.2)
            else:
                cbar.set_label(self._units, fontsize=font_size_cbar*1.2)

    def plot_slice_eoa(self, c_img, ax, cmap='gray', mode=PlotMode.RASTER, mincbar=np.nan, maxcbar=np.nan) -> None:
        """
        Plots a 2D img for EOA data.
        :param c_img: 2D array
        :param ax: geoaxes
        :return:
        """
        c_ax = ax
        if self._flip_data:
            origin = 'lower'
        else:
            origin = 'upper'

        if self._background == BackgroundType.CARTO_DEF:
            c_ax.stock_img()
        else:
            if self._background == BackgroundType.BLUE_MARBLE_LR:
                img = plt.imread(join(self._eoas_pyutils_path,'viz_utils/imgs/bluemarble.png'))
            if self._background == BackgroundType.BLUE_MARBLE_HR:
                img = plt.imread(join(self._eoas_pyutils_path,'viz_utils/imgs/bluemarble_5400x2700.jpg'))
            if self._background == BackgroundType.TOPO:
                img = plt.imread(join(self._eoas_pyutils_path,'viz_utils/imgs/etopo.png'))
            if self._background == BackgroundType.BATHYMETRY:
                img = plt.imread(join(self._eoas_pyutils_path,'viz_utils/imgs/bathymetry_3600x1800.jpg'))
            c_ax.imshow(img, origin='upper', extent=(-180,180,-90,90), transform=ccrs.PlateCarree())

        if mode == PlotMode.RASTER or mode == PlotMode.MERGED:
            if self._contourf:
                im = c_ax.contourf(self._lons, self._lats, c_img, num_colors=255, cmap='inferno', extent=self._extent)
            else:
                if 
np.isnan(mincbar)