import numpy as np
from osgeo import gdal, osr


def write_raster(data):
    pixel_array = []
    pixel_array_alpha = []
    for j in range(0, 36):
        row_array = []
        row_array_alpha = []
        for i, station_id in enumerate(data):
            if i % 36 == 0:
                value = data[i + j][1]
                value = 0 if not value else value
                value = 0 if str(value) == 'NaN' else value
                alpha_value = data[i + j][2]
                alpha_value = 0 if alpha_value else 255
                # print(value)
                row_array.append(value)
                row_array_alpha.append(alpha_value)
                np_row_array = np.array(row_array)
                np_row_array_alpha = np.array(row_array_alpha)
                # np_row_array_alpha = np.where(np_row_array_alpha < 8, 0, 255)
                # print(row_array)
        pixel_array.append(np_row_array)
        pixel_array_alpha.append(np_row_array_alpha)
    np_pixel_array = np.array(pixel_array)
    np_pixel_array_alpha = np.array(pixel_array_alpha)

    xmin, ymin, xmax, ymax = [5.01, 47.15, 14.81, 55.33]
    nrows, ncols = np.shape(np_pixel_array)
    xres = (xmax - xmin) / float(ncols)
    yres = (ymax - ymin) / float(nrows)
    geotransform = (xmin, xres, 0, ymax, 0, -yres)

    output_raster = gdal.GetDriverByName('GTiff').Create('D:/Uni/Master/01_SS2021/Automatisierte_Geodatenprozessierung/temperaturverteilung/dataacquisition/output/myraster.tif', ncols, nrows, 4, gdal.GDT_Float32)  # Open the file
    output_raster.SetGeoTransform(geotransform)
    srs = osr.SpatialReference()
    srs.ImportFromEPSG(4326)
    output_raster.SetProjection(srs.ExportToWkt())
    output_raster.GetRasterBand(1).WriteArray(np_pixel_array)
    # output_raster.GetRasterBand(2).WriteArray(np_pixel_array)
    # output_raster.GetRasterBand(3).WriteArray(np_pixel_array)
    output_raster.GetRasterBand(4).WriteArray(np_pixel_array_alpha)

    output_raster.FlushCache()