import gdal
import osr
import numpy as np
import subprocess
import os
import glob

gdal.VersionInfo()

driver = gdal.GetDriverByName('GTIFF')

xsize = 2048
ysize = 2048
n_bands = 5
size = n_bands * ysize * xsize

#outdir = r'D:\Temp'
outdir = os.getcwd()

# outfile, gdal dtype, numpy dtype, co options
runs = [('gdal_float32_pb.tif',           gdal.GDT_Float32, np.float32, ['COMPRESS=PACKBITS']),
        ('gdal_float32_pb_tiled.tif',     gdal.GDT_Float32, np.float32, ['COMPRESS=PACKBITS', 'TILED=YES']),
        ('gdal_float32_pb_tiled_512.tif', gdal.GDT_Float32, np.float32, ['COMPRESS=PACKBITS', 'TILED=YES', 'BLOCKXSIZE=512', 'BLOCKYSIZE=512']),
        ('gdal_float32_pb_tiled_128.tif', gdal.GDT_Float32, np.float32, ['COMPRESS=PACKBITS', 'TILED=YES', 'BLOCKXSIZE=128', 'BLOCKYSIZE=128']),
        ('gdal_float32_lzw.tif',          gdal.GDT_Float32, np.float32, ['COMPRESS=LZW']),
        ('gdal_float32_deflate.tif',      gdal.GDT_Float32, np.float32, ['COMPRESS=DEFLATE']),
        ('gdal_float32.tif',              gdal.GDT_Float32, np.float32, []),
        ('gdal_float64_packbits.tif',     gdal.GDT_Float64, np.float64, ['COMPRESS=PACKBITS']),
        ('gdal_float64.tif',              gdal.GDT_Float64, np.float64, []),
        ('gdal_uint16_packbits.tif',      gdal.GDT_UInt16,  np.uint16,  ['COMPRESS=PACKBITS']),
        ('gdal_uint16_lzw.tif',           gdal.GDT_UInt16,  np.uint16,  ['COMPRESS=LZW']),
        ('gdal_uint8_packbits.tif',       gdal.GDT_Byte,    np.uint8,   ['COMPRESS=PACKBITS']),
        ('gdal_uint8_deflate.tif',        gdal.GDT_Byte,    np.uint8,   ['COMPRESS=DEFLATE']),        
        ('gdal_uint8.tif',                gdal.GDT_Byte,    np.uint8,   [])]
        
# dummy data
data = np.random.randint(0,255, size=size).reshape(n_bands, ysize, xsize)

print('{filename: <35}{pysize: >15}{gtsize: >15}{diff: >15}'.format(filename='File',
                                                                    pysize='Python',
                                                                    gtsize='Gdal_trans',
                                                                    diff='Difference'))
    
    
for outfile, gdaltype, nptype, options in runs:

    outfile = os.path.join(outdir, outfile)
    outfile_gdaltranslate = outfile.replace('.tif', '_gdal_translate.tif')

    # file creation with Python bindings
    dst_ds = driver.Create(outfile, xsize, ysize, n_bands, gdaltype, options)

    dst_ds.SetGeoTransform([ 444720, 30, 0, 3751320, 0, -30 ])
    srs = osr.SpatialReference()
    srs.SetUTM(11, 1)
    srs.SetWellKnownGeogCS('NAD27')
    dst_ds.SetProjection(srs.ExportToWkt())

    for i in range(n_bands):
        dst_ds.GetRasterBand(i+1).WriteArray(data[i,:,:].astype(nptype))

    dst_ds = None
    
    # convert with gdal_translate as a reference    
    if len(options) > 0:
        res = subprocess.check_output('gdal_translate -co "%s" %s %s' % ('" -co "'.join(options),
                                                                         outfile, 
                                                                         outfile_gdaltranslate))
    else:
        res = subprocess.check_output('gdal_translate %s %s' % (outfile,
                                                                outfile_gdaltranslate))
    
    # get file sizes
    py_size = os.path.getsize(outfile) / 1024**2
    gt_size = os.path.getsize(outfile_gdaltranslate) / 1024**2
    difference = py_size - gt_size
    
    print('{filename: <35}{pysize:12.1f} MB{gtsize:12.1f} MB{dif:12.1f} MB'.format(filename=os.path.basename(outfile),
                                                                                   pysize=py_size,
                                                                                   gtsize=gt_size,
                                                                                   dif=difference))


# check if the properties/options are actually used
# seems all right

def get_props(infile):
    
    ds = gdal.Open(infile)    
    meta = ds.GetMetadata('IMAGE_STRUCTURE')
    bs = ds.GetRasterBand(1).GetBlockSize()
    
    ds = None    
    
    return meta, bs

print('{filename: <47}{compression: >12}{xblock: >12}{yblock: >12}'.format(filename='filename',
                                                                           compression='Compression',
                                                                           xblock='x-block',
                                                                           yblock='y-block'))

for gdal_file in glob.glob(os.path.join(outdir, '*_gdal_translate.tif')):
    
    meta, bs = get_props(gdal_file)
    filename = os.path.basename(gdal_file).replace('.tif', '')
    print('{filename: <47}{compression: >12}{xblock:12.0f}{yblock:12.0f}'.format(filename=filename,
                                                                                 compression=meta.get('COMPRESSION', '(none)'),
                                                                                 xblock=bs[1],
                                                                                 yblock=bs[0]))
    
    python_file = gdal_file.replace('_gdal_translate', '')
    meta, bs = get_props(python_file)
    filename = os.path.basename(python_file).replace('.tif', '')
    print('{filename: <47}{compression: >12}{xblock:12.0f}{yblock:12.0f}'.format(filename=filename,
                                                                                 compression=meta.get('COMPRESSION', '(none)'),
                                                                                 xblock=bs[1],
                                                                                 yblock=bs[0]))
    
    print('\n')