#!/home/bfiedler/anaconda3/bin/python3
# 2025 January 16
import glob, os, sys
import numpy as np
from datetime import datetime, timedelta # date and time helpers
import math
import matplotlib.pyplot as plt
from matplotlib import collections  as mc
import cartopy.crs as ccrs
import argparse

# Define the argument parser
parser = argparse.ArgumentParser(description="npy files to png image")

# Positional arguments (like npyfile1 and optional npyfile2)
parser.add_argument("npydir", help="full path to dir with nyp files")
parser.add_argument("pngdir", help="full path to top dir of pngs")  

# Optional arguments
#parser.add_argument("-i", action="store_true", dest="plotit", help="enable plotting")

# Version argument
parser.add_argument("--version", action="version", version="0.1")

# Parse arguments
args = parser.parse_args()
print(args)


print("look for npy in: "+args.npydir)
print("put pngs in: "+args.pngdir)

splitdir=args.npydir.split('/')
timedir=splitdir[-2]
ftime=splitdir[-1]

print(timedir, ftime)

grids = 'lsm0 u1 v1 t1 u500 v500 gh500 gh200 prmsl0'.split()
#ftimes = 'f000 f012  f024 f036 f048 f060 f072 f084 f096 f108 f120 f132 f144 f156 f168 f180'.split()
a={}
for var in grids:
    npy = args.npydir+'/'+ var+'.npy'
    print(npy)
    a[var] = np.load(npy)

#print(a.keys())

if a['prmsl0'].shape[0]==181: #this is GFS one-degree grid
    lons = np.arange(0.,359.1,1)
    lats = np.arange(90,-90.1,-1)
lonsa,latsa = np.meshgrid(lons,lats)

def uptime(ts, dhr=0):
    year=int( ts[:4] )
    month=int(ts[4:6])
    day=int(ts[6:8])
    hr=int(ts[8:])
    uptime = datetime(year,month,day,hr) + timedelta(hours=dhr)
    return uptime 



#def my_cool_plot(timedir,ftime,showit=False):
if True:
    figsur= plt.figure(figsize=(10, 10))
    ax2 = figsur.add_subplot(1, 1, 1, projection=ccrs.Orthographic(-100,55))
    ax3 = figsur.add_axes([0.8, 0.05, 0.15, .2]) # for colorbar
    ax3.axis('off')

    CF=ax2.contourf( lonsa, latsa, a['prmsl0']/100,
                   levels=[950,960,970,980,990,1000,1010,1020,1030,1040,1050], cmap='nipy_spectral', transform=ccrs.PlateCarree()) #,
                  # colors=['purple']*6+['darkgreen']*5)
    
    CC500 = ax2.contour( lonsa, latsa, a['gh500']/10.,colors='darkgreen',linewidths=3,transform=ccrs.PlateCarree())
    fmtc="%4.0d"
    plt.clabel(CC500, fontsize=12, inline=1,fmt=fmtc)
    
    CC200 = ax2.contour( lonsa, latsa, a['gh200']/10.,colors='darkred',linewidths=3,transform=ccrs.PlateCarree())
    plt.clabel(CC200, fontsize=12, inline=1,fmt=fmtc) 
    
    u1 = a['u1'] 
    v1 = a['v1']

    # skip some u,v data.  Certainly skip the u,v at the poles!
    skip=4
    strt=skip//2
    lonsb=lonsa[strt::skip,strt::skip]
    latsb=latsa[strt::skip,strt::skip]
    u=u1[strt::skip,strt::skip]
    v=v1[strt::skip,strt::skip]

    # calculate displacement vectors:
    coslat = np.cos(np.pi*latsb/180.)
    REarth = 6.37e6
    dlondt = u/coslat #/REarth # no need to divide by REarth
    dlatdt = v  #/REarth
    dt=.2 # this means u = 1 m/s at the equator moves dt degrees longitude
    lone=lonsb+dt*dlondt # longitude end point of vector
    late=latsb+dt*dlatdt # latitude end point of vector

    p1=zip(lonsb.flatten().tolist(),latsb.flatten().tolist()) 
    p2=zip(lone.flatten().tolist(),late.flatten().tolist())
    lines=list(zip(p1,p2))
    
    curlat=91 
    lc=0
    linesthinned=[]
    dotslon=[]
    dotslat=[]
    for line in lines:
        thelat=line[0][1]
        lc+=1
        if thelat<curlat:
            lc=0
            curlat=thelat
            iskip=int(1//math.cos(math.radians(curlat)))
            #print(curlat,iskip)
        if lc%iskip!=0: continue
        linesthinned.append(line)
        dotslon.append(line[1][0])
        dotslat.append(line[1][1])
    
    lc = mc.LineCollection(linesthinned,color='black',transform=ccrs.PlateCarree())
    ax2.add_collection(lc)
    ax2.plot(dotslon,dotslat,'m.',transform=ccrs.PlateCarree())
    
    vtime = uptime(timedir , int(ftime[1:]) )
    vtimestamp = vtime.strftime("%Y %m %d %H")
    titl = "%s valid:   %s" % (ftime, vtimestamp)
    ax2.set_title(titl,fontsize=24)
    ax2.text(0.,-.05,'gh200 gh500 prmsl0 (u1,v1)',fontsize=18, alpha=1., transform = ax2.transAxes) 
    
    ax2.coastlines()
    ax2.gridlines()
    ax2.set_global()
    
    cbar = plt.colorbar(CF,ax=ax3) # use matplotlib colorbar
    cbar.set_label( 'prmsl (hPa)' )
    cbar_exists = True

    pngtdir = args.pngdir+'/'+timedir
    if not os.path.exists(pngtdir): os.makedirs(pngtdir)
    pngname = ftime + '.png'
    pngpath = pngtdir + '/' + pngname
    fc = '#ddffdd'
    if ftime=='f000': fc='lightgrey'
    print("saving "+pngpath)
    figsur.savefig( pngpath, facecolor=fc )
print("done")