#!/usr/bin/env python3

#  Sometime you want to merge files
#
#    example 1.  you have a global model and several regional models
#                and you want a global grid with regional models overlayed
#                on top of the global model.  (This is assuming that the
#                regional models are better than the global model
#
#    example 2.  Some models produce "tiles" rather than a full domain files.
#                The user is suppose to download the "tiles" and "quilt" them
#                together to make a field of the relevant region.  This saves
#                bandwidth as the full domain doesn't have to be downloaded.
#
#                I think that tiles and quilting are a mixed metaphor but
#                that is NCEP terminology.

import pywgrib2_s
import numpy

# want to do:  a[numpy.isnan(a)] = b
#  works if b is a constant, but doesn't work for b[:,:]
def merge(a,b):
    print("a is fortran ",numpy.isfortran(a))
    print("b is fortran ",numpy.isfortran(b))
    nx=a.shape[0]
    ny=a.shape[1]
    for j in range(ny):
        for i in range(nx):
            if numpy.isnan(a[i,j]):
                a[i,j] = b[i,j]

out='b.f006.merged'
# list of files from high to low priority
files=[ 'b.f006.ne', 'b.f006.nw', 'b.f006.se', 'b.f006.sw' ]

# take files and regrid to global 1x1 grid
# make inventories of the new files

n=len(files)
for i in range(n):
    file=files[i]
    file_order=file+'.order'
    file_missing=file+'.missing'
    file_new=file+'.new'
    file_new_inv=file+'.new.inv'
    print("processing ", file)

    # put file in order for regridding

    err=pywgrib2_s.wgrib2([file,'-new_grid_order',file_order,file_missing,'-inv','/dev/null'])
    print("put file in order for -new_grid err=", err, file)

    # call wgrib2 to regrid the file
    # change mode from write to read, no need to rewind
    err = pywgrib2_s.wgrib2( [file_order, '-new_grid_winds', 'earth', '-set_grib_type' , 'c3', 
      '-new_grid', 'latlon', '0:360:1.0', '-90:181:1.0',  file_new, '-inv', '/dev/null' ] )
    print("regrid err=", err)
    if err != 0:
        print("regrid failed")
        quit()

    err = pywgrib2_s.mk_inv(file_new,file_new_inv)
    print("mk_inv err=", err)

# get metadata for files[0]
metas=pywgrib2_s.read_inv(files[0]+'.new.inv')

# change metadata into search string
for i in range(len(metas)):
    # get date code
    meta=metas[i]
    j=meta.find(":D=")
    meta=meta[j+1:]
    j=meta.find("::")
    meta=meta[0:j]
    metas[i]=meta

# process each field in meta
for i in range(len(metas)):

    # read file 0
    nmatch=pywgrib2_s.inq(files[0]+'.new',metas[i],grib='@mem:0',inv=files[0]+'.new.inv',Data=True)
    print('file:  ',i,' nmatch=',nmatch)
    if nmatch != 1:
        print("problem reading file0 ",files[0])
        quit()
    grid = numpy.copy(pywgrib2_s.data)
    print("grid is fortran ",numpy.isfortran(grid))
    print("data is fortran ",numpy.isfortran(pywgrib2_s.data))
    print(pywgrib2_s.nx)
    print(pywgrib2_s.ny)
    print(numpy.shape(grid))
    pywgrib2_s.write('junk2.grb','@mem:0',1,new_data=pywgrib2_s.data)
    pywgrib2_s.write('junk.grb','@mem:0',1,new_data=grid)
    for j in range(1,len(files)):
        nmatch=pywgrib2_s.inq(files[j]+'.new',metas[i],inv=files[j]+'.new.inv',Data=True)
        if nmatch == 1:
            merge(grid,pywgrib2_s.data)
            print(pywgrib2_s.data.shape)
        else:
            print("missing ",metas[i], ' in file: ',files[j])

    # write using file1 metadata, new data
    pywgrib2_s.write(out,'@mem:0',1,new_data=grid)

print("done ",out)