from os.path import expanduser
import pandas as pd, numpy as np,os
import matplotlib.pyplot as plt
import glob, gdal

home = expanduser("~")


images = glob.glob("%s/Dropbox/rs/cheese/data/CHEESE*.tif" % home)

output_dir = "%s/Dropbox/rs/cheese/data/" % (home)
if not os.path.isdir(output_dir):
    os.mkdir(output_dir) 

images.sort()
for image in images:
    
    img_base = os.path.basename(image)
    print(img_base)

    source = gdal.Open(image)

    offset_array = np.ones((source.RasterYSize,source.RasterXSize,3)) * -99       
    
    npz_fold = "%s/Downloads/cheese_npz/" % home 

    npzs = glob.glob("%s/%s*.npz" % (npz_fold,img_base))
    
    for npz in npzs:
        start,end = [int(x) for x in os.path.splitext(npz)[0].split('_')[-2:]]
        temp = np.load(npz)['a']
        offset_array[start:end,:,:] =temp
        
    output_name = "%s/%s" % (output_dir,img_base.replace("g.tif","offset.tif"))    
    datatype = gdal.GDT_Float32
    
    # Set the output raster transform and projection properties
    driver = gdal.GetDriverByName("GTIFF")
    tiff = driver.Create(output_name,source.RasterXSize,source.RasterYSize,3,datatype)
    tiff.SetGeoTransform(source.GetGeoTransform())
    tiff.SetProjection(source.GetProjection())
        
    # Write bands to file
    for band in range(3):
        tiff.GetRasterBand(band +1).WriteArray(offset_array[:,:,band])
        tiff.GetRasterBand(band +1).SetNoDataValue(-99)
    
    del tiff, driver