import numpy as np
import matplotlib as mpl
mpl.use('pdf')
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
from rasterio.crs import CRS
import cartopy.feature as cfeature
import matplotlib.ticker as mticker
import rasterio as rio
from rasterio.transform import Affine
import shapefile
from shapely.geometry import shape
from shapely.ops import transform
import pyproj

def map_lst(data, output_fn, maptitle):

    resolution = 0.02  # 1 arcsecond = 1/3600 degree
    width = 2912
    height = 1250
    invalid = -9999.
    valid_range = [150.0, 1310.7]
    extent = [-124.99, -66.75, 24.86, 49.83]
    region_extent = [-124.99, -66.75, 24.86, 49.83]
    dpiv = 400
    minvalue = 250
    maxvalue = 330

    # Set up the figure and axis
    fig = plt.figure(figsize=(20, 10))
    ax = fig.add_subplot(1, 1, 1, projection=ccrs.PlateCarree())
    ax.set_extent(region_extent, crs=ccrs.PlateCarree())

    #cmap = mpl.cm.jet
    cmap = mpl.cm.bwr
    #cmap = mpl.cm.RdYlGn
    #cmap_r = reverse_colourmap(cmap)
    cmap_r = cmap 

    # Plot the map
    img = ax.imshow(data, extent=extent, transform=ccrs.PlateCarree(), origin='upper', cmap=cmap_r, vmin=minvalue, vmax=maxvalue)

    # Add gridlines
    ax.gridlines(linewidth=0.5, color='gray', linestyle='--')  
 
    #Add country boundaries
    ax.add_feature(cfeature.BORDERS, linewidth=0.8, edgecolor='black')

    # Add coastlines
    ax.coastlines(resolution='50m', linewidth=1.8, color='black')

    # Add longitude and latitude labels
    lonstart = region_extent[0]
    lonend = region_extent[1]
    loninter= (region_extent[1]-region_extent[0])/8
    latstart = region_extent[2]
    latend = region_extent[3]
    loninter = (region_extent[3]-region_extent[2])/8
    ax.set_xticks(np.arange(lonstart, lonend, loninter), crs=ccrs.PlateCarree())
    ax.set_yticks(np.arange(latstart, latend, loninter), crs=ccrs.PlateCarree())
    lon_formatter = mticker.StrMethodFormatter("{x:.2f}\u00b0")
    lat_formatter = mticker.StrMethodFormatter("{x:.2f}\u00b0")
    ax.xaxis.set_major_formatter(lon_formatter)
    ax.yaxis.set_major_formatter(lat_formatter)

    # Add a colorbar
    cbar = plt.colorbar(img, ax=ax, orientation='horizontal', pad=0.05, fraction=0.04)

    # Add map title
    print('maptitle = ' + maptitle)
    ax.set_title(maptitle, fontsize=20, fontdict={'fontweight': 'bold'})

    plt.savefig(output_fn, dpi=dpiv, bbox_inches='tight', pad_inches=0.2)
  
    return

def reverse_colourmap(cmap, name = 'my_cmap_r'):
    """
    In: 
    cmap, name 
    Out:
    my_cmap_r

    Explanation:
    t[0] goes from 0 to 1
    row i:   x  y0  y1 -> t[0] t[1] t[2]
                   /
                  /
    row i+1: x  y0  y1 -> t[n] t[1] t[2]

    so the inverse should do the same:
    row i+1: x  y1  y0 -> 1-t[0] t[2] t[1]
                   /
                  /
    row i:   x  y1  y0 -> 1-t[n] t[2] t[1]
    """        
    reverse = []
    k = []   

    for key in cmap._segmentdata:    
        k.append(key)
        channel = cmap._segmentdata[key]
        data = []

        for t in channel:                    
            data.append((1-t[0],t[2],t[1]))            
        reverse.append(sorted(data))    

    LinearL = dict(zip(k,reverse))
    my_cmap_r = mpl.colors.LinearSegmentedColormap(name, LinearL) 
    return my_cmap_r