"""
Paint correction graphs (png) used on the USNO-B correction pages.
"""

import math
from cStringIO import StringIO

import Image
import matplotlib
import numpy
from matplotlib import cm
from matplotlib import pyplot


from gavo import api

def toMas(arr):
	return arr*3600.*1000.


def toMasAndCosDelta(arr):
	cosDeltaArr = numpy.array(
		[[math.cos((i-90+0.5)/180*math.pi)]*360 for i in range(180)])
	return arr*cosDeltaArr*3600.*1000.


def getData(column, fixerupper):
	with api.getTableConn() as conn:
		data = list(conn.query(
			"select alpha, delta, %s from ppmxl.usnocorr"%column))
	arr = numpy.zeros((180, 360))
	for ra, dec, val in data:
		arr[int(dec+89.5), int(ra-0.5)] = val
	return fixerupper(arr)


def annotateSkyPlot(ax):
	ax.xaxis.set_major_locator(pyplot.FixedLocator(numpy.arange(0, 361, 40)))
	ax.set_xlabel('Right ascension [deg]')
	ax.yaxis.set_major_locator(pyplot.FixedLocator(numpy.arange(-90, 91, 30)))
	ax.set_ylabel('Declination [deg]')
	ax.text(0, -0.15, "One pixel is $1^\\circ\\times\,1^\\circ$", ha="left",
		transform=ax.transAxes, size=8)


def doPlot(data, label, limits, resultName):
	fig = pyplot.figure()
	ax = fig.add_axes([0.1,0.1,0.8,0.9])
	vmin, vmax = limits
	im = ax.imshow(data, origin="lower", extent=(0.5,359.5,-89.5,89.5),
		cmap=cm.jet_r, vmin=vmin, vmax=vmax)
	cbar = fig.colorbar(im, orientation='horizontal', aspect=25)
	cbar.set_label(label, fontsize=20)
	annotateSkyPlot(ax)
	f = StringIO()
	fig.savefig(f, orientation="landscape", format="png")
	return f.getvalue()


def arrange(plots):
	arranged = Image.new("RGB", (800,580))
	for name, pos in [
			("dalpha.png", (0,0)),
			("ddelta.png", (400,0)),
			("dpmalpha.png", (0,280)),
			("dpmdelta.png", (400,280))]:
		f = plots.pop(0)
		arranged.paste(
			Image.open(StringIO(f)).resize((400, 300), Image.ANTIALIAS),
			pos)
	arranged.save("../static/usnocorrplots.png")

def main():
	api.setDBProfile("trustedquery")
	figures = [
		("d_alpha", toMasAndCosDelta,
			r"$(\alpha_{\rm PPMXL}-\alpha_{\rm USNO})\cos\delta{\rm\ [mas]}$",
			(-600, 600), "dalpha.png"),
		("d_delta", toMas,
			r"$\delta_{\rm PPMXL}-\delta_{\rm USNO}{\rm\ [mas]}$",
			(-600, 600), "ddelta.png"),
		("d_pmalpha", toMas,
			r"$\mu_{\alpha,\rm PPMXL}\cos\delta-"
			r"\mu_{\alpha,\rm USNO}\cos\delta{\rm\ [mas/yr]}$",
			(-14, 14), "dpmalpha.png"),
		("d_pmdelta", toMas,
			r"$\mu_{\delta,\rm PPMXL}-"
			r"\mu_{\delta,\rm USNO}{\rm\ [mas/yr]}$",
			(-14, 14), "dpmdelta.png")]
	plots = []
	for valCol, fixerupper, legend, limits, resName in figures:
		plots.append(doPlot(getData(valCol, fixerupper),
			legend, limits, resName))
	arrange(plots)

if __name__=="__main__":
	main()