#!/usr/bin/env python3

import sys
import pywt
from matplotlib.pyplot import (imread, imsave)
from numpy import float32

file = sys.argv[1:]

if file:

	print(f"Using external data from {file[0]}")
	data_image = imread(file[0])

	# check grayscale
	if not isinstance(data_image[0][0], float32):
		print("error: image not gray : convert color.png -grayscale average gray.png")
		exit(1)
else:

	print("Using internal data")
	data_image = [
		[91,   93,  94,  94, 93,   95, 100, 105],
		[93,   94,  93,  91, 89,   89, 93,   97],
		[91,   91,  92,  92, 92,   94, 98,  101],
		[104, 106, 108, 110, 112, 112, 113, 113],
		[123, 125, 128, 131, 131, 129, 126, 123],
		[113, 115, 119, 124, 129, 131, 129, 127],
		[101,  99, 100, 104, 111, 118, 121, 121],
		[113, 107, 100,  97, 97,  104, 107, 108],
	]

print("\nDisplay input data\n")
print(data_image)

coeffs = pywt.dwt2(data_image, 'haar')
LL, (LH, HL, HH) = coeffs

print("\nDisplay wavelet output\n")
print(f"LL=\n{LL}\n")
print(f"LH=\n{LH}\n")
print(f"HL=\n{HL}\n")
print(f"HH=\n{HH}\n")

print("\nSave wavelet output to png...")
imsave("wavelet-output-1-LL.png", arr=LL, cmap='gray')
imsave("wavelet-output-2-LH.png", arr=LH, cmap='gray')
imsave("wavelet-output-3-HL.png", arr=HL, cmap='gray')
imsave("wavelet-output-4-HH.png", arr=HH, cmap='gray')


# ---- reconstruction ----- (to be tested)
# imgr = pywt.idwt2(coeffs,'haar')
# imgr = np.uint8(imgr)
# plt.figure(figsize=(30,30))
# plt.show()
# plt.imshow(imgr,interpolation="nearest",cmap=plt.cm.gray)
# plt.title('reconstruction image',fontsize=40)