import os
from osgeo import gdal
import numpy as np
import time


def get_tifflayer(filepath, dtype=np.float32):
    ds = gdal.Open(filepath)
    return np.asarray(ds.GetRasterBand(1).ReadAsArray(), dtype=dtype)

def get_geo(filepath):
    ds = gdal.Open(filepath)
    return ds.GetGeoTransform(), ds.GetProjection()

def saveGeoTiff(filename, data, dim, trans, prj, dtype=gdal.GDT_Float32,nodata_value=-9999):
    if os.path.exists(filename):
        os.remove(filename)
        time.sleep(1)

    outds = gdal.GetDriverByName('GTiff').Create(filename, dim[0], dim[1], dim[2],dtype)
    outds.SetGeoTransform(trans)
    outds.SetProjection(prj)

    band = outds.GetRasterBand(1)
    band.SetNoDataValue(nodata_value)

    #10/22/2020
    #Hua's code generates nan values
    #then he uses things like np.nanmedian to ignore the nan values
    #maybe his code should be redone so it doesn't generate nans
    #but I haven't looked into it
    #if i had to guess I would say when it does something like
    # mean and there are no valid pixels, then its doing divide by numpixels
    # which is divide by 0
    #and there was a meeting with the UHI team where it was decided
    #that nan would be reassigned to the nodata value
    #so thats what I am doing here
    #Reza

    fixed_data = np.nan_to_num(data,nan=nodata_value)


    if dim[2] == 1:
        outds.GetRasterBand(1).WriteArray(fixed_data)
    else:
        for i in range(dim[2]):
            outds.GetRasterBand(i + 1).WriteArray(fixed_data[:, :, i])
    outds = None
    return True

