# -*- coding: utf-8 -*-
"""
Fujifilm F-Gamut Colourspace
============================

Defines the *Fujifilm F-Gamut* colourspace:

-   :attr:`colour.models.F_GAMUT_COLOURSPACE`.

References
----------
-   :cite:`Fujifilm2016` : Fujifilm. (2016). F-Log Data Sheet Ver.1.0 (pp.
    1-4). https://www.fujifilm.com/support/digital_cameras/software/lut/pdf/\
F-Log_DataSheet_E_Ver.1.0.pdf
"""

from __future__ import division, unicode_literals

import numpy as np

from colour.colorimetry import ILLUMINANTS
from colour.models.rgb import (RGB_Colourspace, log_encoding_FLog,
                               normalised_primary_matrix, log_decoding_FLog)

__author__ = 'Colour Developers'
__copyright__ = 'Copyright (C) 2013-2020 - Colour Developers'
__license__ = 'New BSD License - http://opensource.org/licenses/BSD-3-Clause'
__maintainer__ = 'Colour Developers'
__email__ = 'colour-developers@colour-science.org'
__status__ = 'Production'

__all__ = [
    'F_GAMUT_PRIMARIES', 'F_GAMUT_WHITEPOINT_NAME', 'F_GAMUT_WHITEPOINT',
    'F_GAMUT_TO_XYZ_MATRIX', 'XYZ_TO_F_GAMUT_MATRIX', 'F_GAMUT_COLOURSPACE'
]

F_GAMUT_PRIMARIES = np.array([
    [0.70800, 0.29200],
    [0.17000, 0.79700],
    [0.13100, 0.04600],
])
"""
*Fujifilm F-Gamut* colourspace primaries.

F_GAMUT_PRIMARIES : ndarray, (3, 2)
"""

F_GAMUT_WHITEPOINT_NAME = 'D65'
"""
*Fujifilm F-Gamut* colourspace whitepoint name.

F_GAMUT_WHITEPOINT : unicode
"""

F_GAMUT_WHITEPOINT = (ILLUMINANTS['CIE 1931 2 Degree Standard Observer'][
    F_GAMUT_WHITEPOINT_NAME])
"""
*Fujifilm F-Gamut* colourspace whitepoint.

F_GAMUT_WHITEPOINT : ndarray
"""

F_GAMUT_TO_XYZ_MATRIX = normalised_primary_matrix(F_GAMUT_PRIMARIES,
                                                  F_GAMUT_WHITEPOINT)
"""
*Fujifilm F-Gamut* colourspace to *CIE XYZ* tristimulus values matrix.

F_GAMUT_TO_XYZ_MATRIX : array_like, (3, 3)
"""

XYZ_TO_F_GAMUT_MATRIX = np.linalg.inv(F_GAMUT_TO_XYZ_MATRIX)
"""
*CIE XYZ* tristimulus values to *Fujifilm F-Gamut* colourspace matrix.

XYZ_TO_F_GAMUT_MATRIX : array_like, (3, 3)
"""

F_GAMUT_COLOURSPACE = RGB_Colourspace(
    'F-Gamut',
    F_GAMUT_PRIMARIES,
    F_GAMUT_WHITEPOINT,
    F_GAMUT_WHITEPOINT_NAME,
    F_GAMUT_TO_XYZ_MATRIX,
    XYZ_TO_F_GAMUT_MATRIX,
    log_encoding_FLog,
    log_decoding_FLog,
)
F_GAMUT_COLOURSPACE.__doc__ = """
*Fujifilm F-Gamut* colourspace.

References
----------
:cite:`Fujifilm2016`

F_GAMUT_COLOURSPACE : RGB_Colourspace
"""


#!/usr/bin/env python
"""
Satellite lensing EMCEE wrapper

"""
from __future__ import absolute_import, division, print_function

import emcee
import numpy
import os
import sys
from astropy.io.fits import BinTableHDU, Column, Header, PrimaryHDU
from itertools import count, izip
from numpy import all as npall, array, append, concatenate, dot, inf, isnan
from numpy import isfinite, log, log10, outer, pi, sqrt, transpose, zeros
from os import remove
from os.path import isfile
from time import ctime
import pickle

from . import sampling_utils

def run_emcee(hm_options, sampling_options, args):
    # load halo model setup
    function, params, param_types, prior_types, \
        val1, val2, val3, val4, params_join, hm_functions, \
        starting, meta_names, fits_format = hm_options
    # load MCMC sampler setup
    datafile, datacols, covfile, covcols, exclude_bins, output, \
        sampler, nwalkers, nsteps, nburn, \
        thin, k, threads, sampler_type, update_freq = sampling_options

    #function = cloud.serialization.cloudpickle.dumps(model)
    #del model
    #print(function)

    #pickle.dumps(function)
    #print('pickled')

    if args.demo:
        print(' ** Running demo only **')
    elif isfile(output):
        msg = 'Warning: output file %s exists. Overwrite? [y/N] ' %output
        answer = raw_input(msg)
        if len(answer) == 0:
            exit()
        if answer.lower() not in ('y', 'yes'):
            exit()
    if not args.demo:
        print('Started -', ctime())

    #load data files
    Ndatafiles = len(datafile)
    R, esd = sampling_utils.load_datapoints(datafile, datacols, exclude_bins)
    Nobsbins, Nrbins = esd.shape
    rng_obsbins = xrange(Nobsbins)
    rng_rbins = xrange(Nrbins)
    # load covariance
    cov = sampling_utils.load_covariance(covfile, covcols,
                                         Nobsbins, Nrbins, exclude_bins)
    cov, icov, likenorm, esd_err, cov2d = cov

    # needed for offset central profile
    R, Rrange = sampling_utils.setup_integrand(R, k)
    angles = numpy.linspace(0, 2*pi, 540)
    val1 = numpy.append(val1, [Rrange, angles])

    # identify fixed and free parameters
    jfixed = (prior_types == 'fixed') | (prior_types == 'read') | \
             (prior_types == 'function')
    jfree = ~jfixed
    ndim = len(val1[(jfree)])
    if len(starting) != ndim:
        msg = 'ERROR: Not all starting points defined for free parameters.'
        print(msg)
        exit()
    print('starting =', starting)

    # identify the function. Raises an AttributeError if not found
    #function = model.model()
    #sat_profile = params.sat_profile()
    #group_profile = params.group_profile()
    #function = model

    if not args.demo:
        hdrfile = '.'.join(output.split('.')[:-1]) + '.hdr'
        print('Printing header information to', hdrfile)
        hdr = open(hdrfile, 'w')
        print('Started', ctime(), file=hdr)
        print('datafile', ','.join(datafile), file=hdr)
        print('cols', ','.join([str(c) for c in datacols]), file=hdr)
        print('covfile', covfile, file=hdr)
        print('covcols', ','.join([str(c) for c in covcols]), file=hdr)
        if exclude_bins is not None:
            print('exclude_bins', ','.join([str(c) for c in exclude_bins]),
                  file=hdr)
        print('model %s' %function, file=hdr)
        for p, pt, v1, v2, v3, v4 in izip(params, prior_types,
                                        val1, val2, val3, val4):
            try:
                line = '%s  %s  ' %(p, pt)
                line += ','.join(numpy.array(v1, dtype=str))
            except TypeError:
                line = '%s  %s  %s  %s  %s  %s' \
                    %(p, pt, str(v1), str(v2), str(v3), str(v4))
            print(line, file=hdr)
        print('nwalkers  {0:5d}'.format(nwalkers), file=hdr)
        print('nsteps    {0:5d}'.format(nsteps), file=hdr)
        print('nburn     {0:5d}'.format(nburn), file=hdr)
        print('thin      {0:5d}'.format(thin), file=hdr)
        hdr.close()

    # are we just running a demo?
    if args.demo:
        import pylab
        from matplotlib import cm
        def plot_demo(ax, Ri, gt, gt_err, f, fsat, fhost):
            Ri = Ri[1:]
            ax.errorbar(Ri, gt, yerr=gt_err, fmt='ko', ms=10)
            ax.plot(Ri, f, 'r-', lw=3)
            ax.plot(Ri, fsat, 'b--', lw=2)
            ax.plot(Ri, fhost, 'g-.', lw=2)
            ax.set_xscale('log')
            for x, fi, gti, gei in izip(Ri, f, gt, gt_err):
                ax.annotate('{0:.2f}'.format((fi-gti)/gei),
                            xy=(x,gti+20), ha='center', va='bottom',
                            color='r')
            return
        val1[jfree] = starting
        if params_join is not None:
            v1 = list(val1)
            for p in params_join:
                # without this list comprehension numpy can't keep track of the
                # data type. I believe this is because there are elements of
                # different types in val1 and therefore its type is not
                # well defined (so it gets "object")
                v1[p[0]] = array([val1[pj] for pj in p])
            # need to delete elements backwards to preserve indices
            aux = [[v1.pop(pj) for pj in p[1:][::-1]]
                   for p in params_join[::-1]]
            val1 = v1 #array(v1) ??
        model = function(val1, R)
        residuals = esd - model[0]
        dof = esd.size - starting.size - 1
        chi2 = array([dot(residuals[m], dot(icov[m][n], residuals[n]))
                      for m in rng_obsbins for n in rng_obsbins]).sum()
        print(' ** chi2 = %.2f/%d **' %(chi2, dof))
        fig, axes = pylab.subplots(figsize=(4*Ndatafiles,4), ncols=Ndatafiles)
        if Ndatafiles == 1:
            plot_demo(axes, R, esd, esd_err, model[0], model[1], model[2])
        else:
            for i in izip(axes, R, esd, esd_err, model[0], model[1], model[2]):
                plot_demo(*i)
        if npall(esd - esd_err > 0):
            for ax in axes:
                ax.set_yscale('log')
        fig.tight_layout(w_pad=0.01)
        pylab.show()
        fig, axes = pylab.subplots(figsize=(8,8), nrows=cov.shape[0],
                                   ncols=cov.shape[0])
        for m, axm in enumerate(axes):
            for n, axmn in enumerate(axm):
                axmn.imshow(cov[m][-n-1][::-1], interpolation='nearest',
                            cmap=cm.CMRmap_r)
        fig.tight_layout()
        pylab.show()
        exit()

    # set up starting point for all walkers
    po = starting * numpy.random.uniform(0.99, 1.01, size=(nwalkers,ndim))
    lnprior = zeros(ndim)
    mshape = meta_names.shape
    # this assumes that all parameters are floats -- can't imagine a
    # different scenario
    metadata = [[] for m in meta_names]
    for j in xrange(len(metadata)):
        for f in fits_format[j]:
            if len(f) == 1:
                metadata[j].append(zeros(nwalkers*nsteps/thin))
            else:
                size = [nwalkers*nsteps/thin, int(f[:-1])]
                # only for ESDs. Note that there will be trouble if outputs
                # other than the ESD have the same length, so avoid them at
                # all cost.
                if exclude_bins is not None \
                    and size[1] == esd.shape[-1]+len(exclude_bins):
                    size[1] -= len(exclude_bins)
                metadata[j].append(zeros(size))
    metadata = [array(m) for m in metadata]
    fail_value = []
    for m in metadata:
        shape = list(m.shape)
        shape.remove(max(shape))
        fail_value.append(zeros(shape))
    # the last numbers are data chi2, lnLdata, lnPderived
    for i in xrange(4):
        fail_value.append(9999)

    sampler = emcee.EnsembleSampler(nwalkers, ndim, lnprob,
                                    threads=threads,
                                    args=(R,esd,icov,function,
                                          params,prior_types[jfree],
                                          val1,val2,val3,val4,params_join,
                                          jfree,lnprior,likenorm,
                                          rng_obsbins,fail_value,
                                          array,dot,inf,izip,outer,pi))
                                          #isfinite,log,log10
                                          #outer,sqrt,zeros))
    # burn-in
    if nburn > 0:
        pos, prob, state, blobs = sampler.run_mcmc(po, nburn)
        sampler.reset()
        print('{0} Burn-in steps finished ({1})'.format(nburn, ctime()))
    else:
        pos = po
    # incrementally save output
    chi2 = [zeros(nwalkers*nsteps/thin) for i in xrange(4)]
    nwritten = 0
    for i, result in enumerate(sampler.sample(pos, iterations=nsteps,
                                              thin=thin)):
        # make sure that nwalkers is a factor of this number!
        if i*nwalkers % update_freq == nwalkers:
            out = write_to_fits(output, chi2, sampler, nwalkers, thin,
                                params, jfree, metadata, meta_names, i,
                                nwritten, Nobsbins,
                                array, BinTableHDU, Column, ctime, enumerate,
                                isfile, izip, transpose, xrange)
            metadata, nwriten = out

    hdr = open(hdrfile, 'a')
    try:
        print('acceptance_fraction =', sampler.acceptance_fraction)
        print('acceptance_fraction =', file=hdr, end=' ')
        for af in sampler.acceptance_fraction:
            print(af, file=hdr, end=' ')
    except ImportError:
        pass
    try:
        print('acor =', sampler.acor)
        print('\nacor =', file=hdr, end=' ')
        for ac in sampler.acor:
            print(ac, file=hdr, end=' ')
    except ImportError:
        pass
    try:
        print('acor_time =', sampler.get_autocorr_time())
        print('\nacor_time =', file=hdr, end=' ')
        for act in sampler.get_autocorr_time():
            print(act, file=hdr, end=' ')
    except AttributeError:
        pass
    print('\nFinished', ctime(), file=hdr)
    hdr.close()
    print('Saved to', hdrfile)

    cmd = 'mv {0} {1}'.format(output, output.replace('.fits', '.temp.fits'))
    print(cmd)
    os.system(cmd)
    print('Saving everything to {0}...'.format(output))
    print(i, nwalkers, nwritten)
    write_to_fits(output, chi2, sampler, nwalkers, thin,
                  params, jfree, metadata, meta_names, i+1,
                  nwritten, Nobsbins,
                  array, BinTableHDU, Column, ctime, enumerate,
                  isfile, izip, transpose, xrange)
    os.remove(output.replace('.fits', '.temp.fits'))
    print('Everything saved to {0}!'.format(output))
    return

def lnprob(theta, R, esd, icov, function, params, prior_types,
           val1, val2, val3, val4, params_join, jfree, lnprior, likenorm,
           rng_obsbins, fail_value, array, dot, inf, izip, outer, pi):
           #array, dot, inf, izip, isfinite, log, log10, sqrt):
    """
    Probability of a model given the data, i.e., log-likelihood of the data
    given a model, times the prior on the model parameters.

    Parameters
    ----------
        theta
            whatever *free* parameters are received by the model selected.
        R
            projected distances from the satellite
        esd
            Excess surface density at distances R
        esd_err
            Uncertainties on the ESD
        function
            the model used to calculate the likelihood
        prior_types
            one value per parameter in *theta*, {'normal', 'uniform', 'fixed'}
        val1
            depends on each prior_type:
                -normal : the mean of the gaussian
                -uniform : the minimum allowed value
                -fixed : the fixed value
        val2
            depends on each prior_type:
                -normal : the half-width of the gaussian
                -uniform : the maximum allowed value
                -fixed : ignored (but must be there)
        jfree
            indices of the free values
        lnprior
            just a placeholder, should be an array with a length
            equal to theta, so that it doesn't have to be defined every
            time

    """
    _log = log
    v1free = val1[jfree]
    v2free = val2[jfree]
    v3free = val3[jfree]
    v4free = val4[jfree]
    if not isfinite(v1free.sum()):
        return -inf, fail_value
    # satellites cannot be more massive than the group!
    #if theta[params[jfree] == 'Msat'] >= theta[params[jfree] == 'Mgroup']:
        #return -inf, fail_value
    # not normalized yet
    j = (prior_types == 'normal')
    lnprior[j] = array([-(v-v1)**2 / (2*v2**2) - _log(2*pi*v2**2)/2
                        if v3 <= v <= v4 else -inf
                        for v, v1, v2, v3, v4
                        in izip(theta[j], v1free[j], v2free[j],
                                v3free[j], v4free[j])])
    j = (prior_types == 'lognormal')
    lnprior[j] = array([-(log10(v)-v1)**2 / (2*v2**2) - _log(2*pi*v2**2)/2
                        if v3 <= v <= v4 else -inf
                        for v, v1, v2, v3, v4
                        in izip(theta[j], v1free[j], v2free[j],
                                v3free[j], v4free[j])])
    j = (prior_types == 'uniform')
    lnprior[j] = array([0-_log(v2-v1) if v1 <= v <= v2 else -inf
                        for v, v1, v2
                        in izip(theta[j], v1free[j], v2free[j])])
    # note that exp is not normalized
    j = (prior_types == 'exp')
    lnprior[j] = array([v**v1 if v2 <= v <= v3 else -inf
                        for v, v1, v2, v3
                        in izip(theta[j], v1free[j], v2free[j], v3free[j])])
    lnprior_total = lnprior.sum()
    if not isfinite(lnprior_total):
        return -inf, fail_value
    # all other types ('fixed', 'read') should not contribute to the prior
    # run the given model
    v1 = val1.copy()
    v1[jfree] = theta
    if params_join is not None:
        v1j = list(v1)
        for p in params_join:
            # without this list comprehension numpy can't keep track of the
            # data type. I believe this is because there are elements of
            # different types in val1 and therefore its type is not
            # well defined (so it gets "object")
            v1j[p[0]] = array([v1[pi] for pi in p])
        # need to delete elements backwards to preserve indices
        aux = [[v1j.pop(pi) for pi in p[1:][::-1]]
                for p in params_join[::-1]]
        v1 = v1j #array(v1j) ??

    model = function(v1, R)
    # no covariance
    #chi2 = (((esd-model[0]) / esd_err) ** 2).sum()
    # full covariance included
    residuals = esd - model[0]
    chi2 = array([dot(residuals[m], dot(icov[m][n], residuals[n]))
                  for m in rng_obsbins for n in rng_obsbins]).sum()
    if not isfinite(chi2):
        return -inf, fail_value
    # remember that the last value returned by the models must be a lnprior
    # from the derived parameters
    lnlike = -chi2/2. + likenorm
    model.append(lnprior_total)
    model.append(chi2)
    model.append(lnlike)
    return lnlike + model[-3] + lnprior_total, model

def write_to_fits(output, chi2, sampler, nwalkers, thin, params, jfree,
                  metadata, meta_names, iternum, nwritten,
                  Nobsbins, array, BinTableHDU, Column, ctime, enumerate,
                  isfile, izip, transpose, xrange):
    nexclude = len(chi2)
    lnprior, lnPderived, chi2, lnlike = chi2
    if isfile(output):
        remove(output)
    chain = transpose(sampler.chain, axes=(2,1,0))
    columns = [Column(name=param, format='E', array=data[:iternum].flatten())
               for param, data in izip(params[jfree], chain)]
    columns.append(Column(name='lnprob', format='E',
                          array=sampler.lnprobability.T[:iternum].flatten()))
    if len(meta_names) > 0:
        # save only the last chunk (starting from t),
        # all others are already in metadata.
        # NOTE that this is only implemented for a model with the
        # same format as fiducial()
        for j, blob in izip(xrange(nwritten, iternum),
                            sampler.blobs[nwritten:]):
            data = [transpose([b[i] for b in blob])
                    for i in xrange(len(blob[0])-nexclude)]
            # re-arrange blobs
            if Nobsbins == 1:
                for i in xrange(len(data)):
                    if len(data[i].shape) == 2:
                        data[i] = array([b[i] for b in blob])
            else:
                for i in xrange(len(data)):
                    if len(data[i].shape) == 3:
                        data[i] = transpose([b[i] for b in blob],
                                            axes=(1,0,2))
            # store data
            for k in xrange(len(data)):
                for i in xrange(len(data[k])):
                    metadata[k][i][j*nwalkers:(j+1)*nwalkers] = data[k][i]
            lnPderived[j*nwalkers:(j+1)*nwalkers] = array([b[-4]
                                                           for b in blob])
            lnprior[j*nwalkers:(j+1)*nwalkers] = array([b[-3] for b in blob])
            chi2[j*nwalkers:(j+1)*nwalkers] = array([b[-2] for b in blob])
            lnlike[j*nwalkers:(j+1)*nwalkers] = array([b[-1] for b in blob])
        columns.append(Column(name='lnprior', format='E', array=lnprior))
        columns.append(Column(name='lnPderived', format='E',
                              array=lnPderived))
        columns.append(Column(name='chi2', format='E', array=chi2))
        columns.append(Column(name='lnlike', format='E', array=lnlike))
        # this handles exclude_bins properly
        for name, val in izip(meta_names, metadata):
            for name_i, val_i in izip(name, val):
                try:
                    fmt = '{0}E'.format(val_i.shape[1])
                except IndexError:
                    fmt = 'E'
                columns.append(Column(name=name_i, array=val_i, format=fmt))
        nwritten = iternum * nwalkers
    fitstbl = BinTableHDU.from_columns(columns)
    fitstbl.writeto(output)
    print('Saved to {0} with {1} samples'.format(output, iternum*nwalkers),
          end=' ')
    if thin > 1:
        print('(printing every {0}th sample)'.format(thin), end=' ')
    print('- {0}'.format(ctime()))
    return metadata, nwritten


import theano
import theano.tensor as T

from nn.layers.core import Unit
from nn.utils import logsumexp

class SeqLabelAlg(Unit):
    def __init__(self, name='SeqLabelModel'):
        super(SeqLabelAlg, self).__init__(name=name)

    def viterbi(self, emit_scores, trans_scores):
        """
        :param emit_scores: 1D: n_words, 2D: batch_size, 3D: n_labels
        :param trans_scores: 1D: n_words, 2D: n_labels
        :return: 1D: n_words; 2D: batch_size, elem=label id
        """
        [scores, labels], _ = theano.scan(fn=self._viterbi_forward,
                                          sequences=[emit_scores[1:]],
                                          outputs_info=[emit_scores[0], None],
                                          non_sequences=trans_scores)

        label_max_last = T.argmax(scores[-1], axis=1)
        labels_max, _ = theano.scan(fn=self._viterbi_backward,
                                    sequences=labels[::-1],
                                    outputs_info=label_max_last)

        y = T.zeros(shape=(emit_scores.shape[0], emit_scores.shape[1]), dtype='int32')
        y = T.set_subtensor(y[-1], label_max_last)
        y = T.set_subtensor(y[:-1], labels_max[::-1])
        return y

    @staticmethod
    def _viterbi_forward(e_t, score_prev, trans):
        """
        :param e_t: 1D: batch_size, 2D: n_labels
        :param score_prev: 1D: batch_size, 2D: n_labels
        :param trans: 1D: n_labels, 2D, n_labels
        :return: max_scores_t: 1D: batch_size, 2D: n_labels
        :return: max_labels_t: 1D: batch_size, 2D: n_labels
        """
        score = score_prev.dimshuffle(0, 'x', 1) + trans + e_t.dimshuffle(0, 1, 'x')
        max_scores_t, max_labels_t = T.max_and_argmax(score, axis=2)
        return max_scores_t, max_labels_t

    @staticmethod
    def _viterbi_backward(labels_t, label_max):
        """
        :param labels_t: 1D: batch_size, 2D: n_labels; elem=label id
        :param label_max: 1D: batch_size; elem=label id
        :return: 1D: batch_size; elem=label id
        """
        return labels_t[T.arange(labels_t.shape[0]), label_max]

class CRF(SeqLabelAlg):
    def __init__(self,
                 input_dim,
                 output_dim,
                 use_bias=True,
                 weight_init='xavier',
                 bias_init='zero'):
        super(CRF, self).__init__(name='CRF(%dx%d)' % (input_dim, output_dim))
        self.W = self._set_param(shape=(input_dim, output_dim),
                                 init_type=weight_init,
                                 name='W_crf')
        self.W_t = self._set_param(shape=(output_dim, output_dim),
                                   init_type=weight_init,
                                   name='W_tran_crf')

        if use_bias:
            self.b = self._set_param(shape=output_dim,
                                     init_type=bias_init,
                                     name='b_crf')
            self.params = [self.W, self.W_t, self.b]
        else:
            self.b = None
            self.params = [self.W, self.W_t]

    def forward(self, x):
        emit_scores = T.dot(x, self.W)
        if self.b:
            emit_scores = emit_scores + self.b
        return emit_scores

    def get_y_proba(self, emit_scores, y_true):
        """
        :param emit_scores: 1D: n_words, 2D: batch_size, 3D: n_labels
        :param y_true: 1D: n_words, 2D: batch_size
        :return: 1D: batch_size; elem=log probability
        """
        # 1D: batch_size, 2D: n_labels
        z_score0 = emit_scores[0]
        # 1D: batch_size; elem=path score
        y_score0 = z_score0[T.arange(z_score0.shape[0]), y_true[0]]

        inputs = [emit_scores[1:], y_true[1:]]
        [_, y_scores, z_scores], _ = theano.scan(fn=self._forward_step,
                                                 sequences=inputs,
                                                 outputs_info=[y_true[0], y_score0, z_score0],
                                                 non_sequences=self.W_t)

        y_score = y_scores[-1]
        z_score = logsumexp(z_scores[-1], axis=1).flatten()

        return y_score - z_score

    @staticmethod
    def _forward_step(h_t, y_t, y_prev, y_score_prev, z_score_prev, trans):
        """
        :param h_t: 1D: batch_size, 2D: n_labels
        :param y_t: 1D: batch_size
        :param y_prev: 1D: batch_size
        :param y_score_prev: 1D: batch_size
        :param z_score_prev: 1D: batch_size, 2D: n_labels
        :param trans: 1D: n_labels, 2D, n_labels
        """
        # 1D: batch_size
        y_score_t = y_score_prev + trans[y_t, y_prev] + h_t[T.arange(h_t.shape[0]), y_t]
        # 1D: batch_size, 2D: n_labels, 3D: n_labels
        z_sum = z_score_prev.dimshuffle(0, 'x', 1) + trans
        # 1D: batch_size, 2D: n_labels
        z_score_t = logsumexp(z_sum, axis=2).reshape(h_t.shape) + h_t
        return y_t, y_score_t, z_score_t

    def get_y_pred(self, emit_scores):
        """
        :param emit_scores: 1D: n_words, 2D: batch_size, 3D: n_labels
        :return: 1D: batch_size, 2D: n_words; elem=label id
        """
        return self.viterbi(emit_scores=emit_scores, trans_scores=self.W_t).dimshuffle(1, 0)


# -*- coding: utf-8 -*-
import numpy as np
from pyGRBz.utils import mag2Jy, convAB
from astropy.table import Table, Column
from pyGRBz.extinction_correction import correct_MW_ext
from pyGRBz.io_grb import load_telescope_transmissions
import imp

def load_sys_response(data, wavelength, path):
    """ Load the system throuput curves for each filter in the data

    Returns
    -------
    sys_rep: astropy Table
    """
    # works for constant wvl_step
    dwvl = np.gradient(wavelength)

    sys_res = []
    tel_name = []
    tel_band = []
    wvl_eff = []
    width = []
    zp = []

    # Sort the telescope used
    for tel in data.group_by(["telescope", "band"]).groups.keys:
        # Import the filter throughput curve only once if filter used several
        # times (for a light curve for instance)
        tel_name.append(tel["telescope"])
        tel_band.append(tel["band"])

        # Import the throughput curve
        filter_trans = load_telescope_transmissions(
            {"telescope": tel["telescope"], "band": tel["band"], "path": path},
            wavelength,
        )
        sys_res.append(filter_trans)

        # calculate the effective wavelength
        a = np.trapz(wavelength * filter_trans, wavelength)
        b = np.trapz(filter_trans, wavelength)
        wvl_eff.append(a / b)

        # Calculate the width of the band
        mask = filter_trans > 0.05 * max(filter_trans)

        width.append(wavelength[mask][-1] - wavelength[mask][0])

        # print (tel['telescope'], tel['band'], a/b, width[-1])
        # Not sure it is used anymore ... to be checked. Formula also to check
        zp.append(2.5 * np.log10(np.sum(filter_trans * dwvl, axis=0)) + 23.9)

    sys_res_table = Table(
        [tel_name, tel_band, wvl_eff, width, sys_res, zp],
        names=["telescope", "band", "wvl_eff", "band_width",
               "sys_response", "zeropoint"])

    # Sort the table by telescope names and ascending eff. wavelength
    sys_res_table.sort(["telescope", "wvl_eff"])
    return sys_res_table

def formatting_data(data, system_response, grb_info, wavelength,
                    dustrecalib="yes"):
    """ """
    try:
        _, path_dust_map, _ = imp.find_module("pyGRBaglow")
    except:
        print("path to pyGRBaglow can not be found.")

    dustmapdir = path_dust_map + "/galactic_dust_maps"

    #  Add filter info to data (throughut curve,eff. wvl and width)
    col_band_width = Column(name="band_width", data=np.zeros(len(data)))
    col_band_effwvl = Column(name="eff_wvl", data=np.zeros(len(data)))
    col_band_zp = Column(name="zeropoint", data=np.zeros(len(data)))
    col_band_sysres = Column(
        name="sys_response", data=np.zeros((len(data), len(wavelength)))
    )
    data.add_columns([col_band_effwvl, col_band_width, col_band_zp,
                      col_band_sysres])

    for table in data.group_by(["telescope", "band"]).groups.keys:
        # print (table)
        mask1 = data["telescope"] == table["telescope"]
        mask1[mask1 == True] = data[mask1]["band"] == table["band"]
        mask2 = system_response["telescope"] == table["telescope"]
        mask2[mask2 == True] = system_response[mask2]["band"] == table["band"]
        # print (system_response[mask3][mask4]['sys_response'])
        # print (system_response[mask3][mask4]['sys_response'][0])

        width = []
        effwvl = []
        zp = []
        sys_res = []
        for i in range(np.sum(mask2)):
            width.append(system_response[mask2]["band_width"][0])
            effwvl.append(system_response[mask2]["wvl_eff"][0])
            zp.append(system_response[mask2]["zeropoint"][0])
            sys_res.append(system_response[mask2]["sys_response"][0])
        data["band_width"][mask1] = width
        data["eff_wvl"][mask1] = effwvl
        data["zeropoint"][mask1] = zp
        data["sys_response"][mask1] = sys_res

    # Convert vega magnitudes in AB if needed
    mask1 = data["phot_sys"] == "vega"
    if mask1.any():
        # print ('some vega')

        #  If a Vega-AB correction is present in the file use this value

        if "ABcorr" in data.colnames:
            mask2 = (data["phot_sys"] == "vega") & (~data["ABcorr"].mask)
            if mask2.any():
                # print ('AB corr')
                for table in data[mask2]:
                    mask3 = mask2.copy()

                    mask3[mask3 == True] = data[mask3]["Name"] == table["Name"]
                    mask3[mask3 == True] = (
                        data[mask3]["telescope"] == table["telescope"]
                    )
                    mask3[mask3 == True] = data[mask3]["band"] == table["band"]
                    mask3[mask3 == True] = (
                        data[mask3]["time_since_burst"] == table["time_since_burst"]
                    )

                    newABmag = table["mag"] + table["ABcorr"]
                    photsys = "AB"
                    # substitute the vega magnitudes by AB ones
                    data["mag"][mask3] = newABmag
                    data["phot_sys"][mask3] = photsys

            #  When no AB correction is given in input file, compute it
            mask2 = (data["phot_sys"] == "vega") & (data["ABcorr"].mask)
            if mask2.any():

                # print ('convAB')
                for table in data[mask2]:
                    mask3 = mask2.copy()

                    mask3[mask3 == True] = data[mask3]["Name"] == table["Name"]
                    mask3[mask3 == True] = (
                        data[mask3]["telescope"] == table["telescope"]
                    )
                    mask3[mask3 == True] = data[mask3]["band"] == table["band"]
                    mask3[mask3 == True] = (
                        data[mask3]["time_since_burst"] == table["time_since_burst"]
                    )

                    newABmag = table["mag"] + convAB(wavelength, table["sys_response"])
                    photsys = "AB"
                    # substitute the vega magnitudes by AB ones
                    data["mag"][mask3] = newABmag
                    data["phot_sys"][mask3] = photsys

        else:
            # print ('convAB')
            for table in data[mask1]:
                mask3 = mask1.copy()

                mask3[mask3 == True] = data[mask3]["Name"] == table["Name"]
                mask3[mask3 == True] = data[mask3]["telescope"] == table["telescope"]
                mask3[mask3 == True] = data[mask3]["band"] == table["band"]
                mask3[mask3 == True] = (
                    data[mask3]["time_since_burst"] == table["time_since_burst"]
                )

                newABmag = table["mag"] + convAB(wavelength,
                                                 table["sys_response"])
                photsys = "AB"
                # substitute the vega magnitudes by AB ones
                data["mag"][mask3] = newABmag
                data["phot_sys"][mask3] = photsys

    # Correct for galactic extinction
    data = correct_MW_ext(data, grb_info, wavelength,
                          dustmapdir=dustmapdir,
                          recalibration=dustrecalib)

    # Add Flux to the seds
    convert_dict = {"photometry_system": "AB"}

    flux = mag2Jy(convert_dict, data["mag"] - data["ext_mag"]) * 1e6
    flux_err = np.array(abs(flux * -0.4 * np.log(10) * data["mag_err"]))
    mask = data["detection"] == -1
    if mask.any():
        flux_err[mask] = flux[mask] / 2.0

    col_flux = Column(name="flux", data=flux, unit="microJy")
    col_flux_err = Column(name="flux_err", data=flux_err, unit="microJy")

    data.add_columns([col_flux, col_flux_err])

    data.sort(["Name", "eff_wvl"])
    return data


"""
Object to Object relationship models
"""
import numpy as np
from scipy.spatial.distance import euclidean
import ipdb
from collections import OrderedDict

# Params
OBJ_W_SIZE = 0.02  # Object width size
NEAR_DIST_TH = 0.1  # Near distance threshold
HOLD_DIST_TH = 0.02 # Holding distance threshold
IN_DIST_TH = 0.01 # Object in target distance threshold

class SpatialObject(object):

    OBJ_CLASSES = OrderedDict([
        ("eef", 0),
        ("can", 1),
        ("milk", 2),
        ("bread", 3),
        ("cereal", 4),
        ("roundnut", 5),
        ("squarenut", 6)
    ])

    def __init__(self, name, pos, orient, is_robot=False):
        """ Constructor

            Args
            -----
            name: str
                Object name
            pos: np.ndarray
                Position vector (3,1)
            orient: np.ndarray
                Orientation Quaternion  (4,1)
        """
        self.name = name
        self.pos = pos
        self.orient = orient
        self.is_robot = is_robot
        self.encode_onehot()

    def encode_onehot(self):
        """ Encode into One-hot vector
        """
        self.encoding = np.array([
                1. if self.name == _c else 0. \
                    for _c in self.OBJ_CLASSES
            ])

    def __repr__(self):
        return "<" + self.name + ">"

class Relation(object):

    def __init__(self, name, validator):
        """ Constructor

            Args
            -----
            name : str
                Name of the relation
            validator : function | lambda
                Validation function
        """
        self._name = name
        self._f_validator = validator

    @property
    def name(self):
        """ Name getter
        """
        return self._name

    @property
    def validator(self):
        """ Validator Function getter
        """
        return self._f_validator

    def __call__(self, *args):
        """ Relation validation call

            Returns
            -----
            int (0 or 1)
                Returns weather the relation applies or not
        """
        return self._f_validator(*args)

class O2ORelation(object):

    # Object to object relations
    O2O_RELATIONS = OrderedDict([
        ("on", "is_on"),
        ("in", "is_in"),
        ("front", "is_front"),
        ("behind", "is_behind"),
        ("right_side", "is_right"),
        ("left_side", "is_left"),
        ("above", "is_above"),
        ("below", "is_below"),
        ("holding", "is_holding"),
        ("near", "is_near")
    ])
    # o2o_index = {_o: _j for _j, _o in enumerate(O2O_RELATIONS)}

    def __init__(self, *args):
        """ Constructor
        """
        self._rep = [] # Representation vector
        self.objs = args  # Objects to compute the relationship between
        # initializes all possible relations
        for _n, _func in self.O2O_RELATIONS.items():
            self.__dict__[_n] = Relation(
                _n,
                getattr(self, _func)
            )
        # Compute rules
        self.evaluate()

    @property
    def rep(self):
        """ Relations getter
        """
        return self._rep

    def decode(self):
        """ Decode relationships for human interpretation
        """
        o_reps = [_o.__repr__() for _o  in self.objs]
        return [("_"
                + list(self.O2O_RELATIONS.keys())[_i].upper()
                + "_").join(o_reps) \
            for _i, _r in enumerate(self._rep) if _r
        ]

    def __repr__(self):
        return '<' + ','.join(self.decode()) \
                +  ":" + super().__repr__().split('.')[-1]

    def evaluate(self):
        """ Evaluate over all registered relations,
            set which of them are valid with 1 and 0 otherwise,
            to return a one-hot vector.
        """
        # Call all rules
        _evals = [self.__dict__[_r]() \
            for _r in self.O2O_RELATIONS.keys()]
        # Set the relations representation
        self._rep = np.array(_evals)

    @staticmethod
    def zeros():
        """ Return zero vector with
            O2O Relations cardinality

            Returns
            -----
            np.ndarray
                Zero Vector
        """
        return np.zeros((len(O2ORelation.O2O_RELATIONS),))

    def default(self):
        """ Default Rule function
        """
        # print("Executing DEFAULT..")
        return 0.

    def is_on(self):
        """ ON rule validator: returns 1 in case
            the 1st object has greater Z and
            the X and Y absolute difference is less or equal
            to the OBJ_W_SIZE, otherwise returns 0
        """
        # print("Executing IS_ON..")
        if self.objs[0].pos[2] > self.objs[1].pos[2]:
            _xdif = abs(self.objs[0].pos[0] - self.objs[1].pos[0])
            _ydif = abs(self.objs[0].pos[1] - self.objs[1].pos[1])
            if (_xdif <= OBJ_W_SIZE) and (_ydif <= OBJ_W_SIZE):
                return 1.
        return 0.

    def is_below(self):
        """ BELOW rule validator: 1 if the 1st object
            has lower Z than the 2nd one, otherwise returns 0
        """
        # print("Executing IS_BELOW..")
        if self.objs[0].pos[2] < self.objs[1].pos[2]:
            return 1.
        return 0.

    def is_above(self):
        """ ABOVE rule validator: 1 if the 1st object
            has greater Z than the 2nd one, otherwise returns 0
        """
        # print("Executing IS_ABOVE..")
        if self.objs[0].pos[2] > self.objs[1].pos[2]:
            return 1.
        return 0.

    def is_front(self):
        """ FRONT rule validator: 1 if the 1st object
            has greater X than the 2nd one, otherwise returns 0
        """
        # print("Executing IS_FRONT..")
        if self.objs[0].pos[0] > self.objs[1].pos[0]:
            return 1.
        return 0.

    def is_behind(self):
        """ BEHING rule validator: 1 if the 1st object
            has lower X than the 2nd one, otherwise returns 0
        """
        # print("Executing IS_BEHIND..")
        if self.objs[0].pos[0] < self.objs[1].pos[0]:
            return 1.
        return 0.

    def is_right(self):
        """ RIGHT rule validator: 1 if the 1st object
            has greater Y than the 2nd one, otherwise returns 0
        """
        # print("Executing IS_RIGHT..")
        if self.objs[0].pos[1] > self.objs[1].pos[1]:
            return 1.
        return 0.

    def is_left(self):
        """ LEFT rule validator: 1 if the 1st object
            has lower Y than the 2nd one, otherwise returns 0
        """
        # print("Executing IS_LEFT..")
        if self.objs[0].pos[1] < self.objs[1].pos[1]:
            return 1.
        return 0.

    def is_near(self):
        """ NEAR rule validator: 1 if the euclidean distance
            between the 2 objects is less or equal to the
            NEAR_DIST_TH and greater than HOLD_DIST_TH,
            otherwise 0
        """
        # print("Executing IS_NEAR..")
        _dist = euclidean(self.objs[0].pos, self.objs[1].pos)
        if  _dist <= NEAR_DIST_TH and _dist > HOLD_DIST_TH:
            return 1.
        return 0.

    def is_holding(self):
        """ HOLD rule validator: 1 if the euclidean distance
            between the 2 objects is less or equal to the
            HOLD_DIST_TH, otherwise 0
        """
        # print("Executing IS_HOLDING..")
        if not self.objs[0].is_robot:
            return 0.
        _dist = euclidean(self.objs[0].pos, self.objs[1].pos)
        if  _dist <= HOLD_DIST_TH:
            return 1.
        return 0.

    def is_in(self, ):
        """ IN rule validator: 1 if the euclidean distance
            between the 1st object and the 2nd one is less
            or equal to the IN_DIST_TH
        """
        # print("Executing IS_IN..")
        _dist = euclidean(self.objs[0].pos, self.objs[1].pos)
        if  _dist <= IN_DIST_TH:
            return 1.
        return 0.


"""
Utilities for making finder charts and overlay images for ALMA proposing
"""
import string

import numpy as np

from astropy import wcs
from astropy import log
from astropy import units as u
from astropy.io import fits

from astroquery.skyview import SkyView
from astroquery.alma import Alma

def pyregion_subset(region, data, mywcs):
    """
    Return a subset of an image (`data`) given a region.

    Parameters
    ----------
    region : `pyregion.parser_helper.Shape`
        A Shape from a pyregion-parsed region file
    data : np.ndarray
        An array with shape described by WCS
    mywcs : `astropy.wcs.WCS`
        A world coordinate system describing the data
    """
    import pyregion

    shapelist = pyregion.ShapeList([region])
    if shapelist[0].coord_format not in ('physical','image'):
        # Requires astropy >0.4...
        # pixel_regions = shapelist.as_imagecoord(self.wcs.celestial.to_header())
        # convert the regions to image (pixel) coordinates
        celhdr = mywcs.sub([wcs.WCSSUB_CELESTIAL]).to_header()
        pixel_regions = shapelist.as_imagecoord(celhdr)
    else:
        # For this to work, we'd need to change the reference pixel after cropping.
        # Alternatively, we can just make the full-sized mask... todo....
        raise NotImplementedError("Can't use non-celestial coordinates with regions.")
        pixel_regions = shapelist

    # This is a hack to use mpl to determine the outer bounds of the regions
    # (but it's a legit hack - pyregion needs a major internal refactor
    # before we can approach this any other way, I think -AG)
    mpl_objs = pixel_regions.get_mpl_patches_texts()[0]

    # Find the minimal enclosing box containing all of the regions
    # (this will speed up the mask creation below)
    extent = mpl_objs[0].get_extents()
    xlo, ylo = extent.min
    xhi, yhi = extent.max
    all_extents = [obj.get_extents() for obj in mpl_objs]
    for ext in all_extents:
        xlo = xlo if xlo < ext.min[0] else ext.min[0]
        ylo = ylo if ylo < ext.min[1] else ext.min[1]
        xhi = xhi if xhi > ext.max[0] else ext.max[0]
        yhi = yhi if yhi > ext.max[1] else ext.max[1]

    log.debug("Region boundaries: ")
    log.debug("xlo={xlo}, ylo={ylo}, xhi={xhi}, yhi={yhi}".format(xlo=xlo,
                                                                  ylo=ylo,
                                                                  xhi=xhi,
                                                                  yhi=yhi))

    subwcs = mywcs[ylo:yhi, xlo:xhi]
    subhdr = subwcs.sub([wcs.WCSSUB_CELESTIAL]).to_header()
    subdata = data[ylo:yhi, xlo:xhi]

    mask = shapelist.get_mask(header=subhdr,
                              shape=subdata.shape)
    log.debug("Shapes: data={0}, subdata={2}, mask={1}".format(data.shape, mask.shape, subdata.shape))
    return (xlo,xhi,ylo,yhi),mask

def parse_frequency_support(frequency_support_str):
    """
    Given a "Frequency Support" string from ALMA queries, parse it into a set
    of frequency ranges

    Example input:

        '[86.26..88.14GHz,976.56kHz, XX YY] U [88.15..90.03GHz,976.56kHz, XX YY] U [98.19..100.07GHz,976.56kHz, XX YY] U [100.15..102.03GHz,976.56kHz, XX YY]'
    """
    supports = frequency_support_str.split("U")
    freq_ranges = [(float(sup.strip('[] ').split("..")[0]),
                    float(sup.strip('[] ').split("..")[1].split(',')[0].strip(string.ascii_letters)))
                   *u.Unit(sup.strip('[] ').split("..")[1].split(',')[0].strip(string.punctuation+string.digits))
                   for sup in supports]
    return u.Quantity(freq_ranges)

def approximate_primary_beam_sizes(frequency_support_str):
    """
    Given a frequency support string, return the approximate 12m array beam
    size using 1.22 lambda / D
    """
    freq_ranges = parse_frequency_support(frequency_support_str)
    beam_sizes = [(1.22*fr.mean().to(u.m, u.spectral())/(12*u.m)).to(u.arcsec,
                                                                     u.dimensionless_angles())
                  for fr in freq_ranges]
    return u.Quantity(beam_sizes)

def make_finder_chart(target, radius, save_prefix, service=SkyView.get_images,
                      service_kwargs={'survey':['2MASS-K'], 'pixels':500},
                      alma_kwargs={'public':False, 'science':False},
                      private_band_colors=('red','darkred','orange','brown','maroon'),
                      public_band_colors=('blue','cyan','green','turquoise','teal'),
                      integration_time_contour_levels=np.logspace(0,5,base=2, num=6),
                     ):
    """
    Create a "finder chart" showing where ALMA has pointed in various bands,
    including different color coding for public/private data and each band.

    Contours are set at various integration times.

    Parameters
    ----------
    target : `astropy.coordinates` or str
        A legitimate target name
    radius : `astropy.units.Quantity`
        A degree-equivalent radius
    save_prefix : str
        The prefix for the output files.  Both .reg and .png files will be written.
        The .reg files will have the band numbers and public/private appended,
        while the .png file will be named prefix_almafinderchart.png
    service : function
        The `get_images` function of an astroquery service, e.g. SkyView.
    service_kwargs : dict
        The keyword arguments to pass to the specified service.  For example,
        for SkyView, you can give it the survey ID (e.g., 2MASS-K) and the
        number of pixels in the resulting image.  See the documentation for the
        individual services for more details.
    alma_kwargs : dict
        Keywords to pass to the ALMA archive when querying.
    private_band_colors / public_band_colors : tuple
        A tuple or list of colors to be associated with private/public observations
        in the various bands
    integration_time_contour_levels : list or np.array
        The levels at which to draw contours in units of seconds.  Default is
        log-spaced (2^n) seconds: [  1.,   2.,   4.,   8.,  16.,  32.])
    """
    import aplpy

    import pyregion
    from pyregion.parser_helper import Shape

    log.info("Querying {0} for images".format(service))
    images = service(target, radius=radius, **service_kwargs)

    log.info("Querying ALMA around {0}".format(target))
    catalog = Alma.query_region(coordinate=target, radius=radius,
                                **alma_kwargs)

    primary_beam_radii = [approximate_primary_beam_sizes(row['Frequency support'])
                          for row in catalog]

    bands = np.unique(catalog['Band'])
    log.info("The bands used include: {0}".format(bands))
    band_colors_priv = dict(zip(bands, private_band_colors))
    band_colors_pub = dict(zip(bands, public_band_colors))

    private_circle_parameters = {band: [(row['RA'],row['Dec'],np.mean(rad).to(u.deg).value)
                                 for row,rad in zip(catalog, primary_beam_radii)
                                 if row['Release date']!='' and row['Band']==band]
                                 for band in bands}
    public_circle_parameters = {band: [(row['RA'],row['Dec'],np.mean(rad).to(u.deg).value)
                                 for row,rad in zip(catalog, primary_beam_radii)
                                 if row['Release date']=='' and row['Band']==band]
                                 for band in bands}

    unique_private_circle_parameters = {band:
                                        np.array(list(set(private_circle_parameters[band])))
                                        for band in bands}
    unique_public_circle_parameters = {band:
                                       np.array(list(set(public_circle_parameters[band])))
                                       for band in bands}

    for band in bands:
        log.info( "BAND {0}".format(band) )
        privrows = sum((catalog['Band']==band) & (catalog['Release date'] != ''))
        pubrows  = sum((catalog['Band']==band) & (catalog['Release date'] == ''))
        log.info("PUBLIC:  Number of rows: {0}.  Unique pointings: "
                 "{1}".format(pubrows,
                 len(unique_public_circle_parameters[band])))
        log.info( "PRIVATE: Number of rows: {0}.  Unique pointings: "
                 "{1}".format(privrows,
                 len(unique_private_circle_parameters[band])))

    prv_regions = {band: pyregion.ShapeList([Shape('circle',[x,y,r]) for x,y,r
                                             in
                                             private_circle_parameters[band]])
                   for band in bands}
    pub_regions = {band: pyregion.ShapeList([Shape('circle',[x,y,r]) for x,y,r
                                             in
                                             public_circle_parameters[band]])
                   for band in bands}
    for band in bands:
        circle_pars = np.vstack([x for x in (private_circle_parameters[band],
                                        public_circle_parameters[band]) if any(x)])
        for r,(x,y,c) in zip(prv_regions[band]+pub_regions[band],
                             circle_pars):
            r.coord_format = 'fk5'
            r.coord_list = [x,y,c]
            r.attr = ([], {'color': 'green',  'dash': '0 ',  'dashlist': '8 3',
                           'delete': '1 ',  'edit': '1 ', 'fixed': '0 ',
                           'font': '"helvetica 10 normal roman"',  'highlite':
                           '1 ', 'include': '1 ',  'move': '1 ',  'select': '1',
                           'source': '1',  'text': '', 'width': '1 '})

        if prv_regions[band]:
            prv_regions[band].write('{0}_band{1}_private.reg'.format(save_prefix, band))
        if pub_regions[band]:
            pub_regions[band].write('{0}_band{1}_public.reg'.format(save_prefix, band))

    prv_mask = {band: fits.PrimaryHDU(prv_regions[band].get_mask(images[0][0]).astype('int'),
                               header=images[0][0].header) for band in bands
                if prv_regions[band]}
    pub_mask = {band: fits.PrimaryHDU(pub_regions[band].get_mask(images[0][0]).astype('int'),
                               header=images[0][0].header) for band in bands
                if pub_regions[band]}

    hit_mask_public = {band: np.zeros_like(images[0][0].data) for band in pub_mask}
    hit_mask_private = {band: np.zeros_like(images[0][0].data) for band in prv_mask}
    mywcs = wcs.WCS(images[0][0].header)

    for band in bands:
        log.debug('Band: {0}'.format(band))
        for row,rad in zip(catalog, primary_beam_radii):
            shape = Shape('circle', (row['RA'], row['Dec'],np.mean(rad).to(u.deg).value))
            shape.coord_format = 'fk5'
            shape.coord_list = (row['RA'], row['Dec'],np.mean(rad).to(u.deg).value)
            shape.attr = ([], {'color': 'green',  'dash': '0 ',  'dashlist': '8 3 ',
                               'delete': '1 ',  'edit': '1 ', 'fixed': '0 ',
                               'font': '"helvetica 10 normal roman"',
                               'highlite': '1 ', 'include': '1 ',  'move': '1 ',
                               'select': '1 ',  'source': '1',  'text': '',
                               'width': '1 '})
            log.debug('{1} {2}: {0}'.format(shape, row['Release date'], row['Band']))
            if row['Release date']!='' and row['Band']==band and band in prv_mask:
                (xlo,xhi,ylo,yhi),mask = pyregion_subset(shape,
                                                         hit_mask_private[band],
                                                         mywcs)
                log.debug("{0},{1},{2},{3}: {4}".format(xlo,xhi,ylo,yhi,mask.sum()))
                hit_mask_private[band][ylo:yhi,xlo:xhi] += row['Integration']*mask
            if row['Release date']=='' and row['Band']==band and band in pub_mask:
                (xlo,xhi,ylo,yhi),mask = pyregion_subset(shape,
                                                         hit_mask_public[band],
                                                         mywcs)
                log.debug("{0},{1},{2},{3}: {4}".format(xlo,xhi,ylo,yhi,mask.sum()))
                hit_mask_public[band][ylo:yhi,xlo:xhi] += row['Integration']*mask

    fig = aplpy.FITSFigure(images[0])
    fig.show_grayscale(stretch='arcsinh')
    for band in bands:
        if band in pub_mask:
            fig.show_contour(fits.PrimaryHDU(data=hit_mask_public[band],
                                             header=images[0][0].header),
                             levels=integration_time_contour_levels,
                             colors=[band_colors_pub[band]]*6)
        if band in prv_mask:
            fig.show_contour(fits.PrimaryHDU(data=hit_mask_private[band],
                                             header=images[0][0].header),
                             levels=integration_time_contour_levels,
                             colors=[band_colors_priv[band]]*6)

    fig.save('{0}_almafinderchart.png'.format(save_prefix))

    return images, catalog, hit_mask_public, hit_mask_private