aegean

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

"""
The Aegean source finding program.

"""

import sys
import os
import numpy as np
import scipy
import lmfit
import astropy
import logging
import logging.config
import argparse

# need Region in the name space in order to be able to unpickle it
try:
    from AegeanTools.regions import Region

    region_available = True
    try:
        import cPickle as pickle
    except ImportError:
        import pickle
except ImportError:
    region_available = False

from AegeanTools.source_finder import get_aux_files
from AegeanTools.wcs_helpers import Beam
from AegeanTools.catalogs import show_formats, check_table_formats, save_catalog
from AegeanTools import fitting
import multiprocessing

from AegeanTools import __version__, __date__, __citation__

header = """#Aegean version {0}
# on dataset: {1}"""


def check_projection(filename, options, log):
    """
    Kindly let the user know that projections other than SIN need some careful thought.

    parameters
    ----------
    filename : str
        The input fits filename

    options : argparse options
        Options from the command line
    """
    header = astropy.io.fits.getheader(filename)
    if not( "SIN" in header['CTYPE1']):
        if options.imgpsf is None:
            projection = header['CTYPE1'].split('-')[-1]
            log.warning("For projection {0} you should consider supplying a psf via --psf".format(projection))
    return


if __name__ == "__main__":
    parser = argparse.ArgumentParser(prog='aegean', prefix_chars='-')
    parser.add_argument('image', nargs='?', default=None)
    group1 = parser.add_argument_group('Configuration Options')

    group1.add_argument("--find", dest='find', action='store_true', default=False,
                        help='Source finding mode. [default: true, unless --save or --measure are selected]')
    group1.add_argument("--cores", dest="cores", type=int, default=None,
                        help="Number of CPU cores to use for processing [default: all cores]")
    group1.add_argument("--hdu", dest="hdu_index", type=int, default=0,
                        help="HDU index (0-based) for cubes with multiple images in extensions. [default: 0]")

    group1.add_argument('--beam', dest='beam', type=float, nargs=3, default=None,
                        help='The beam parameters to be used is "--beam major minor pa" all in degrees. ' +
                             '[default: read from fits header].')
    group1.add_argument('--telescope', dest='telescope', type=str, default=None,
                        help='DEPRECATED')
    group1.add_argument('--lat', dest='lat', type=float, default=None,
                        help='DEPRECATED')
    group1.add_argument('--slice', dest='slice', type=int, default=None,
                        help='If the input data is a cube, then this slice will determine the array index of the image which will be processed by aegean')


    # Input
    group2 = parser.add_argument_group("Input Options")
    group2.add_argument("--forcerms", dest='rms', type=float, default=None,
                        help="Assume a single image noise of rms. [default: None]")
    group2.add_argument("--forcebkg", dest='bkg', type=float, default=None,
                        help="Assume a single image background of bkg. [default: None]")
    group2.add_argument("--noise", dest='noiseimg', default=None, type=str,
                        help="A .fits file that represents the image noise (rms), created from Aegean with --save " +
                             "or BANE. [default: none]")
    group2.add_argument('--background', dest='backgroundimg', default=None, type=str,
                        help="A .fits file that represents the background level, created from Aegean with --save " +
                             "or BANE. [default: none]")
    group2.add_argument('--psf', dest='imgpsf', default=None, type=str,
                        help="A .fits file that represents the local PSF. ")
    group2.add_argument('--autoload', dest='autoload', action="store_true", default=False,
                        help="Automatically look for background, noise, region, and psf files "+
                             "using the input filename as a hint. [default: don't do this]")

    # Output
    group3 = parser.add_argument_group("Output Options")
    group3.add_argument("--out", dest='outfile', default=None, type=str,
                        help="Destination of Aegean catalog output. [default: No output]")
    group3.add_argument("--table", dest='tables', default=None, type=str,
                        help="Additional table outputs, format inferred from extension. [default: none]")
    group3.add_argument("--tformats", dest='table_formats', action="store_true", default=False,
                        help='Show a list of table formats supported in this install, and their extensions')
    group3.add_argument('--blankout', dest='blank', action="store_true", default=False,
                        help="Create a blanked output image. [Only works if cores=1].")
    group3.add_argument('--colprefix', dest='column_prefix', default=None, type=str,
                        help='Prepend each column name with "prefix_". Default = prepend nothing')

    # SF config options
    group4 = parser.add_argument_group("Source finding/fitting configuration options")
    group4.add_argument("--maxsummits", dest='max_summits', type=float, default=None,
                        help="If more than *maxsummits* summits are detected in an island, no fitting is done, " +
                             "only estimation. [default: no limit]")
    group4.add_argument('--seedclip', dest='innerclip', type=float, default=5,
                        help='The clipping value (in sigmas) for seeding islands. [default: 5]')
    group4.add_argument('--floodclip', dest='outerclip', type=float, default=4,
                        help='The clipping value (in sigmas) for growing islands. [default: 4]')
    group4.add_argument('--island', dest='doislandflux', action="store_true", default=False,
                        help='Also calculate the island flux in addition to the individual components. [default: false]')
    group4.add_argument('--nopositive', dest='nopositive', action="store_true", default=False,
                        help="Don't report sources with positive fluxes. [default: false]")
    group4.add_argument('--negative', dest='negative', action="store_true", default=False,
                        help="Report sources with negative fluxes. [default: false]")
    group4.add_argument('--region', dest='region', type=str, default=None,
                        help="Use this regions file to restrict source finding in this image.\nUse MIMAS region (.mim) files.")
    group4.add_argument('--nocov', dest='docov', action="store_false", default=True,
                        help="Don't use the covariance of the data in the fitting proccess. [Default = False]")
    group4.add_argument('--condon', dest='condon', action="store_true", default=False,
                        help="replace errors with those suggested by Condon'97. [Default = False]")

    # priorized fitting
    group5 = parser.add_argument_group("Priorized Fitting config options","in addition to the above source fitting options")
    group5.add_argument('--priorized', dest='priorized', default=0, type=int,
                        help="Enable priorized fitting level n=[1,2,3]. 1=fit flux, 2=fit flux/position, 3=fit flux/position/shape. " + 
                             "See the GitHub wiki for more details." )
    group5.add_argument('--ratio', dest='ratio', default=None, type=float,
                        help="The ratio of synthesized beam sizes (image psf / input catalog psf). " +
                             "For use with priorized.")
    group5.add_argument('--noregroup', dest='regroup', default=True, action='store_false',
                        help='Do not regroup islands before priorized fitting.')
    group5.add_argument('--input', dest='input', type=str, default=None,
                        help='If --priorized is used, this gives the filename for a catalog of locations at which ' +
                             'fluxes will be measured.')
    group5.add_argument('--catpsf', dest='catpsf', type=str, default=None,
                        help='A psf map corresponding to the input catalog. This will allow for the correct resizing of' +
                             ' sources when the catalog and image psfs differ.')


    # Debug and extras
    group6 = parser.add_argument_group("Extra options")
    group6.add_argument('--save', dest='save', action="store_true", default=False,
                        help='Enable the saving of the background and noise images. Sets --find to false. ' +
                             '[default: false]')
    group6.add_argument('--outbase', dest='outbase', type=str, default=None,
                        help='If --save is True, then this specifies the base name of the background and noise images. ' +
                             '[default: inferred from input image]')
    group6.add_argument("--debug", dest="debug", action="store_true", default=False,
                        help="Enable debug mode. [default: false]")
    group6.add_argument('--versions', dest='file_versions', action="store_true", default=False,
                        help='Show the file versions of relevant modules. [default: false]')
    group6.add_argument('--cite', dest='cite', action="store_true", default=False,
                        help='Show citation information.')

    options = parser.parse_args()

    invocation_string = " ".join(sys.argv[:])

    # configure logging
    logging.basicConfig(format="%(module)s:%(levelname)s %(message)s")
    log = logging.getLogger("Aegean")
    logging_level = logging.DEBUG if options.debug else logging.INFO
    log.setLevel(logging_level)
    log.info("This is Aegean {0}-({1})".format(__version__, __date__))

    log.debug("Run as:\n{0}".format(invocation_string))

    # deprecation warnings
    if options.telescope is not None:
        log.warn("--telescope has been deprecated. Ignoring.")

    if options.lat is not None:
        log.warn("--lat has been deprecated. Ignoring.")
  
    # options that don't require image inputs
    if options.cite:
        print(__citation__)
        sys.exit(0)

    import AegeanTools
    from AegeanTools.source_finder import SourceFinder, check_cores
    # source finding object
    sf = SourceFinder(log=log)

    if options.table_formats:
        show_formats()
        sys.exit(0)

    if options.file_versions:
        log.info("AegeanTools {0} from {1}".format(AegeanTools.__version__, AegeanTools.__file__))
        log.info("Numpy {0} from {1} ".format(np.__version__, np.__file__))
        log.info("Scipy {0} from {1}".format(scipy.__version__, scipy.__file__))
        log.info("AstroPy {0} from {1}".format(astropy.__version__, astropy.__file__))
        log.info("LMFit {0} from {1}".format(lmfit.__version__, lmfit.__file__))
        sys.exit(0)

    # print help if the user enters no options or filename
    if options.image is None:
        parser.print_help()
        sys.exit(0)

    from AegeanTools.source_finder import SourceFinder, check_cores
    # source finding object
    sf = SourceFinder(log=log)

    # check that a valid filename was entered
    filename = options.image
    if not os.path.exists(filename):
        log.error("{0} not found".format(filename))
        sys.exit(1)

    # check to see if the user has supplied --telescope/--psf when required
    check_projection(filename, options, log)

    # tell numpy to shut up about "invalid values encountered"
    # Its just NaN's and I don't need to hear about it once per core
    np.seterr(invalid='ignore', divide='ignore')

    # check for nopositive/negative conflict
    if options.nopositive and not options.negative:
        log.warning('Requested no positive sources, but no negative sources. Nothing to find.')
        sys.exit()

    # if priorized/save are enabled we turn off "find" unless it was specifically set
    if (options.save or options.priorized) and not options.find:
        options.find = False
    else:
        options.find = True

    # debugging in multi core mode is very hard to understand
    if options.debug:
        log.info("Setting cores=1 for debugging")
        options.cores = 1

    # check/set cores to use
    if options.cores is None:
        options.cores = multiprocessing.cpu_count()
        log.info("Found {0} cores".format(options.cores))
    if options.cores > 1:
        options.cores = check_cores(options.cores)
    log.info("Using {0} cores".format(options.cores))

    hdu_index = options.hdu_index
    if hdu_index > 0:
        log.info("Using hdu index {0}".format(hdu_index))

    # create a beam object from user input
    if options.beam is not None:
        beam = options.beam
        if len(beam) != 3:
            beam = beam.split()
            print(("Beam requires 3 args. You supplied '{0}'".format(beam)))
            sys.exit(1)
        options.beam = Beam(beam[0], beam[1], beam[2])
        log.info("Using user supplied beam parameters")
        log.info("Beam is {0} deg x {1} deg with pa {2}".format(options.beam.a, options.beam.b, options.beam.pa))

    # auto-load background, noise, psf and region files
    basename = os.path.splitext(filename)[0]
    if options.autoload:
        files = get_aux_files(filename)
        if files['bkg'] and not options.backgroundimg:
            options.backgroundimg = files['bkg']
            log.info("Found background {0}".format(options.backgroundimg))
        if files['rms'] and not options.noiseimg:
            options.noiseimg = files['rms']
            log.info("Found noise {0}".format(options.noiseimg))
        if files['mask'] and not options.region:
            options.region = files['mask']
            log.info("Found region {0}".format(options.region))
        if files['psf'] and not options.imgpsf:
            options.imgpsf = files['psf']
            log.info("Found psf {0}".format(options.imgpsf))

    # check that the aux input files exist
    if options.backgroundimg and not os.path.exists(options.backgroundimg):
        log.error("{0} not found".format(options.backgroundimg))
        sys.exit(1)
    if options.noiseimg and not os.path.exists(options.noiseimg):
        log.error("{0} not found".format(options.noiseimg))
        sys.exit(1)
    if options.imgpsf and not os.path.exists(options.imgpsf):
        log.error("{0} not found".format(options.imgpsf))
        sys.exit(1)
    if options.catpsf and not os.path.exists(options.catpsf):
        log.error("{0} not found".format(options.catpsf))
        sys.exit(1)

    if options.region is not None:
        if not os.path.exists(options.region):
            log.error("Region file {0} not found".format(options.region))
            sys.exit(1)
        if not region_available:
            log.error("Could not import AegeanTools/regions.py")
            log.error("(you probably need to install HealPy)")
            sys.exit(1)

    # Generate and save the background FITS files with the Aegean default calculator
    if options.save:
        sf.save_background_files(filename, hdu_index=hdu_index, cores=options.cores, beam=options.beam,
                                 outbase=options.outbase, bkgin=options.backgroundimg, rmsin=options.noiseimg,
                                 rms=options.rms, bkg=options.bkg)
        sys.exit(0)

    # check that the output table formats are supported (if given)
    # BEFORE any cpu intensive work is done
    if options.tables is not None:
        if not check_table_formats(options.tables):
            log.critical("One or more output table formats are not supported: Exiting")
            sys.exit(1)


    # if an outputfile was specified open it for writing
    if options.outfile == 'stdout':
        options.outfile = sys.stdout
    elif options.outfile is not None:
        options.outfile = open(options.outfile, 'w')

    sources = []

    if options.find:
        log.info("Finding sources.")
        found = sf.find_sources_in_image(filename, outfile=options.outfile, hdu_index=options.hdu_index,
                                         rms=options.rms, bkg=options.bkg,
                                         max_summits=options.max_summits,
                                         innerclip=options.innerclip,
                                         outerclip=options.outerclip, cores=options.cores, rmsin=options.noiseimg,
                                         bkgin=options.backgroundimg, beam=options.beam,
                                         doislandflux=options.doislandflux,
                                         nonegative=not options.negative, nopositive=options.nopositive,
                                         mask=options.region, imgpsf=options.imgpsf, blank=options.blank,
                                         docov=options.docov, cube_index=options.slice)
        if options.blank:
            outname = basename+'_blank.fits'
            sf.save_image(outname)
        if len(found) == 0:
            log.info("No sources found in image")

    if options.priorized > 0:
        if options.ratio is not None:
            if options.ratio <= 0:
                log.error("ratio must be positive definite")
                sys.exit(1)
            if options.ratio < 1:
                log.error("ratio <1 is not advised. Have fun!")
        if options.input is None:
            log.error("Must specify input catalog when --priorized is selected")
            sys.exit(1)
        if not os.path.exists(options.input):
            log.error("{0} not found".format(options.input))
            sys.exit(1)
        log.info("Priorized fitting of sources in input catalog.")

        log.info("Stage = {0}".format(options.priorized))
        if options.doislandflux:
            log.warning("--island requested but not yet supported for priorized fitting")
        sf.priorized_fit_islands(filename, catalogue=options.input, hdu_index=options.hdu_index,
                                 rms=options.rms, bkg=options.bkg,
                                 outfile=options.outfile, bkgin=options.backgroundimg,
                                 rmsin=options.noiseimg, beam=options.beam, imgpsf=options.imgpsf,
                                 catpsf=options.catpsf,
                                 stage=options.priorized, ratio=options.ratio, outerclip=options.outerclip,
                                 cores=options.cores, doregroup=options.regroup, docov=options.docov,
                                 cube_index=options.slice)


    sources = sf.sources
    # if --condon is set then we replace all the errors with those described by Condon'97
    if options.condon:
        # theta_N is the FWHM of the smoothing kernel (the noise correlation)
        # which in this case is the same as the synthesized beam FWHM
        if options.beam:
            theta_n = np.sqrt(options.beam.a * options.beam.b)
            psf = None
        else:
            psf = sf.global_data.psfhelper
            theta_n = None
        for s in sources:
            fitting.condon_errors(s, theta_n=theta_n, psf=psf)

    log.info("found {0} sources total".format(len(sources)))
    if len(sources) > 0 and options.tables:
        meta = {"PROGRAM": "Aegean",
                "PROGVER": "{0}-({1})".format(__version__, __date__),
                "FITSFILE": filename,
                "RUN-AS": invocation_string}
        for t in options.tables.split(','):
            save_catalog(t, sources, prefix=options.column_prefix, meta=meta)
    sys.exit()