import xarray as xr
import rioxarray as rio
from matplotlib import cm
import numpy as np
import time
import matplotlib
import matplotlib.image as imgg
import os
from PIL import Image, ImageDraw, ImageFont





def write_png(data,name, origin='upper', colormapD=None):
    """
    Transform an array of data into a PNG string.
    This can be written to disk using binary I/O, or encoded using base64
    for an inline PNG like this:

    >>> png_str = write_png(array)
    >>> "data:image/png;base64,"+png_str.encode('base64')

    Inspired from
    https://stackoverflow.com/questions/902761/saving-a-numpy-array-as-an-image

    Parameters
    ----------
    data: numpy array or equivalent list-like object.
         Must be NxM (mono), NxMx3 (RGB) or NxMx4 (RGBA)

    origin : ['upper' | 'lower'], optional, default 'upper'
        Place the [0,0] index of the array in the upper left or lower left
        corner of the axes.

    colormap : callable, used only for `mono` image.
        Function of the form [x -> (r,g,b)] or [x -> (r,g,b,a)]
        for transforming a mono image into RGB.
        It must output iterables of length 3 or 4, with values between
        0. and 1.  Hint: you can use colormaps from `matplotlib.cm`.

    Returns
    -------
    PNG formatted byte string

    """

    if colormapD is None:
        def colormapD(x):
            return (x, x, x, 1)

    arr = np.atleast_3d(data)
    height, width, nblayers = arr.shape

    if nblayers not in [1, 3, 4]:
        raise ValueError('Data must be NxM (mono), '
                         'NxMx3 (RGB), or NxMx4 (RGBA)')
    assert arr.shape == (height, width, nblayers)


    if nblayers == 1:
        arr = np.array(list(map(colormapD, arr.ravel())))
        nblayers = arr.shape[1]
        if nblayers not in [3, 4]:
            raise ValueError('colormap must provide colors of r'
                             'length 3 (RGB) or 4 (RGBA)')
        arr = arr.reshape((height, width, nblayers))
    assert arr.shape == (height, width, nblayers)

    if nblayers == 3:
        arr = np.concatenate((arr, np.ones((height, width, 1))), axis=2)
        nblayers = 4
    assert arr.shape == (height, width, nblayers)
    assert nblayers == 4

    # Normalize to uint8 if it isn't already.
    if arr.dtype != 'uint8':
        with np.errstate(divide='ignore', invalid='ignore'):
            arr = arr * 255./np.array([1., 1., 1., 1.]).reshape((1, 1, 4))
            arr[~np.isfinite(arr)] = 0
        arr = arr.astype('uint8')

    # Eventually flip the image.
    if origin == 'lower':
        arr = arr[::-1, :, :]
    r3 = arr.copy(order='C')
    matplotlib.image.imsave(name, r3)


def image_to_url(image,name, colormapD=None, origin='upper'):
    """
    Infers the type of an image argument and transforms it into a URL.

    Parameters
    ----------
    image: string, file or array-like object
        * If string, it will be written directly in the output file.
        * If file, it's content will be converted as embedded in the
          output file.
        * If array-like, it will be converted to PNG base64 string and
          embedded in the output.
    origin: ['upper' | 'lower'], optional, default 'upper'
        Place the [0, 0] index of the array in the upper left or
        lower left corner of the axes.
    colormap: callable, used only for `mono` image.
        Function of the form [x -> (r,g,b)] or [x -> (r,g,b,a)]
        for transforming a mono image into RGB.
        It must output iterables of length 3 or 4, with values between
        0. and 1.  You can use colormaps from `matplotlib.cm`.

    """

    if 'ndarray' in image.__class__.__name__:
        img = write_png(image, origin=origin, colormapD=colormapD,name=name)


def get_color(x,colormap,Min,Max):
            decimals = 2
            try:
                #print(x,Min,Max)
                x=(x*1.-Min)/(Max*1.-Min)
                #print(x)
                if x < 0:
                    x=0.0
                if x > 1.0:
                    x=1.0
                x = np.around(x, decimals=decimals)
                if colormap=="Spectral" or colormap=='ocean' or colormap=="RdYlBu":
                    Tempcm=cm.get_cmap(colormap).reversed()
                else:
                    Tempcm=cm.get_cmap(colormap)
                #print(x)
                ls = np.around(np.linspace(0,1,10**decimals+1),decimals=decimals)
                #print(ls)
                if 0 <= x <= 1:
                    #print(cm.get_cmap('viridis')(ls)[np.argwhere(ls==x)][0][0].shape,cm.get_cmap('viridis')(ls)[np.argwhere(ls==x)])
                    #print(x,np.argwhere(ls==x))
                    return Tempcm(ls)[np.argwhere(ls==x)][0][0]
                else:
                    #print(np.array(np.zeros(4)).shape,np.array(np.zeros(4)))
                    #print(x)
                    return np.zeros(4)
            except:
                print(x,np.argwhere(ls==x))
                return np.zeros(4)

def ExtractMapImage(file,colormap,countyear,name,nc,geometry=""):
   
    nc = nc.rio.write_crs(4326)
    if geometry=="":
        try:
            nc = nc.rio.clip(geometries)
        except:
            pass
    else:
        pass
    bounds=[[float(nc.lat.min().values),float(nc.lon.min().values)],[float(nc.lat.max().values),float(nc.lon.max().values)]]
    i=countyear

    year=int(nc.time[i].values)
    print(i,year)
    ncVar=nc.DHW_q99
    data = ncVar[i,:,:].values
        

    Min=np.around(np.nanquantile(ncVar[int(year)-1986,:,:].values, 0.01),2)
    Max=np.around(np.nanquantile(ncVar[int(year)-1986,:,:].values, 0.99),2)#np.around(np.nanmax(ncVar[year-1985,:,:].values),2)

    image_to_url(data,name, colormapD=lambda x: get_color(x,colormap,Min,Max), origin="lower")


def ProcessAllImage(ssp,model,Colorpalete,ExportDirectory,DataDirectory):
    cc=0
    Var="DHW"
    for i in ssp:
        for j in model:
            ff=DataDirectory+"%s_%s_%s_DHW.nc"%(Var,i,j)
            for CM in Colorpalete:
                for countyear in range(115):
                    nc = xr.open_dataset(ff, decode_coords="all")
                    year=int(nc.time[countyear].values)
                    Evaluado=ExportDirectory+"%s_%s"%(CM,ff.split("/")[-1].replace(".nc","_%s.png"%(year)))
                    if not os.path.isfile(Evaluado):
                        cc=cc+1

    Total=cc
    start=time.time()
    cc=0
    for i in ssp:
        for j in model:
            ff=DataDirectory+"%s_%s_%s_DHW.nc"%(Var,i,j)
            for CM in Colorpalete:
                for countyear in range(115):
                    nc = xr.open_dataset(ff, decode_coords="all")
                    year=int(nc.time[countyear].values)
                    Evaluado=ExportDirectory+"%s_%s"%(CM,ff.split("/")[-1].replace(".nc","_%s.png"%(year)))
                    if not os.path.isfile(Evaluado):
                        ExtractMapImage(ff,CM,countyear,Evaluado,nc)
                        cc=cc+1
                        now=time.time()
                        print(cc,(start-now)/cc,(start-now)/cc*(Total-cc))

ssp=["ssp245","ssp370","ssp585"]
model=("BCC-CSM2-MR","CESM2","CanESM5","EC-Earth3","IPSL-CM6A-LR","MIROC6","MRI-ESM2-0","NorESM2-MM")#,"ensemble5","ensemble8")
Colorpalete=["RdYlBu",'Spectral','ocean',"coolwarm"]
ExportDirectory="img/"
DataDirectory="../Data/"

ProcessAllImage(ssp,model,Colorpalete,ExportDirectory,DataDirectory)




