#!/usr/bin/env python3

# This is an example of writing a grib message with various packing
#
# input: a.grb
#
# output: 

import pywgrib2_s
import numpy
import math

in_file='a.grb'

# read a field TMPsfc, save grib in @mem:0 (memory file 0)
# and raw field in pywgrib2_s.data

nmatch = pywgrib2_s.inq(in_file,':SHTFL:',':surface:',time0=2020010200,grib='@mem:0',Data=True)
if nmatch != 1:
    print("read of TMPsfc failed nmatch=",nmatch)
    quit()

# write original grib message
result = pywgrib2_s.write('a_shtfl_original.grb','@mem:0',1)
if result is None:
    quit()
print("write original:", result)

shtfl = pywgrib2_s.data

# write jpeg2000
result = pywgrib2_s.write('a_shtfl_jpeg2000.grb','@mem:0',1,new_data=shtfl,packing='j')
if result is None:
    quit()
print("write jpeg2000:", result)

# write complex-3
result = pywgrib2_s.write('a_shtfl_c3.grb','@mem:0',1,new_data=shtfl,packing='c3')
if result is None:
    quit()
print("write c3:", result)


print("encode grib messages with c1/c2/c3, save in mem1,mem2,mem3")

# write complex-1 to @mem:1
result = pywgrib2_s.write('@mem:1','@mem:0',1,new_data=shtfl,packing='c1')
c1=pywgrib2_s.mem_size(1)
print("size for c1 = ", c1)

# write complex-2 to @mem:2
result = pywgrib2_s.write('@mem:2','@mem:0',1,new_data=shtfl,packing='c2')
c2=pywgrib2_s.mem_size(2)
print("size for c2 = ", c2)

# write complex-3 to @mem:2
result = pywgrib2_s.write('@mem:3','@mem:0',1,new_data=shtfl,packing='c3')
c3=pywgrib2_s.mem_size(3)
print("size for c3 = ", c3)

if c1 <= c2 and c1 <= c3:
  print("c1 is smallest")
  grb_mess = pywgrib2_s.get_bytes_mem(1)
elif c2 <= c3:
  print("c2 is smallest")
  grb_mess = pywgrib2_s.get_bytes_mem(2)
else:
  print("c3 is smallest")
  grb_mess = pywgrib2_s.get_bytes_mem(3)

file = open('a_shtfl_smallest.grb','wb')
file.write(grb_mess)
file.close()
quit()