#! /usr/bin/env python

# Copyright (C) 2019 Collin Capano
#
# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the
# Free Software Foundation; either version 3 of the License, or (at your
# option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General
# Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.

"""Wrapper around ligo-skymap-plot that creates a skymap plot from a fits file.
The main purpose of this is to add metadata to the resulting png plot, so
that it can be used in an inference results page. If you don't care about that,
and want more control on the plot options, run ligo-skaymap-plot instead.

To create a fits file from an inference or posterior hdf file, see
pycbc_inference_create_fits.

This requires ligo.skymap to be installed, which requires python 3.
"""
import argparse
import sys
import subprocess
from PIL import Image, PngImagePlugin


parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--input-file', required=True,
                    help='Input fits file"')
parser.add_argument('--output-file', required=True,
                    help='Output png file."')
parser.add_argument('--colormap',
                    help='Specify the colormap to use.')

opts = parser.parse_args()

cmd = 'ligo-skymap-plot {} -o {} --annotate --contour 50 90'.format(
    opts.input_file, opts.output_file)
if opts.colormap is not None:
    cmd += ' --colormap {}'.format(opts.colormap)
ret = subprocess.run(cmd.split())
ret.check_returncode()

im = Image.open(opts.output_file)
meta = PngImagePlugin.PngInfo()

kwds = {'cmd': " ".join(sys.argv),
        'title': "Sky map",
        'caption': "Sky location posterior probability."}
for key in kwds:
    meta.add_text(str(key), str(kwds[key]))
im.save(opts.output_file, "png", pnginfo=meta)
