#!/usr/bin/env python3 

# read and write a grib message
#  1 make a copy using grib=file option in pywgrib2_s.inq
#  2 write using pywgrib2_s.write
#     a) input is ndarray fortran, float32
#     b) input is ndarray C, float64
#     c)

import pywgrib2_s
import numpy

in_file='a.grb'

out_copy='a_copy.grb'
out_write='a_write.grb'
out_write_c='a_write_c.grb'
out_write_l='a_write_l.grb'

# only doing one pywgrib2_s.inq, save time by not making inv file
# time0 takes YYYYMMDDHH and YYYYMMDDHHmmSS
#    time0 requires year >= 1000 and year <= 9999

# ok nmatch=pywgrib2_s.inq(in_file, ':d=2020013118:',':WEASD:',grib=out_copy,Data=True,Latlon=True)
# ok nmatch=pywgrib2_s.inq(in_file, ':D=20200131180000:',':WEASD:',grib=out_copy,Data=True,Latlon=True)
# ok nmatch=pywgrib2_s.inq(in_file,':WEASD:',time0=2020013118,grib=out_copy,Data=True,Latlon=True)
nmatch=pywgrib2_s.inq(in_file,':WEASD:',time0=20200131180000,grib=out_copy,Data=True,Latlon=True)
if nmatch != 1:
    print("expected 1 match, found nmatch=",nmatch)
    quit()

data=pywgrib2_s.data
print("data shape=",numpy.shape(data))
print("data isfortran=",numpy.isfortran(data))
print("bottom left lat/lon,val=",pywgrib2_s.lat[0,0],pywgrib2_s.lon[0,0], data[0,0])

# write with data=ndarray, Fortran, float32
field=pywgrib2_s.write(out_write,out_copy,1,new_data=data)

# data2 is in C order, float 64
data2=numpy.zeros(numpy.shape(data),dtype=numpy.float64,order='C')

s=numpy.shape(data)
for i in range(s[0]):
    for j in range(s[1]):
        data2[i,j] = data[i,j]

print("data2 shape=",numpy.shape(data2))
print("data2 isfortran=",numpy.isfortran(data2))
# write with data2=ndarray, c, float64
field=pywgrib2_s.write(out_write_c,out_copy,1,new_data=data2)
print("write field=",field)

# now for numpy 1-d array of integers

vector=numpy.zeros((numpy.size(data)),dtype=numpy.int)
k=0
for j in range(s[1]):
    for i in range(s[0]):
        vector[k] = int(data[i,j])
        k += 1
field=pywgrib2_s.write(out_write_l,out_copy,1,new_data=vector)
print("write field=",field)

print("done copy=", out_copy, " write=",out_write, " write_c=",out_write_c, " write_l=",out_write_l)