import gdal


def write_tiff_mask(filename, image):
    rows = image.shape[0]
    cols = image.shape[1]
    print "image size" + str(rows) + str(cols)
    driver = gdal.GetDriverByName("GTiff")
    outdata = driver.Create(filename, cols, rows, 1, gdal.GDT_UInt16)
    if outdata is None:
        print 'Could not create' + filename
        return

    outband = outdata.GetRasterBand(1)
    outband.WriteArray(image)

    # flush data to disk, set the NoData value and calculate stats
    outband.FlushCache()
    outband.SetNoDataValue(255)


def write_tiff(filename, image, type='uint16'):
    if type == 'float':
        dtype = gdal.GDT_Float64
    else:
        dtype = gdal.GDT_UInt16
    channels = image.shape[0]
    rows = image.shape[1]
    cols = image.shape[2]
    driver = gdal.GetDriverByName("GTiff")
    outdata = driver.Create(filename, cols, rows, channels, dtype)
    if outdata is None:
        print 'Could not create' + filename
        return

    for i in range(channels):
        outBand = outdata.GetRasterBand(i + 1)
        outBand.WriteArray(image[i, :, :])
        # flush data to disk, set the NoData value and calculate stats
        outBand.FlushCache()
        outBand.SetNoDataValue(255)
    outdata = None
    del outdata, outBand