import os

import numpy

from gavo.helpers import processing
from gavo.utils import imgtools
from gavo import api
api.LoggingUI(api.ui)


class PreviewMaker(api.PreviewMaker):
	def getPreviewData(self, srcName):
		if "cube" in srcName:
			arr = api.pyfits.open(
				os.path.join(api.getConfig("inputsDir"), srcName))[0].data
			arr = numpy.ma.masked_invalid(arr)
			return imgtools.colorJpegFromNumpyArrays(
				sum(arr[:10, :, :], 0),
				sum(arr[500:550, :, :], 0),
				sum(arr[990:1000, :, :], 0), gamma=0.2)

		else:
			arr = numpy.ma.masked_invalid(api.pyfits.open(
				os.path.join(api.getConfig("inputsDir"), srcName))[0].data)
			return imgtools.jpegFromNumpyArray(arr)



if __name__=="__main__":
	processing.procmain(PreviewMaker, "ppakm31/q", "import_cubes")
	processing.procmain(PreviewMaker, "ppakm31/q", "import")