'''
Author: Alex Crawford
Date Created: 2017
Date Modified: 11 May 2018

Purpose: To extract the spatial dimensions in a ERA netCDF file, and then
create a raster of the lat and long values for ERA in the EASE2 projection.
'''
'''*******************************************
Set up Modules
*******************************************'''
from osgeo import gdal, gdalconst, gdalnumeric
from mpl_toolkits.basemap import Basemap
import numpy as np
import netCDF4 as nc
import os

'''############
Load Variables
############'''
# Projection info
projP = '''PROJCS["EASE",GEOGCS["WGS 84",DATUM["WGS_1984",SPHEROID["WGS 84",
6378137,298.257223563,AUTHORITY["EPSG","7030"]],AUTHORITY["EPSG","6326"]],
PRIMEM["Greenwich",0],UNIT["degree",0.0174532925199433],AUTHORITY["EPSG","4326"]],
PROJECTION["Lambert_Azimuthal_Equal_Area"],PARAMETER["latitude_of_center",90],
PARAMETER["longitude_of_center",0],PARAMETER["false_easting",0],
PARAMETER["false_northing",0],UNIT["metre",1,AUTHORITY["EPSG","9001"]]]'''

# Inputs for reprojection
bb = [-84.614,-45,-84.614,135] # in degrees
ulpnt = (-9039270.143837992,9039270.143837988) # in meters (yields a 181 by 181 grid)
xsize, ysize = 100000, -100000 # in meters
nx, ny = 181, 181 # number of grid cells
lon_0 = 0 # Central Meridian (which longitude is at the 6 o'clock position)
lat_0 = 90 # Reference Latitude (center of projection)

ext = '.nc'
dtype = gdal.GDT_Float32

# Path Variables
ra = 'ERAI'
invar = 'SLP'
inpath = '/Volumes/Ferdinand/'+ra+'_nc/'+invar
outpath = '/Volumes/Ferdinand/Projections'

outNameLat = outpath+"/EASE2_N0_100km_Lats.tif"
outNameLon = outpath+"/EASE2_N0_100km_Lons.tif"

'''################
Main Analysis
################'''
# Load File
nclist = os.listdir(inpath)
nclist = [n for n in nclist if n.startswith('E') and n.endswith(ext)]
ncf = nc.Dataset(inpath+'/'+nclist[0])

# Load variables
lons = ncf.variables['longitude'][:]
lats = ncf.variables['latitude'][:]

#Shift to be -180 to 180 and -90 to 90
nLg2 = len(lons)/2
inLonsA = np.where(lons >= 180, lons - 360, lons)
inLons = np.zeros_like(lons)
inLons[:nLg2], inLons[nLg2:] = inLonsA[nLg2:], inLonsA[:nLg2]

inLats = np.flipud(lats)

# Create 2-D Array Inputs
inLats2 = np.repeat(inLats,len(lons)).reshape(len(lats),len(lons))
inLons2 = np.rot90(np.repeat(inLons,len(lats)).reshape(len(lons),len(lats)))

### Reprojection of Latitudes
# Set up the map for reprojecting
mp = Basemap(projection='laea',lat_0=lat_0,lon_0=lon_0,\
llcrnrlat=bb[0], llcrnrlon=bb[1],urcrnrlat=bb[2], urcrnrlon=bb[3],resolution='c')

# Transform data
outArr, xs, ys = mp.transform_scalar(inLats2,inLons,inLats,nx,ny,returnxy=True)

# Write to file
driver = gdal.GetDriverByName('GTiff')
outFile = driver.Create(outNameLat,outArr.shape[1],outArr.shape[0],1,dtype) # Create file
outFile.GetRasterBand(1).WriteArray(np.flipud(outArr),0,0) # Write array to file
outFile.GetRasterBand(1).ComputeStatistics(False) # Compute stats for display purposes
outFile.SetGeoTransform((ulpnt[0],xsize,0,ulpnt[1],0,ysize)) # Set geotransform (those six needed values)
outFile.SetProjection(projP)  # Set projection
outFile = None

### Reprojection of Longitudes
# Set up the map for reprojecting
mp = Basemap(projection='laea',lat_0=lat_0,lon_0=lon_0,\
llcrnrlat=bb[0], llcrnrlon=bb[1],urcrnrlat=bb[2], urcrnrlon=bb[3],resolution='c')

# Transform data
outArr, xs, ys = mp.transform_scalar(inLons2,inLons,inLats,nx,ny,returnxy=True)

# Write to file
driver = gdal.GetDriverByName('GTiff')
outFile = driver.Create(outNameLon,outArr.shape[1],outArr.shape[0],1,dtype) # Create file
outFile.GetRasterBand(1).WriteArray(np.flipud(outArr),0,0) # Write array to file
outFile.GetRasterBand(1).ComputeStatistics(False) # Compute stats for display purposes
outFile.SetGeoTransform((ulpnt[0],xsize,0,ulpnt[1],0,ysize)) # Set geotransform (those six needed values)
outFile.SetProjection(projP)  # Set projection
outFile = None