#!/home/bfiedler/anaconda3/bin/python3
# v0.3 3 July 2017
import pygrib
import sys,os,glob
import numpy as np
def dir2npy(dirname,overwrite=False):
    look = dirname + '/' + '*.grb2'
    grnames = glob.glob(look)
    print(grnames)
    for grname in grnames:
        print(grname)
        ret=gr2npy(grname,overwrite)
        print(ret)
def gr2npy(gribname, overwrite=False):
    top = "/home/bfiedler/play/grib/gfs1degnpy/"
    print("open: ",gribname)
    a = pygrib.open(gribname)
    names = []
    wantlev = [1000,850,700,500,300,200]
    wantsn = 'absv u v w t r gh'.split()
    othersn = 'cape cin lsm cprat prate tp acpcp cpofp pwat prmsl'.split()
    g = a[1] # sample first record in grib file, to get info
    spl = repr(g).split(':')
    start = spl[-1]
    start = start.split()[-1]
    start = start.strip()
    start = start[0:10]
    gname = gribname.split('.')[-2]
    ftime = 'f'+gname.split('_')[-1] 
    startdir = top+start # start means start time
    print(startdir)
    summaries=[]
    if not os.path.exists(startdir):
        os.mkdir(startdir)
        command = "ln -s -f "+startdir+" "+top+"latest"
        #os.system(command)
    outdir = startdir+'/'+ftime
    if not os.path.exists(outdir):
        os.mkdir(outdir)
    elif not overwrite:
        return outdir+' exists, no overwrite permission' 
    
     
    for g in a:
        spl = repr(g).split(':')
        num = spl[0]
        sn  = g.shortName 
        lev = g.level 
        tlev = g.typeOfLevel
        iso = tlev == "isobaricInhPa" and lev in wantlev and sn in wantsn
        surf = lev==1 and tlev=="sigma" and sn in wantsn
        if iso or surf or sn in othersn:
            fn = sn+str(lev)+'.npy'
            if fn in names:
                print("duplication:",fn)
            else:
                print(fn,repr(g))
                summaries.append(fn+' '+repr(g))
                names.append(fn)
                data = g.values
                np.save(outdir+'/'+fn,data)
        readme = '\n'.join(summaries)+'\n'
        outreadme = open(outdir+'/'+'Readme.txt','w')
        outreadme.write(readme)
 
        outreadme.close()
    return outdir+' filled'

if __name__=='__main__':
    from optparse import OptionParser
    usage="usage: %prog [options] filenames"
    parser = OptionParser(usage=usage, version="%prog 0.2")
    parser.add_option("--ov", dest="ov",action="store_true",
                    help="overwrite existing")
    (opt, arg) = parser.parse_args()
    names = arg[0:]
    for name in names: 
        if name[-5:]=='.grb2':
            ret = gr2npy(name,overwrite=opt.ov)
            print(ret)
        else:
            dir2npy(name, overwrite=opt.ov)