# test alpha blending when both alpha and array alpha are passed
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap
import numpy as np

fig, axs = plt.subplots(3, 4, figsize=(12, 10), layout='compressed')

for ax in axs.flat:
    ax.set(facecolor='red', xticks=[], yticks=[])

mapped = np.array([
    [0.1, 1.0],
    [1.0, 0.1]
])
rgb = np.repeat(mapped[:, :, np.newaxis], 3, axis=2)
rgba = np.concatenate(
    [
        rgb,
        [
            [[1.0], [0.9]],
            [[0.8], [0.7]],
        ]
    ],
    axis=2
)

alpha_scalar = 0.5
alpha_2d = np.full_like(mapped, alpha_scalar)

cmap_with_alpha = ListedColormap(
    np.concatenate([plt.cm.viridis.colors,
                    np.full((len(plt.cm.viridis.colors), 1), alpha_scalar)],
                   axis=1),
)



# axs[0,0].imshow(rgb, alpha=alpha_scalar)
# axs[0,1].imshow(rgb, alpha=alpha_2d)
for ax, alpha, t in zip(axs, [None, alpha_scalar, alpha_2d], ['off', 'float', 'array']):
    ax[0].imshow(mapped, alpha=alpha)
    ax[0].set_title(f'2D, alpha={alpha_scalar} {t}')

    ax[1].imshow(mapped, cmap=cmap_with_alpha, alpha=alpha)
    ax[1].set_title(f'2D with {alpha_scalar} alpha cmap, alpha={alpha_scalar} {t}')

    ax[2].imshow(rgb, alpha=alpha)
    ax[2].set_title(f'RGB, alpha={alpha_scalar} {t}')

    ax[3].imshow(rgba, alpha=alpha)
    ax[3].set_title(f'RGBA, alpha={alpha_scalar} {t}')

plt.show()