#! /users/fearon/software/Anaconda2/bin/python
#
#  File:
#    arctic2d.py
#
#  Synopsis:
#    Use Ngl and Nio with Python to map a scalar field and vector overlay over the North Pole
#
#  Author:    
#    Matt Fearon, March 2018
#  
#
import Ngl
import Nio
import numpy as np
import os
import sys



def map_contour_vector(cvar,u,v,lat,lon,ny,nx,fhr,mname):

  
#  Send graphics to PNG

    wks_type = "ps"
    wks = Ngl.open_wks(wks_type,mname+"_"+fhr)

    cnres = Ngl.Resources()
    vcres = Ngl.Resources()
    mpres = Ngl.Resources()

    cnres.nglFrame                   = False
    cnres.nglDraw                    = False
    vcres.nglFrame                   = False
    vcres.nglDraw                    = False
    mpres.nglFrame                   = False
    mpres.nglDraw                    = False

    cnres.sfXArray                   = lon
    cnres.sfYArray                   = lat
    vcres.vfXArray                   = lon
    vcres.vfYArray                   = lat

    mpres.mpFillOn                   = True
    mpres.mpLandFillColor            = "Tan1"
    mpres.mpOceanFillColor           = "SkyBlue"
    mpres.mpInlandWaterFillColor     = "SkyBlue"
    mpres.mpGeophysicalLineThicknessF= 2.5
    mpres.mpLimitMode                = "LatLon"

    mpres.mpCenterLonF               = 340.
    mpres.mpCenterLatF               = 90.

#    mpres.mpCenterLonF               = 300.
#    mpres.mpCenterLatF               = 75.

    mpres.mpMinLatF                  = 48.
    mpres.mpMaxLatF                  = 80. 

    mpres.mpDataBaseVersion          = "MediumRes"
    mpres.mpProjection               = "Stereographic"
    mpres.tmXBMajorLengthF           = 0.005
    mpres.tmXBLabelFontHeightF       = 0.008
    mpres.tmXBLabelDeltaF            = -0.09
    mpres.pmLabelBarWidthF           = 0.018
    map_plot                         = Ngl.map(wks,mpres)

    
    cnres.cnFillOn                   = True
    cnres.cnLinesOn                  = False
    cnres.cnLineLabelsOn             = False
    cnres.lbAutoManage               = False
    cnres.pmLabelBarWidthF           = 0.018
    cnres.lbLabelFontHeightF         = 0.009
    contour_plot                     = Ngl.contour(wks,cvar,cnres) 

#    vcres.vcMinFracLengthF           = 0.013          # Increase length of
#    vcres.vcMinMagnitudeF            = 0.2        # vectors.
#    vcres.vcVectorDrawOrder          = "PostDraw"
#    vcres.vcLineArrowColor           = "Pink"
    vcres.vcGlyphStyle               = "FillArrow" 
    vcres.vcFillArrowFillColor       = "DeepPink"
    vcres.vcFillArrowEdgeThicknessF  = 0.03
    vcres.vcMinDistanceF             = 0.009
    vcres.vcRefLengthF               = 0.006
    vcres.vcRefMagnitudeF            = 10.0
    vcres.vcRefAnnoOn                = False
    vector_plot                      = Ngl.vector(wks,u,v,vcres)

    Ngl.overlay(map_plot,contour_plot)
    Ngl.overlay(map_plot,vector_plot)
    Ngl.maximize_plot(wks,map_plot)
    Ngl.draw(map_plot) 
    Ngl.frame(wks)
    Ngl.end()
        
    return


def extract_field(filename, model_grid):

    record_length = model_grid['num_bytes']

    f = open( filename, 'rb' )

    data = np.fromstring(f.read(model_grid['num_bytes']), dtype='float32')
    if sys.byteorder == 'little':
        data = data.byteswap()
    data = data.reshape(model_grid['jm'], model_grid['im'])

    return data


def main():

    #Initialization dtg
    yyyymmddhh="2015072600"

    #case name
    case_name="realtest3"

    #grid dimensions and #of nests
    ny=278  #390  #278
    nx=262  #390  #262
    nk=40
    nnest=1

    #field
    scalar_fieldname="relvor"  #wwwind"  #relvor"
    uvector_fieldname="uuwind"
    vvector_fieldname="vvwind"
 
    #datapaths
    _datapath="/scratch/fearon/coamps"
    src_datapath=_datapath+"/data/"+case_name+"/atmos"
    dst_datapath=_datapath+"/run/"+case_name
 
    # define the model grid
    model_grid={}
    model_grid['im'] = nx
    model_grid['jm'] = ny
#    model_grid['lm'] = nk
    model_grid['num_bytes'] = model_grid['im'] * model_grid['jm'] * 4

    gridstring = ("%da%.4ix%.4i" % ( nnest, model_grid['im'], model_grid['jm'] ) )

    #lat, lon filenames and data
    lat_file = '_'.join( [ 'latitu', 'sfc', '000000', '000000', gridstring, yyyymmddhh, '00000000', 'fcstfld'] )
    lat_file = os.path.join(src_datapath, lat_file )
    lon_file = '_'.join( [ 'longit', 'sfc', '000000', '000000', gridstring, yyyymmddhh, '00000000', 'fcstfld'] )
    lon_file = os.path.join(src_datapath, lon_file )
    lat = extract_field(lat_file, model_grid)
    lon = extract_field(lon_file, model_grid)

 
    #forecast hours
    fhrs=["00000000", \
          "00030000", \
          "00060000", \
          "00090000", \
          "00120000"]

    levtype="pre"   #pre,sig,sfc,msl
    level="000500"  #"000000"
    for it in range (2,len(fhrs)):
 
        fhr=fhrs[it]
        #scalar
        input_file = '_'.join([scalar_fieldname, levtype, level, '000000', gridstring, yyyymmddhh, fhr, 'fcstfld'])
        input_file = os.path.join(src_datapath, input_file )
        cvar = extract_field(input_file, model_grid)

        # u wind
        input_file = '_'.join([uvector_fieldname, levtype, level, '000000', gridstring, yyyymmddhh, fhr, 'fcstfld'])
        input_file = os.path.join(src_datapath, input_file )
        u = extract_field(input_file, model_grid)

        #v wind
        input_file = '_'.join([vvector_fieldname, levtype, level, '000000', gridstring, yyyymmddhh, fhr, 'fcstfld'])
        input_file = os.path.join(src_datapath, input_file )
        v = extract_field(input_file, model_grid)

        if (scalar_fieldname == "relvor"):
          cvar*=10**4

        print "forecast hour: ", fhr
        mname=scalar_fieldname
        if (fhr == "00000000" and scalar_fieldname == "wwwind"):
          continue  
        else:
          map_contour_vector(cvar,u,v,lat,lon,ny,nx,fhr,mname)
        print cvar

if __name__ == '__main__':
     main()