# Simple demo script for retrieving LS colour data. By default, retrieves a colour LS
# image from specific RA, Dec, FOV. Commented sections illustrate how to use an input
# FITS file instead of coordinates - will grid the retrieved image to the correct pixel
# dimensions.
# Full documentation on HIPS2FITS : 
# https://astroquery.readthedocs.io/en/latest/api/astroquery.hips2fits.hips2fitsClass.html


import astroquery
import PIL
import astropy
from astroquery.hips2fits import hips2fits
from matplotlib.colors import Colormap
import astropy.units as u
from astropy.coordinates import Longitude, Latitude, Angle
from astropy.io import fits as pyfits
from astropy.wcs import WCS
import numpy
import matplotlib.pyplot as plt
import PIL
from PIL import Image


# 1) Simple input parameters from position, FOV
name = 'VCC1964'	# Source name, for saving the output file
ra  =  190.833		# Decimal degrees
dec =  8.95306		# Decimal degrees
size = 14.0		# FOV (square) in arcminutes
# Can also specify width, height of the output image in pixels (these are set below).
# Doesn't appear to have a parameter to directly set the pixel scale, but documentation
# says that FOV is always that of the largest dimension.


# 2) ALTERNATIVE : Specify a FITS file and it will retrieve an image of the same WCS
#FitsFileName = 'SDSS_g.fits'
#FitsFile = pyfits.open(FitsFileName)
#FitsData = FitsFile[0].data
#proj1 = WCS(FitsFile[0].header)


# SELECT DATA SOURCE
# Some examples below. Note can also specify individual wavebands and many other data
# sources e.g. Planck, 2MASS, HST...
# Full list : https://aladin.cds.unistra.fr/hips/list
# 1) DESI Legacy Survey
hips = 'CDS/P/DESI-Legacy-Surveys/DR10/color'
# 2) SDSS DR9 RGB images (no other data releases available)
#hips = 'CDS/P/SDSS9/color'
# 3) SDSS DR9 alternative RGB colour scheme
#hips = 'CDS/P/SDSS9/color-alt'
# 4) DSS2 colour 
#hips = 'CDS/P/DSS2/color'
# 5) GALXEX
#hips = 'CDS/P/GALEXGR6/AIS/color'


# 2) FOR USING FITS FILES AS INPUT
#result = hips2fits.query_with_wcs(
#    hips=hips,
#    wcs=proj1,
#    get_query_payload=False,
#    format='jpg',
#    #min_cut=0.5,
#    #max_cut=99.5,
#    cmap=Colormap('viridis'),
#)

#result_flipped=numpy.flip(result,axis=0) #otherwise the color image plots flipped

#fig = plt.figure(figsize=(100, 100))
#wcs1 = proj1
#ax1 = plt.subplot(4,3,1,projection=wcs1,label='overlays')
#ax1.imshow(result_flipped)
##im1 = ax1.contour(FitsData/1000.0,transform=ax1.get_transform(WCS(FitsFile[0].header)), cmap='autumn',linewidths=2.0, levels=numpy.logspace(-2., 4.,10))

#plt.savefig('TEST.png',dpi=100,facecolor='white',bbox_inches='tight')


# Coordinate-based query
result = hips2fits.query(
   hips=hips,
   width=512,
   height=512,
   ra=Longitude(ra * u.deg), 
   dec=Latitude(dec* u.deg),
   fov=Angle(size * u.arcmin),
   projection="TAN",
   get_query_payload=False,
   format='jpg',
   min_cut=0.0,
   max_cut=100.0,
   cmap=Colormap('viridis'),
)

result_flipped=numpy.flip(result,axis=0) # Otherwise the color image plots flipped

im = plt.imshow(result_flipped)


# Could simply save the image directly with matplotlib, but this will add padding and other
# crap to the image, which is annoying, so don't do that
#plt.savefig('VSCE_Extra_006_RGB.png',dpi=100,facecolor='white',bbox_inches='tight')


# Using PIL - easiest way to ensure exact image dimensions with no axes plotted
colmap = Image.fromarray(result)

colmap.save(name+'.png')

# To save the result as a FITS file, specify a HIPS2FITS single waveband as the source, e.g.
#hips = 'CDS/P/DESI-Legacy-Surveys/DR10/g'
# Change the "format" parameter to FITS but keep the rest of the query unchanged. Save
# the result simply with :
#result.writeto(name+'.fits')