
# coding: utf-8

# In[1]:

# Functions


# In[ ]:
import netCDF4    # for dealing with the datasets
import matplotlib.pyplot as plt # for plotting the data
from mpl_toolkits.basemap import Basemap, cm # for showing the data over a map
import numpy as np 
import pylab as plb
import matplotlib as mpl
import seaborn as svn

def plot_nc(fig, data, latvals, lonvals, zord, alp, colmap, label, title, axname):
    
# Create Equidistant Cylindrical Projection 'cyl' Basemap instance specifying
# the coordinates of the corners.
    m = Basemap(projection='cyl',llcrnrlat=latvals[0],            urcrnrlat=latvals[-1], llcrnrlon=lonvals[0],urcrnrlon=lonvals[-1],resolution='c')

# Use makegrid to get lon,lat coordinates of an equally spaced native projection grid
# based on the specified corners and then project them on tehe map in use
    ny = latvals.shape[0]; nx = lonvals.shape[0] # Dimensions
    lons, lats = m.makegrid(nx, ny) # get lat/lons of ny by nx evenly space grid.
    x, y = m(lons, lats) # compute map projected coordinates.

# Draw filled contours.

    cs = m.contourf(x, y, data, cmap=colmap, zorder=zord, alpha=alp) # item with higher zorder stays on top

# Draw coastlines and country boundaries, edge of map.
    m.drawcoastlines()

# Draw parallels.
    parallels = np.arange(-90.,91.,30.) # specify parallels stepsize
    m.drawparallels(parallels,labels=[1,0,0,0],fontsize=10)

# Draw meridians
    meridians = np.arange(-180.,181.,30.) # specify meridians stepsize
    m.drawmeridians(meridians,labels=[0,0,0,1],fontsize=10)

# Add colorbar
    cbar = fig.colorbar(cs,orientation='horizontal',pad=0.05, shrink = 0.8, aspect = 22,                     anchor = (0.3, 1.0), use_gridspec=False)
    cbar.set_label(label)


# Add title
    axname.set_title(title)
    return;


def plot_north_polar(fig, data, lon, lat, blat = 45, lon0 = 60, res = 'l', clevs = np.arange(0), cmap =plt.cm.RdBu_r, par = np.arange(-90.,91.,20.), mer = np.arange(-180.,181.,20.), cbar=True):
    m = Basemap(projection='npaeqd',boundinglat=blat,lon_0=lon0,resolution='l')
    m.drawcoastlines()
    #m.fillcontinents(color='coral',lake_color='aqua')
    # draw parallels and meridians.
    m.drawparallels(par)
    m.drawmeridians(mer)
    m.drawmapboundary(fill_color='white')
    m.drawcoastlines(linewidth=1.5)
    lons,lats = np.meshgrid(lon,lat)
    x, y = m(lons,lats)
    if (len(clevs)>0): 
        cs = m.contourf(x,y,data,clevs=clevs,cmap=cmap)
    else:
        cs = m.contourf(x,y,data,cmap=cmap)
    
    # cs2 = m.contour(x,y,data,,linewidths=0.5,colors='k',animated=True)
    if (cbar==True): cb = m.colorbar(cs,"right", size="5%", pad="4%")
    return m;


def reverse_arr_k(arr,k):
    # reverse dimension k of an n-dimensional array
    reversed = np.swapaxes(np.swapaxes(arr, 0, k)[::-1], 0, k)
    return reversed



def write_asc2nc(ifile, ofile, varname, reversed, description, units, longname):
    # Read ASCII GRID ESRI Dataset and Convert it to NETCDF4 dataset with rectangular grid
    # load the dataset 
    data  = np.loadtxt(ifile, skiprows=6)
    root_grp = netCDF4.Dataset(ofile, 'w', format='NETCDF4', clobber=True)
    root_grp.description = description # string description of content of dataset
    
    # dimensions
    # root_grp.createDimension('time', None)
    ny = data.shape[:][0]
    nx = data.shape[:][1]
    root_grp.createDimension('lat', ny)
    root_grp.createDimension('lon', nx)
    
    # variables
    # times = root_grp.createVariable('time', 'f8', ('time',))
    longitudes = root_grp.createVariable('lon', 'f4', ('lon',))
    latitudes = root_grp.createVariable('lat', 'f4', ('lat',))
    var = root_grp.createVariable(varname, 'f4', ('lat', 'lon',))
    
    # data
    cy = 180./ny
    cx = 360./nx
    lats =  np.arange(-90, -90+ny*cy, cy)
    lons =  np.arange(-180, -180+nx*cx, cx)
    latitudes[:] = lats
    longitudes[:] = lons
    longitudes.units = 'degrees_east'
    latitudes.units = 'degrees_north'
    longitudes.long_name = 'Longitudes'
    latitudes.long_name = 'Latitudes'
    if reversed==True: # in case dataset is bottom up, reverse longitudes
        my_data = np.reshape(data[::-1],(ny,nx),'F')
    else:
        my_data = np.reshape(data,(ny,nx),'F')
    var[:] = my_data
    var.units = units
    var.long_name = longname
    root_grp.close()
