#!/usr/bin/env python3

# This is an example of writing a grib message with varying precision
#
# input: a.grb
#
# output: 
# a_tmpsfc_original.grb             original grib message before random values added
# a_tmpsfc_default_precision.grb    12 bits binary precision,  ECMWF style
# a_tmpsfc_21bin_precision.grb      17 bits binary precision,  ECMWF style
# a_tmpsfc_1C_keep_min.grb          1 degree K resolution, minimum value retained
# a_tmpsfc_1C_int_min.grb           value to the nearest degree K
# a_tmpsfc_0.1C_int_min.grb         value to the nearest 0.1 degree K

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,':TMP:',':2 m above ground:',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_tmpsfc_original.grb','@mem:0',1)
if result is None:
    quit()
print("write 1:", result)

# add random noise to get bits of precision
TMPsfc = pywgrib2_s.data + numpy.random.random_sample((pywgrib2_s.nx, pywgrib2_s.ny))

# write default: 12 bit binary precistion
result = pywgrib2_s.write('a_tmpsfc_default_precision.grb','@mem:0',1,new_data=TMPsfc)
if result is None:
    quit()
print("write 2:", result)

# write default: 21 bit binary precistion .. ECMWF style
result = pywgrib2_s.write('a_tmpsfc_21bin_precision.grb','@mem:0',1,new_data=TMPsfc,encode_bits=21)
if result is None:
    quit()
print("write 3:", result)

# write with 1K resolution, preserve min value
result = pywgrib2_s.write('a_tmpsfc_1C_keep_min.grb','@mem:0',1,new_data=TMPsfc,d_scale=0,b_scale=0)
if result is None:
    quit()
print("write 4:", result)

# write with 1K, resolution, min is integer
tmp = numpy.round(TMPsfc,0)
result = pywgrib2_s.write('a_tmpsfc_1C_int_min.grb','@mem:0',1,new_data=tmp,d_scale=0,b_scale=0)
if result is None:
    quit()
print("write 5:", result)

# write with 0.1K, resolution, min is integer*0.1K
tmp = numpy.round(TMPsfc,1)
result = pywgrib2_s.write('a_tmpsfc_0.1C_int_min.grb','@mem:0',1,new_data=tmp,d_scale=-1,b_scale=0)
if result is None:
    quit()
print("write 5:", result)


print("done")