
import numpy as np
import pytest
from osgeo import gdal

from yirgacheffe.operators import DataType
from yirgacheffe.backends import backend, BACKEND
from yirgacheffe.layers import RasterLayer

from helpers import gdal_dataset_with_data

@pytest.mark.parametrize("gtype", [
	gdal.GDT_Int8,
	gdal.GDT_Int16,
	gdal.GDT_Int32,
	gdal.GDT_Int64,
	gdal.GDT_Byte,
	gdal.GDT_UInt16,
	gdal.GDT_UInt32,
	gdal.GDT_UInt64,
	gdal.GDT_Float32,
])
def test_round_trip(gtype) -> None:
	ytype = DataType.of_gdal(gtype)
	backend_type = backend.dtype_to_backed(ytype)
	assert backend.backend_to_dtype(backend_type) == ytype

@pytest.mark.parametrize("ytype", [
	DataType.Int8,
	DataType.Int16,
	DataType.Int32,
	DataType.Int64,
	DataType.UInt8,
	DataType.UInt16,
	DataType.UInt32,
	DataType.UInt64,
	DataType.Float32,
	DataType.Float64,
])
def test_round_trip_from_gdal(ytype) -> None:
	gtype = ytype.to_gdal()
	assert DataType.of_gdal(gtype) == ytype

def test_round_trip_float64() -> None:
	backend_type = backend.dtype_to_backed(DataType.Float64)
	ytype = backend.backend_to_dtype(backend_type)
	print(BACKEND, "sad")
	match BACKEND:
		case "NUMPY":
			assert ytype == DataType.Float64
		case "MLX":
			assert ytype == DataType.Float32

def test_float_to_int() -> None:
	data1 = np.array([[1.0, 2.0, 3.0, 4.0], [5.5, 6.5, 7.5, 8.5]])

	layer1 = RasterLayer(gdal_dataset_with_data((0.0, 0.0), 0.02, data1))
	# Note the float 32 here is to rule out that writing the result to the
	# new dataset was what caused the truncation
	result = RasterLayer.empty_raster_layer_like(layer1, datatype=DataType.Float32)

	comp = layer1.astype(DataType.UInt8)
	comp.save(result)

	expected = backend.promote(np.array([[1, 2, 3, 4], [5, 6, 7, 8]]))
	actual = backend.demote_array(result.read_array(0, 0, 4, 2))
	assert (expected == actual).all()
