"""
Some miscellaneous helpers for making images and such.

As this may turn into a fairly expensive import, this should *not* be imported
by utils.__init__.   Hence, none of these functions are in gavo.api or
gavo.utils.
"""

#c Copyright 2008-2024, the GAVO project <gavo@ari.uni-heidelberg.de>
#c
#c This program is free software, covered by the GNU GPL.  See the
#c COPYING file in the source distribution.


import io

from PIL import Image
import numpy

from gavo.utils.dachstypes import (Filename, NDArray)


def _normalizeForImage(pixels: NDArray, gamma: float) -> NDArray:
	"""helps jpegFromNumpyArray and friends.
	"""
	pixels = numpy.flipud(pixels)
	pixMax, pixMin = numpy.max(pixels), numpy.min(pixels)
	pixels = numpy.power((pixels-pixMin)/(pixMax-pixMin), gamma)*255
	return numpy.asarray(pixels, numpy.uint8)


def jpegFromNumpyArray(pixels: NDArray, gamma: float=0.25) -> bytes:
	"""returns a normalized JPEG for numpy pixels.

	pixels is assumed to come from FITS arrays, which are flipped wrt to
	jpeg coordinates, which is why we're flipping here.

	The normalized intensities are scaled by v^gamma; we believe the default
	helps with many astronomical images
	"""
	f = io.BytesIO()
	Image.fromarray(_normalizeForImage(pixels, gamma)
		).save(f, format="jpeg")
	return f.getvalue()


def colorJpegFromNumpyArrays(
		rPix: NDArray, gPix: NDArray, bPix: NDArray, gamma: float=0.25
		) -> bytes:
	"""as jpegFromNumpyArray, except a color jpeg is built from red, green,
	and blue pixels.
	"""
	pixels = numpy.array([
		_normalizeForImage(rPix, gamma),
		_normalizeForImage(gPix, gamma),
		_normalizeForImage(bPix, gamma)]).transpose(1,2,0)

	f = io.BytesIO()
	Image.fromarray(pixels, mode="RGB").save(f, format="jpeg")
	return f.getvalue()


def scaleNumpyArray(arr: NDArray, destSize: int) -> NDArray:
	"""returns the numpy array arr scaled down to approximately destSize.
	"""
	origWidth, origHeight = arr.shape
	size = max(origWidth, origHeight)
	scale = max(1, size//destSize+1)
	destWidth, destHeight = origWidth//scale, origHeight//scale

	# There's very similar code in fitstools.iterScaledRows
	# -- it would be nice to refactor things so this can be shared.
	img = numpy.zeros((destWidth, destHeight), 'float32')

	for rowInd in range(destHeight):
		wideRow = (numpy.sum(
			arr[:,rowInd*scale:(rowInd+1)*scale], 1, 'float32'
			)/scale)[:destWidth*scale]
		# horizontal scaling via reshaping to a matrix and then summing over
		# its columns.
		newRow = numpy.sum(
			numpy.transpose(wideRow.reshape((destWidth, scale))), 0)/scale
		img[:,rowInd] = newRow

	return img


def getScaledPNG(srcFile: Filename, newWidth: int) -> bytes:
	"""returns PNG bytes that is a scaled version of an image in srcFile.

	srcFile must correspond to something that PIL can read.  Since scaling
	really sucks for non-RGB images, we unconditionally convert whatever we
	get to 3-band RGB.
	"""
	im = Image.open(srcFile).convert("RGB")
	scale = newWidth/float(im.size[0])
	scaled = im.resize((newWidth, int(im.size[1]*scale)), Image.LANCZOS)
	f = io.BytesIO()
	scaled.save(f, format="png")
	return f.getvalue()
