#!/home/hancock/V-Py/bin/python
"""
The Aegean source finding program.

"""

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

from optparse import OptionParser

# 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 scope2lat, get_aux_files
from AegeanTools.fits_image import Beam
from AegeanTools.catalogs import show_formats, check_table_formats, save_catalog
import multiprocessing

from AegeanTools import __version__, __date__

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


def check_projection(filename, options):
    """
    If the user has supplied --telescope when not required issue a warning
    If the user has not supplied --telescope when it is required issue a critical warning
    This depends on the projection of the image and if a psf map has been supplied:
    If ZEA then --telescope is required
    If SIN or a psf map has been supplied then --telescope is not required
    :param options: options from the option parser
    :return:
    """
    header = astropy.io.fits.getheader(filename)
    if "ZEA" in header['CTYPE1']:
        if options.telescope is None and options.imgpsf is None:
            logging.warning("For projection ZEA it is advised that you supply either")
            logging.warning("--telescope or --psf")
    if "SIN" in header['CTYPE1']:
        if options.telescope is not None:
            logging.warning("For projection SIN --telescope is not required. Ignoring.")
            options.telescope = None
    return

if __name__ == "__main__":
    usage = "usage: %prog [options] FileName.fits"
    parser = OptionParser(usage=usage)
    parser.add_option("--find", dest='find', action='store_true', default=False,
                      help='Source finding mode. [default: true, unless --save or --measure are selected]')
    parser.add_option("--cores", dest="cores", type="int", default=None,
                      help="Number of CPU cores to use for processing [default: all cores]")
    parser.add_option("--debug", dest="debug", action="store_true", default=False,
                      help="Enable debug mode. [default: false]")
    parser.add_option("--hdu", dest="hdu_index", type="int", default=0,
                      help="HDU index (0-based) for cubes with multiple images in extensions. [default: 0]")
    parser.add_option("--out", dest='outfile', default=None,
                      help="Destination of Aegean catalog output. [default: No output]")
    parser.add_option("--table", dest='tables', default=None,
                      help="Additional table outputs, format inferred from extension. [default: none]")
    parser.add_option("--tformats", dest='table_formats', action="store_true", default=False,
                      help='Show a list of table formats supported in this install, and their extensions')
    parser.add_option("--forcerms", dest='rms', type='float', default=None,
                      help="Assume a single image noise of rms, and a background of zero. [default: false]")
    parser.add_option("--noise", dest='noiseimg', default=None,
                      help="A .fits file that represents the image noise (rms), created from Aegean with --save " +
                           "or BANE. [default: none]")
    parser.add_option('--background', dest='backgroundimg', default=None,
                      help="A .fits file that represents the background level, created from Aegean with --save " +
                           "or BANE. [default: none]")
    parser.add_option('--psf', dest='imgpsf', default=None,
                      help="A .fits file that represents the local PSF. ")

    parser.add_option('--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]")
    parser.add_option("--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]")
    parser.add_option('--seedclip', dest='innerclip', type='float', default=5,
                      help='The clipping value (in sigmas) for seeding islands. [default: 5]')
    parser.add_option('--floodclip', dest='outerclip', type='float', default=4,
                      help='The clipping value (in sigmas) for growing islands. [default: 4]')
    parser.add_option('--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].')
    parser.add_option('--telescope', dest='telescope', type=str, default=None,
                      help='The name of the telescope used to collect data. [MWA|VLA|ATCA|LOFAR]')
    parser.add_option('--lat', dest='lat', type=float, default=None,
                      help='The latitude of the telescope used to collect data.')
    parser.add_option('--versions', dest='file_versions', action="store_true", default=False,
                      help='Show the file versions of relevant modules. [default: false]')
    parser.add_option('--island', dest='doislandflux', action="store_true", default=False,
                      help='Also calculate the island flux in addition to the individual components. [default: false]')
    parser.add_option('--nopositive', dest='nopositive', action="store_true", default=False,
                      help="Don't report sources with positive fluxes. [default: false]")
    parser.add_option('--negative', dest='negative', action="store_true", default=False,
                      help="Report sources with negative fluxes. [default: false]")
    parser.add_option('--blankout', dest='blank', action="store_true", default=False,
                      help="Create a blanked output image. [Only works if cores=1].")

    parser.add_option('--region', dest='region', default=None,
                      help="Use this regions file to restrict source finding in this image.")

    parser.add_option('--save', dest='save', action="store_true", default=False,
                      help='Enable the saving of the background and noise images. Sets --find to false. ' +
                           '[default: false]')
    parser.add_option('--outbase', dest='outbase', default=None,
                      help='If --save is True, then this specifies the base name of the background and noise images. ' +
                           '[default: inferred from input image]')
    parser.add_option('--measure', dest='measure', action='store_true', default=False,
                      help='Enable forced measurement mode. Requires an input source list via --input. ' +
                           'Sets --find to false. [default: false]')
    parser.add_option('--priorized', dest='priorized', default=0, type=int,
                      help="Enable priorized fitting, with stage = n [default=1]")
    parser.add_option('--ratio', dest='ratio', default=None, type=float,
                      help="The ratio of synthesized beam sizes (image psf / input catalog psf). " +
                           "For use with priorized.")
    parser.add_option('--noregroup', dest='regroup', default=True, action='store_false',
                      help='Do not regroup islands before priorized fitting.')
    parser.add_option('--input', dest='input', default=None,
                      help='If --measure is true, this gives the filename for a catalog of locations at which ' +
                           'fluxes will be measured. [default: none]')
    parser.add_option('--catpsf', dest='catpsf', 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.')

    (options, args) = parser.parse_args()

    # 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__))

    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("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__))
        try:
            import h5py

            log.info("h5py {0} from {1}".format(h5py.__version__, h5py.__file__))
        except ImportError:
            log.info("h5py not found")
        sys.exit(0)

    # print help if the user enters no options or filename
    if len(args) == 0:
        parser.print_help()
        sys.exit(0)

    # check that a valid filename was entered
    filename = args[0]
    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)

    # 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 measure/save are enabled we turn off "find" unless it was specifically set
    if (options.measure or 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))

    # determine the latitude of the telescope
    if options.telescope is not None:
        lat = scope2lat(options.telescope)
    elif options.lat is not None:
        lat = options.lat
    else:
        lat = None

    # 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)
        sys.exit(0)

    # 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")
            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)

    # 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 = []

    # do forced measurements using catfile
    if options.measure and options.priorized == 0:
        raise NotImplementedError("forced measurements are not supported")

    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.warn("--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,
                                 outfile=options.outfile, bkgin=options.backgroundimg,
                                 rmsin=options.noiseimg, beam=options.beam, lat=lat, imgpsf=options.imgpsf,
                                 catpsf=options.catpsf,
                                 stage=options.priorized, ratio=options.ratio, outerclip=options.outerclip,
                                 cores=options.cores, doregroup=options.regroup)

    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,
                                         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, lat=lat, imgpsf=options.imgpsf, blank=options.blank)
        if options.blank:
            outname = basename+'_blank.fits'
            sf.save_image(outname)
        if len(found) == 0:
            log.info("No sources found in image")

    sources = sf.sources
    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}
        for t in options.tables.split(','):
            save_catalog(t, sources)
    sys.exit()
