import gdal,glob, sys
import math
import numpy as np, pandas as pd
from os.path import expanduser
home = expanduser("~")
from scipy import interpolate
import matplotlib.pyplot as plt

def create_mesh(Xgeo,Ygeo,resolution):
    """
    Create meshgrid based on input X,Y geo coordinate and the target resolution
    :param Xgeo: np array. X geo coordinate (longitude) for pixels from the raw img.
    :param Ygeo: np array. Y geo coordinate (latitude) for pixels from the raw img.
    :param resolution: target resolution in [x_resolution,y_resolution], should be sharing the same unit as the Xgeo and Ygeo.
    :return: grid mesh for the target image.
    """
    # Locate the ul coordinate
    x_min = np.min(Xgeo)
    x_max = np.max(Xgeo)
    y_min = np.min(Ygeo)
    y_max = np.max(Ygeo)


    # create the mesh grid:
    x_range = np.arange(x_min,x_max,resolution[0])
    y_range = np.arange(y_min,y_max,resolution[1])
    xx, yy = np.meshgrid(x_range, y_range, sparse=True)

    return xx, yy




# Directories:
images = glob.glob("%s/GoogleDrive/CheeseHead/Data/Hyspex/cheese/data/*Refl_g.tif" % home)
images.sort()
nodata = -1
datatype = gdal.GDT_Int16
for i, image in enumerate(images):
    print(image)
    raw = gdal.Open(image)
    GT = raw.GetGeoTransform()
    value = raw.GetRasterBand(1).ReadAsArray()
    NewCoor = gdal.Open(image.replace('Refl_g.tif','Refl_newCoordinate'))
    Xgeo = NewCoor.GetRasterBand(1).ReadAsArray()
    Ygeo = NewCoor.GetRasterBand(2).ReadAsArray()



    xx, yy = create_mesh(Xgeo, Ygeo,[GT[1],abs(GT[5])])


    # Interpolate the image:
    new = interpolate.griddata((Xgeo.flatten(), Ygeo.flatten()), value.flatten(), (xx, yy), method = 'nearest', fill_value=nodata )

    # Write to disk:
    # Set the output raster transform and projection properties
    output_name = image.replace('Refl_g', 'Refl_g_corrected')
    driver = gdal.GetDriverByName("GTIFF")
    tiff = driver.Create(output_name, new.shape[1], new.shape[0], 1, datatype)
    new_GT = list(GT)
    new_GT[0] = np.min(xx)
    new_GT[3] = np.max(yy)
    tiff.SetGeoTransform(new_GT)
    tiff.SetProjection(raw.GetProjection())

    # Write bands to file
    # flip the data:
    new_flip = np.flipud(new)
    tiff.GetRasterBand(1).WriteArray(new_flip)
    tiff.GetRasterBand(1).SetNoDataValue(nodata)

    del tiff, driver