import os
import numpy as np
from osgeo import gdal

def get_dir_list(dir_path, extension):
    raw_list = os.listdir(dir_path)

    filenames = []
    for _raw in raw_list:
        if _raw.endswith(extension):
            filenames.append(_raw)

    return filenames

def save_array_as_tif(array, proj, trans, no_datas, file_path):
    gdal.UseExceptions()

    paths = file_path.split('/')
    file_dir = '/'.join(paths[:-1])
    file_name = paths[-1]

    if not file_name.endswith('.tiff'):
        raise Exception("filename must end up with '.tiff'")

    array = np.array(array)
    shape = array.shape
    if len(shape) != 3:
        raise Exception("shape of array must be (band, width, height)")

    os.makedirs(file_dir, exist_ok=True)
    band_cnt, height, width = shape

    driver = gdal.GetDriverByName('GTiff')
    output_ds = driver.Create(file_path, width, height, band_cnt, gdal.GDT_Float32)
    output_ds.SetGeoTransform(trans)
    output_ds.SetProjection(proj)

    if len(no_datas) != band_cnt:
        raise Exception("no_datas length must equal to band_cnt")

    for i, no_data in enumerate(no_datas):
        band = output_ds.GetRasterBand(i + 1)
        band.WriteArray(array[i, :, :])
        band.SetNoDataValue(no_data)

    output_ds.FlushCache()