#!/usr/bin/env python
"""
 command line program to change packing scheme in a grib2 file

 'grib_repack -h' for help.
"""
from optparse import OptionParser
import sys
import pygrib

usage = """
 %prog [options] grib2filein grib2fileout

 repack the data in grib2filein with a different packing scheme

 example:

'%prog -p3 grib2filein grib2fileout' converts grib2filein to complex packing with spatial differencing, saves result to grib2fileout."""

parser = OptionParser(usage=usage)
parser.add_option("-p", default=False, action="store",dest="ipack",
                        help="desired packing scheme - 0 for simple packing, 1 for simple matrix packing, 2 for complex packing, 3 for complex packing with spatial differencing, 4 for no packing (ieee), 40 for JPEG2000 packing, 41 for PNG packing, 61 for simple log pre-processing")

(options, args) = parser.parse_args()

if not options.ipack:
    sys.stdout.write('must specify a packing scheme with -p\n')
    sys.exit(1)

ipack = int(options.ipack)
if ipack not in [0,1,2,3,4,40,41,61]:
    raise ValueError('packing scheme must be one of 0,1,2,3,4,40,41,61')

# open grib file
try:
    filein = args[0]
    fileout = args[1]
except:
    sys.stdout.write('need to specify input and output grib2 filenames\n')
    sys.exit(1)
sys.stdout.write('input grib file: %s\n' % filein)
grbs = pygrib.open(filein)

# open the file for output.
sys.stdout.write('output grib file: %s\n' % fileout)
f=open(fileout,'wb')

#  "grid_simple"                       = { dataRepresentationTemplateNumber = 0;     }
#  "spectral_complex"                  = { dataRepresentationTemplateNumber = 51; spectralType=1; spectralMode=1; }
#  "spectral_simple"                   = { dataRepresentationTemplateNumber = 50; spectralType=1; spectralMode=1; }
#  "grid_simple_matrix"                = { dataRepresentationTemplateNumber = 1;     }
#  "grid_complex"                      = { dataRepresentationTemplateNumber = 2;     }
#  "grid_complex_spatial_differencing" = { dataRepresentationTemplateNumber = 3;     }
#  "grid_jpeg"                         = { dataRepresentationTemplateNumber = 40000; }
#  "grid_jpeg"                         = { dataRepresentationTemplateNumber = 40;    }
#  "grid_png"                          = { dataRepresentationTemplateNumber = 40010; }
#  "grid_png"                          = { dataRepresentationTemplateNumber = 41;    }
#  "grid_ieee"                         = { dataRepresentationTemplateNumber = 4;     }
#  "grid_simple_log_preprocessing"     = { dataRepresentationTemplateNumber = 61;    }
#  "grid_second_order"                 = { dataRepresentationTemplateNumber = 0;     }
#  "grid_second_order_different_width" = { dataRepresentationTemplateNumber = 0;  }

for grb in grbs:
    if grb['editionNumber'] != 2:
        sys.stdout.write('not a GRIB2 message, skipping ..\n')
        continue
    if grb['dataRepresentationTemplateNumber'] == ipack:
        sys.stdout.write('no repacking required, skipping ..\n')
    grb['dataRepresentationTemplateNumber'] = ipack
    msg = grb.tostring()
    f.write(msg)

f.close()
grbs.close()
