import numpy as np, pickle, os, re, sys, subprocess, math, json, shutil
import netCDF4, glob, csv, arcpy, requests, logging, traceback
from osgeo import gdal, osr, ogr, gdal_array
#from datetime import datetime, timedelta
import datetime
from time import gmtime, strftime
import xarray as xr
import pandas as pd
from os import path

try:
    if arcpy.CheckExtension("Spatial") == "Available":
        arcpy.CheckOutExtension("Spatial")
    else:
        # Raise a custom exception
        #
        raise LicenseError

except LicenseError:
    print ("Spatial Analyst license is unavailable")  

wsa_data = "C:\\Users\\joshua.kastman\\gis\\wsa\\data\\"
wsa_scripts = "C:\\Users\\joshua.kastman\\gis\\wsa\\scripts\\"
wsa_proj = "C:\\Users\\joshua.kastman\\gis\\wsa\\projects\\"

arcpy.env.overwriteOutput = True

srs = osr.SpatialReference()
srs.ImportFromProj4("+proj=lcc +units=m +lat_0=25 +lon_0=-95 +lat_1=25 +lat_2=25 +a=6371000 +b=6371000 +ellps=sphere")

def clean(clean_path):
    print ("Cleaning out old files")
    try:
        for root, dirs, files in os.walk(clean_path):
            for f in files:
                os.unlink(os.path.join(root, f))
            for d in dirs:
                shutil.rmtree(os.path.join(root, d))
    except Exception as e:
        print(traceback.format_exc())

def validTXT(date):
    vt = datetime.datetime.strptime(date, '%Y%m%d%H')
    vtD4 = vt + datetime.timedelta(days = 4)
    valid = vt.strftime("%HZ %m/%d/%Y")
    validD4 = vtD4.strftime("%HZ %m/%d/%Y")
    print ("Valid "+valid+" - "+validD4)
    out=open(wsa_data+"valid.txt", 'w')
    out.write("Valid "+valid+" - "+validD4)
    out.close()

def createGeoTiffRaster(data,filename,geoTransform,noData,dataType,srs):
    
    colors = readColorFile()

    #Open GeoTiff driver
    driver = gdal.GetDriverByName("GTiff")

    #Create our image
    if len(np.shape(data)) > 2:
        zSize,ySize,xSize = np.shape(data)
        data = data[1:]
    else:
        ySize,xSize = np.shape(data)
    dsOut = driver.Create(filename, xSize, ySize, 1, gdal.GDT_Byte)
    dsOut.SetProjection(srs.ExportToWkt())
    dsOut.SetGeoTransform(geoTransform)

    #Set the noData value on the raster
    dsOut.GetRasterBand(1).SetNoDataValue(noData)
    bandOut=dsOut.GetRasterBand(1)
    bandOut.SetRasterColorTable(colors)

    bandOut.SetRasterColorInterpretation(gdal.GCI_PaletteIndex)   
    gdal_array.BandWriteArray(bandOut, data)

    bandOut = None
    dsOut.FlushCache()
    dsOut = None


def warpGeoTiff(infile):

    tmpfilename = infile.replace("LCC","webMerc") 
    arcpy.env.overwriteOutput = True
    arcpy.management.ProjectRaster(in_raster=infile, out_raster=tmpfilename, out_coor_system="PROJCS[\"WGS_1984_Web_Mercator_Auxiliary_Sphere\",GEOGCS[\"GCS_WGS_1984\",DATUM[\"D_WGS_1984\",SPHEROID[\"WGS_1984\",6378137.0,298.257223563]],PRIMEM[\"Greenwich\",0.0],UNIT[\"Degree\",0.0174532925199433]],PROJECTION[\"Mercator_Auxiliary_Sphere\"],PARAMETER[\"False_Easting\",0.0],PARAMETER[\"False_Northing\",0.0],PARAMETER[\"Central_Meridian\",0.0],PARAMETER[\"Standard_Parallel_1\",0.0],PARAMETER[\"Auxiliary_Sphere_Type\",0.0],UNIT[\"Meter\",1.0]]")
    return tmpfilename

def createGeotransform(lat, lon):

    #Find number of columns and rows in raster
    rows,cols = np.shape(lat)
    print ('In CreateGeotransform')
    
    llLat = float(lat[-1][0])
    llLon = float(lon[-1][0])
    urLat = float(lat[0][-1])
    urLon = float(lon[0][-1])

    print("[[{},{}],[{},{}]]".format(llLat,llLon,urLat,urLon))
    #Convert the lat/lons to meters on the projection
    urX,urY = projPoints(urLon,urLat)
    llX,llY = projPoints(llLon,llLat)

    ulLat = float(lat[0][0])
    ulLon = float(lon[0][0])
    lrLat = float(lat[0][0])
    lrLon = float(lon[0][0])

    #Convert the lat/lons to meters on the projection

    ulX,ulY = projPoints(ulLon,ulLat)
    lrX,lrY = projPoints(lrLon,lrLat)

    #Find cell column and row height based on the width/height of the grid and the difference in meters
    colWidth = math.fabs((ulX - urX)/cols)
    rowHeight = math.fabs((ulY - llY)/rows)
    
    print ('Raster size in meters:',rowHeight, colWidth)
    #Create the array of the location of the raster on the projection (meters)
    geoTransform = [None] * 6
    geoTransform[0] = ulX # top left x (note that the netCDF files are upside down, so I used the lowerleft instead of using the rotation variable [3])
    geoTransform[1] = colWidth #w-e pixel resolution
    geoTransform[2] = 0 #rotation, 0 if image is "north up"
    geoTransform[3] = ulY #top left y
    geoTransform[4] = 0 #rotation, 0 if image is "north up"
    geoTransform[5] = -rowHeight #n-s pixel resolution
    return geoTransform

def projPoints(lon,lat):

  srs = osr.SpatialReference()
  srs.ImportFromProj4("+proj=lcc +units=m +lat_0=25 +lon_0=-95 +lat_1=25 +lat_2=25 +a=6371000 +b=6371000 +ellps=sphere")

  target = osr.SpatialReference()
  target.ImportFromEPSG(4326)

  transform = osr.CoordinateTransformation(target, srs)
  p = ogr.CreateGeometryFromWkt("POINT ("+str(lon)+" "+str(lat)+")")
  p.Transform(transform)
  point = p.ExportToWkt()
  print ('transformed point {}'.format(point))
  logging.debug('transformed point {}'.format(point))
  plst = point.split()
  x = float(plst[1].split('(')[1])
  y = float(plst[2].split(')')[0])
  print ('projected x value: {}'.format(x))
  print ('projected y value: {}'.format(y))
  
  return x,y


def readColorFile():
    '''
    Returns a gdal colortable using the HeatRisk_Color.txt values.
    '''
    fn = "{}HeatRisk_Color.txt".format(wsa_proj)

    colors = gdal.ColorTable()

    with open(fn) as f:
        lines = f.readlines()
        f.close()
    for line in lines:
        tmp = line.split()
        colors.SetColorEntry(int(tmp[0]),(int(tmp[1]),int(tmp[2]),int(tmp[3])))

    return colors

def make_webpngs(file):

    aprx = arcpy.mp.ArcGISProject(wsa_proj+"wsa_web_png.aprx")
    m = aprx.listMaps("WSA_web")[0]
    l = m.listLayers()[0]
    mv = m.defaultView
    layers = m.listLayers()
    web_png = wsa_data+"wsa_latest.png"
    print (web_png)
    mv.exportToPNG(web_png, width=5000, height=3719, resolution=600, world_file=False, color_mode="32-BIT_WITH_ALPHA")

def send2web(ldir, ext):
    key = "C:\\Users\\joshua.kastman\\gis\\keys\\kastman_c_nc_vs_wwddev_private.ppk"
    linux_user = "hpc@vm-cprk-rzdm01"
    wbdir = "//home//people//hpc//www//htdocs//wwd//internal//wsa//images//"
    cmd = "pscp -i {} {}{} {}:{}".format(key,ldir,ext,linux_user,wbdir)
    print (cmd)
    os.system(cmd)
 
def main():
    # clean dirs
    clean(wsa_data)
    subprocess.call(wsa_scripts+"\\dl_wsa.bat")
    wsa_nc = glob.glob(os.path.join(wsa_data, "*WSA*"+'*.nc'))[-1]
    wsa_fn = wsa_nc.split("\\")[-1].split(".")[0]
    date = wsa_fn.split("_")[-1]
    # Open netcdf file in Xarray
    ds_wsa = xr.open_dataset(wsa_nc)
    # get array values for WSA
    wsa_vals = ds_wsa.WPCWSA_SFC.values
    # remove meaningless time dimension 
    wsa_2d = np.squeeze(wsa_vals)
    # flip cords and adjust lons for degrees west
    wsa_vals_flip = np.flipud(wsa_2d)
    lats = np.flipud(ds_wsa['latitude'].values)
    lons = np.flipud(ds_wsa['longitude'].values - 360)
    # transform image cords to georeferenced cords
    geoTransform = createGeotransform(lats,lons)
    noDataValue = 9999
    tif_name = wsa_data+"wsa_latest_LCC.tif"
    # create tif
    createGeoTiffRaster(wsa_vals_flip,tif_name,geoTransform,noDataValue,1,srs)
    mercTif = warpGeoTiff(tif_name)
    make_webpngs(mercTif)
    validTXT(date)
    ds_wsa.close()
    send2web(wsa_data, "*png")
    send2web(wsa_data, "*txt")

if __name__ == '__main__':
    main()