#!/usr/bin/env python
from __future__ import print_function

if __name__ == '__main__':
    from optparse import OptionParser

    parser = OptionParser()

    parser.add_option('--outdir', help='output directory', default='.')
    parser.add_option('--samples', help='posterior samples file')

    parser.add_option('--fitsoutname', help='filename for the FITS file',
                      default='skymap.fits.gz')

    parser.add_option('--loadpost',
                      help='filename for pickled posterior state')

    parser.add_option('--maxpts', type='int',
                      help='maximum number of posterior points to use')

    parser.add_option('--trials', type='int', default=5,
                      help='number of trials at each clustering number '
                      '[default: %default]')

    parser.add_option('--disable-distance-map', dest='enable_distance_map',
                      action='store_false', default=True,
                      help='disable output of healpy map of '
                      'distance mean and s.d.')

    parser.add_option('-j', default=False, action='store_true',
                      help='Use all available cores [default: use one core]')

    parser.add_option('--seed', type=int, help='use specified random seed '
                      '[default: use system entropy]')

    parser.add_option('--objid', help='event ID to store in FITS header')

    (args, remaining) = parser.parse_args()

    if args.samples is None:
        parser.error('missing required argument: --samples')

    # Late imports
    from distutils.dir_util import mkpath
    from lalinference.io import fits
    from lalinference.io import hdf5
    from lalinference.bayestar.sky_map import rasterize
    from astropy.table import Table
    from astropy.time import Time
    import numpy as np
    import os
    import sys
    from collections import namedtuple
    import six.moves.cPickle as pickle
    import sky_area
    import sky_area.sky_area_clustering as sac

    print('reading samples ...')
    try:
        data = hdf5.read_samples(args.samples)
    except IOError:
        # FIXME: remove this code path once we support only HDF5
        data = Table.read(args.samples, format='ascii')

    if args.maxpts is not None:
        # Shuffle the data and take a random subsample.
        # Note that if arg.maxpts > len(data), then this
        # just selects the entire array.
        print('taking random subsample of chain ...')
        if args.seed is not None:
            # use fixed seed so that results are reproducible
            np.random.seed(args.seed)
        data = data[np.random.choice(len(data), args.maxpts, replace=False)]
    try:
        dist = data['dist']
    except KeyError:
        try:
            dist = data['distance']
        except KeyError:
            dist = None

    mkpath(args.outdir)

    if args.loadpost is None:
        if dist is None:
            if args.enable_distance_map:
                parser.error("The posterior samples file '{0}' does not have a distance column named 'dist' or 'distance'.  Cannot generate distance map.".format(args.samples))
            pts = np.column_stack((data['ra'], data['dec']))
        else:
            pts = np.column_stack((data['ra'], data['dec'], dist))
        if args.enable_distance_map:
            cls = sac.Clustered2Plus1DSkyKDE
        else:
            cls = sac.Clustered2DSkyKDE
        skypost = cls(pts, trials=args.trials, multiprocess=args.j)

        print('pickling ...')
        with open(os.path.join(args.outdir, 'skypost.obj'), 'w') as out:
            pickle.dump(skypost, out)
    else:
        with open(args.loadpost, 'r') as inp:
            skypost = pickle.load(inp)
        skypost.multiprocess = args.j

    print('making skymap ...')
    hpmap = rasterize(skypost.as_healpix())
    hpmap.meta['creator'] = parser.get_prog_name()
    hpmap.meta['origin'] = 'LIGO/Virgo'
    hpmap.meta['gps_creation_time'] = Time.now().gps
    # FIXME: Fake the LALSuite VCS info structure.
    # In the future, this code will probably live inside LALInference
    # and use the LALInference VCS info structure.
    VCSInfo = namedtuple(
        'VCSInfo', 'name version vcsStatus vcsId buildDate'.split())
    hpmap.meta['vcs_info'] = VCSInfo(
        'skyarea', sky_area.__version__, None, None, None)
    hpmap.meta['history'] = [
        '', 'Generated by running the following script:',
        ' '.join([parser.get_prog_name()] + sys.argv[1:])]
    if args.objid is not None:
        hpmap.meta['objid'] = args.objid
    if args.enable_distance_map:
        hpmap.meta['distmean'] = np.mean(dist)
        hpmap.meta['diststd'] = np.std(dist)

    names = data.dtype.names
    if 'time' in names:
        hpmap.meta['gps_time'] = data['time'].mean()
    elif 'time_mean' in names:
        hpmap.meta['gps_time'] = data['time_mean'].mean()
    elif 'time_maxl' in names:
        hpmap.meta['gps_time'] = data['time_maxl'].mean()
    else:
        parser.error(
            "Cannot find time, time_mean, or time maxl variable in posterior.")

    fits.write_sky_map(os.path.join(args.outdir, args.fitsoutname),
                       hpmap, nest=True)
