import numpy as np
import camitk

print("-- Create a numpy array (2D 1 component, e.g. black and white picture")
row=10
col=9
count = row*col
a = np.arange(count).reshape(row,col)
# Convert to ImageComponent
imageComp = camitk.newImageComponentFromNumpy(a, "test")
print("-- Convert back, check it is identical to source")
aFromImage = imageComp.getImageDataAsNumpy()
assert(a.shape == aFromImage.shape), f"conversion back failed {a.shape} is not equals to {aFromImage.shape}"
assert (aFromImage == a).all(), f"exported/imported images are different: {aFromImage} != {a}"

print("-- Testing transposed array to check for axis inversion error with non-default stride")
b = np.transpose(a, (1,0))
assert(a.strides != b.strides), f"The transpose should change the strides, not the raw data.\na.strides={a.strides}\nb.strides={b.strides}"
imageCompTr = camitk.newImageComponentFromNumpy(b, "testTr")
bFromImage = imageCompTr.getImageDataAsNumpy()
assert (b.shape == bFromImage.shape), f"shape should be the same {b.shape} != {bFromImage.shape}"
assert (bFromImage == b).all(), f"The strides should be different but the content interpreted the same. {bFromImage.strides} != {b.strides}"
bFromImageTranspose = np.transpose(bFromImage, (1,0))
assert (bFromImageTranspose == a).all(), f"Transpose should be equal to a. a.shape {a.shape} vs bFromImageTranspose.shape {bFromImageTranspose.shape}"

## Same test using 3D 1 component (CT scan)
a = np.arange(300).reshape(5,6,10)
imageComp = camitk.newImageComponentFromNumpy(a, "test3D")
aFromImage = imageComp.getImageDataAsNumpy()
assert (aFromImage == a).all()
b = np.transpose(a, (1,2,0))
assert(a.strides != b.strides) # The transpose should change the strides, not the raw data
imageCompTr = camitk.newImageComponentFromNumpy(b, "test3DTr")
bFromImage = imageCompTr.getImageDataAsNumpy()
assert (bFromImage == b).all() # The strides should be different but the content interpreted the same
assert (np.transpose(bFromImage, (2,0,1)) == a).all() # Transpose should be equal to a

## Same test using 2D 3 components (RGB image)
a = np.arange(300).reshape(10,10,3)
imageComp = camitk.newImageComponentFromNumpy(a, "testRGB")
imageDims = imageComp.getPropertyValue("Image Dimensions")
assert imageDims == "(10, 10, 1)", f"Unexpected Image dimensions: {imageDims} {type(imageDims)}"
aFromImage = imageComp.getImageDataAsNumpy()
assert (aFromImage == a).all()
b = np.transpose(a, (1,2,0))
assert(a.strides != b.strides) # The transpose should change the strides, not the raw data
imageCompTr = camitk.newImageComponentFromNumpy(b, "testRGBTr")
bFromImage = imageCompTr.getImageDataAsNumpy() # as the 3rd dimension size is 10, should be interpreted as 3D b&w image
assert (bFromImage == b).all() # The strides should be different but the content interpreted the same
assert (np.transpose(bFromImage, (2,0,1)) == a).all() # Transpose should be equal to a

## Same test using 4D (3D vector field or 3D with 3 color components) in double
a = np.arange(900, dtype='d').reshape(5,6,10,3)
imageComp = camitk.newImageComponentFromNumpy(a, "test4D-double")
aFromImage = imageComp.getImageDataAsNumpy()
assert (aFromImage == a).all()
b = np.transpose(a, (1,2,0,3))
assert(a.strides != b.strides) # The transpose should change the strides, not the raw data
imageCompTr = camitk.newImageComponentFromNumpy(b, "test4D-doubleTr")
bFromImage = imageCompTr.getImageDataAsNumpy()
assert (bFromImage == b).all(), f"Error: the strides should be different but the content interpreted the same.\nbFromImage={bFromImage}\nb={b}\n"
bTranspose = np.transpose(bFromImage, (2,0,1,3))
assert (bTranspose == a).all(), f"Error: transpose should be equal to a.\ntranspose={bTranspose}\na={a}\n"