import os
from osgeo import gdal


# set proj folder
_gdal_data = os.path.join(os.path.dirname(gdal.__file__), "data")
os.environ.setdefault("PROJ_DATA", os.path.join(_gdal_data, "proj"))
os.environ.setdefault("GDAL_DATA", os.path.join(_gdal_data, "gdal"))
os.environ.setdefault("CPL_LOG", os.devnull)

# gdal config
gdal.DontUseExceptions()
gdal.SetConfigOption("SHAPE_ENCODING", "utf-8")
gdal.PushErrorHandler("CPLQuietErrorHandler")
gdal.SetConfigOption("GDAL_FILENAME_IS_UTF8", "YES")
gdal.SetConfigOption("GDAL_NUM_THREADS", "ALL_CPUS")

# default spatial reference system
WGS84 = "EPSG:4326"

# creation options
CREATION = [
    "BIGTIFF=YES",
    "TILED=YES",
    "NUM_THREADS=ALL_CPUS",
    "COMPRESS=ZSTD",
    "PREDICTOR=1",
    "ZSTD_LEVEL=1",
]

# mapping between gdal type and ogr field type
TYPE_MAP = {
    "bool": gdal.GDT_Byte,
    "uint8": gdal.GDT_Byte,
    "int8": gdal.GDT_Byte,
    "uint16": gdal.GDT_UInt16,
    "int16": gdal.GDT_Int16,
    "uint32": gdal.GDT_UInt32,
    "int32": gdal.GDT_Int32,
    "float32": gdal.GDT_Float32,
    "float64": gdal.GDT_Float64,
    "complex64": gdal.GDT_CFloat32,
    "complex128": gdal.GDT_CFloat64,
}

# cpu used
N_CPU = max(
    int(
        int(os.environ.get("SLURM_CPUS_PER_TASK", os.cpu_count()))
        * int(os.environ.get("SLURM_NTASKS", 1))
        * 3
        // 4
    ),
    1,
)
