#
# POIS: Pan-STARRS (or Pricey's) Optimal Image Subtraction
#
import os, re, sys, time
import ds9
import ps
import psLibSwig as psLib; from psLibSwig import psTraceSetLevel, psTrace
import poisSwig as pois

def getTime():
    return time.clock()

def displayImage(image, frame = 0, caption = None, ctype = "green", init = False):
    global init_display
    if init_display:
        init = True
        
    ds9.mtv(image, init = init, frame = frame)
    if caption:
        ds9.dot(caption, 10, image.numRows + 5, ctype = ctype, frame = frame)

    init_display = False

def run(showLeaks = True, *args, **nargs):
    #"""Wrapper for run, with a check on memory leaks after out-of-scope
    #objects have been garbage collected."""
    #
    # Only look for leaks of memory allocated after this point
    #
    id0 = psLib.psMemGetId()
    #
    # Do the work
    #
    value = _run(*args, **nargs)
    #
    # Look for otherwise-undetected problems
    #
    ps.raisePsError()

    nleak = psLib.psMemCheckLeaks(id0, None, None, False)
    if nleak:
        if showLeaks != False:
            ps.memCheckLeaks(sys.stderr, id0 = id0, addresses = True)
            
        raise IOError, "Memory Leak (%d pointers); id0 = %d" % (nleak, id0)

    return value

def _run(fileDir = ".", refFile = "test1.fits", inFile = "test2.fits", outFile=None,
         section = "[100:500,100:500]",
         testing = False, verbose = 0, psLib_verbose = 0,
         display = {}, setTrace = True):
    """POIS: Pan-STARRS (or Pricey's) Optimal Image Subtraction.

Attempts to fit for a general convolution kernel that will match the PSFs between two images.
Arguments:
   fileDir	Directory for input and output files (default: .)
   refFile      Reference image (default: test1.fits)
   inFile       File to be analysed (default: test2.fits)
   outFile	Write output to this file (default: none)
   section	Section of datafiles to process, or None (=> all) (default: [100:500,100:500])

   testing	Write some extra test information/files (default: False)
   verbose      Verbosity passed to psTrace (default: 0).0, psLib_verbose = 0,
   setTrace	Should the code fiddle with trace levels? (default: True)

   display	Dictionary of things that can be displayed; use display="help" for help (default: {})

E.g.
   pois.run(fileDir="data", verbose=3, display={'stamps':1, 'reference':1, 'subtracted':1})
"""

    startTime = getTime();

    if setTrace:
        psTraceSetLevel(".", 0)
        psTraceSetLevel("pois", verbose)
        psTraceSetLevel("pois.config", verbose)
        psTraceSetLevel("pois.time", verbose + 2)
        psTraceSetLevel("pois.solution", 0)
        psTraceSetLevel("pois.kernelBasisFunctions", verbose)
        psTraceSetLevel("pois.calculateEquation", verbose)
        psTraceSetLevel("pois.convolveImage", verbose)
        psTraceSetLevel("pois.findStamps", verbose)
        psTraceSetLevel("pois.parseConfig", verbose)
        psTraceSetLevel("pois.makeMask", verbose)
        psTraceSetLevel("pois.extractKernel", verbose)
        psTraceSetLevel("pois.calculateDeviations", verbose)

        psTraceSetLevel("ps.utils", psLib_verbose)

        if verbose > 3:
            print "Trace levels:"
            psLib.psTracePrintLevels()

    # Set logging level
    psLib.psLogSetDestination("dest:stderr")

    #
    # Set missing keys in display
    keys = ['all', 'pause']            # permitted keywords in display
    subkeys = ['convolved', 'input', 'kernel', 'mask', 'reference', 'stamps', 'subtracted']

    if display == "help":
        print """Usage:
        display={'key1':1, 'key2':True, 'key3':0, ...}   (1==True, 0==False)
Keys:
    all		Display everything that can be displayed
    pause	Pause after every iteration over the stamp list

    convolved	Show the kernel-convolved reference image
    input	Show the input image
    kernel	Show the convolution kernel
    mask	Show the image mask
    reference	Show the reference image
    stamps	Overlay the reference image with the centroids of stamp stars (circle bad ones in red)
    subtracted	Show the subtracted image
        """
        return            

    for k in (keys + subkeys):
        if not k in display:
            display[k] = False

    # Are illegal keys present in display[]?
    illegal = []
    for k in display.keys():
        if not k in keys + subkeys:
            illegal.append(k)

    if illegal:
        raise ("Unrecognised keywords %s are present in display[]" % (illegal.__str__()))

    # Deal with display values
    if display['all']:
        for k in subkeys:
            display[k] = True

    if display['stamps']:
        display['input'] = True

    global init_display;
    init_display = True            # have we initialised the display?

    # Set desired variables from config file
    if not re.search("/$", fileDir):
        fileDir += "/"

    config = pois.poisConfigAlloc()      # Configuration values
    config.inFile = fileDir + inFile
    config.refFile = fileDir + refFile
    try:
        config.outFile = fileDir + outFile
    except:
        config.outFile = outFile

    # Read inputs
    imageRegion = psLib.psRegionAlloc(0, 0, 0, 0)

    fd = psLib.psFitsAlloc(config.refFile)
    refImage = psLib.psFitsReadImage(None, fd, imageRegion, 0)
    del fd
    psTrace("pois", 3, ("Read %s: %dx%d\n" % (config.refFile, refImage.numCols, refImage.numRows)))

    fd = psLib.psFitsAlloc(config.inFile)
    inImage = psLib.psFitsReadImage(None, fd, imageRegion, 0)
    del fd
    psTrace("pois", 3, ("Read %s: %dx%d\n" % (config.inFile, inImage.numCols, inImage.numRows)))
    
    del imageRegion

    if section:
        inImage = psLib.psImageSubsection(inImage, section)
        refImage = psLib.psImageSubsection(refImage, section)

    if display['reference']:
        displayImage(refImage, 1, "Reference Image", "red")
    if display['input']:
        displayImage(inImage, 2, "Input Image", "red")

    if (refImage.numCols != inImage.numCols) or (refImage.numRows != inImage.numRows):
	raise IOError, \
              ("Reference and input images must have the same dimensions: %dx%d vs %dx%d\n" %
               (refImage.numCols, refImage.numRows, inImage.numCols, inImage.numRows))

    config.xImage = refImage.numCols
    config.yImage = refImage.numRows

    psTrace("pois.time", 1, ("Inputs read at %f sec\n" % (getTime() - startTime)))

    # Add a background offset to get the image all above zero
    # An image that goes negative (e.g., a background-subtracted image) produces
    # horrible results --- the matrix goes NaN because the noise is determined
    # as the square root of the value.
    if config.background != 0.0:
            psBinaryOp(refImage, refImage, "+", psScalarAlloc(config.background, PS_TYPE_F32))
            psBinaryOp(inImage, inImage, "+", psScalarAlloc(config.background, PS_TYPE_F32))

    # Generate the kernel parameters
    kernelBasisFunctions = pois.poisKernelBasisFunctions(config)

    # Generate a mask
    mask = pois.poisMakeMask(None, refImage, inImage, config)

    stamps = None                       # Array of stamps
    stampMask = None                    # Mask for stamps
    solution = None                     # Solution vector

    psLib.psMemCheckCorruption(1)

    # Iterate for a good solution
    for iterNum in range(0, config.numIter):
	psTrace("pois", 1, ("Iteration %d...\n" % (iterNum)))

	# Find stamps
        stamps = pois.poisFindStamps(stamps, refImage, mask, config);
        
	numStamps = 0;
        for s in range(0, stamps.n):
	    stamp = pois.poisStamp_Cast(stamps.get_data(s)) # Stamp of interest
            if stamp:
                if stamp.status == pois.POIS_STAMP_BAD:
                    ctype = "red"; dotType = "x"
                else:
                    ctype = "green"; dotType = "+"
                    numStamps += 1
                    
                if display['stamps']:
                    ds9.dot(dotType, stamp.x, stamp.y, size = 5, ctype = ctype, frame = 1)

	psTrace("pois.time", 1, ("%d stamps found at %f sec\n" % (numStamps, getTime() - startTime)))
        if numStamps == 0:
            raise IOError, "No stamps found.  Check input parameters"

	# Calculate equation
	pois.poisCalculateEquation(stamps, refImage, inImage, kernelBasisFunctions, config)
	psTrace("pois.time", 1, ("Equation calculated at %f sec\n" % (getTime() - startTime)))

	# Solve the equation
	solution = pois.poisSolveEquation(solution, stamps, config)
	psTrace("pois.time", 1, ("Matrix equation solved at %f sec\n" % (getTime() - startTime)))

        if testing:
            # Print solution
            print "Solution:"
            for i in range(0, solution.n - 1):
                kernel = kernelBasisFunctions.get_data(i) # The i-th kernel basis function
                kernel = pois.poisKernelBasis_Cast(kernel)
                print "%d: (%d,%d) x^%d y^%d - %f" % \
                      (i, kernel.u, kernel.v, kernel.xOrder, kernel.yOrder, solution.get_data_F64()[i])


        if testing or display['kernel']:
            kernelImage = pois.poisExtractKernel(solution, kernelBasisFunctions, 0.0, 0.0, config)

            if display['kernel']:
                displayImage(kernelImage, 5, "Kernel")

            if testing:
                # Save the kernel postage stamp
                kernelName = "kernel%d.fits" % iterNum
                kernelFile = psLib.psFitsAlloc(kernelName)
                psLib.psFitsWriteImage(kernelFile, None, kernelImage, 0, None)
                psTrace("pois", 2, ("Kernel written to %s\n" %kernelName))

	# Calculate deviations for the stamps
	deviations = pois.poisCalculateDeviations(None, stamps, refImage, inImage, mask,
						       kernelBasisFunctions, solution, config)
        # Have we converged?
        rejected = pois.poisRejectStamps(stamps, mask, deviations, config)

        if rejected == None:
            break
        
        if display['stamps']:
            iter = psLib.psListIteratorAlloc(rejected, psLib.PS_LIST_HEAD, False)
            while True:
                stamp = pois.poisStamp_Cast(psLib.psListGetAndIncrement(iter))
                if not stamp:
                    break
                ds9.dot("o", stamp.x, stamp.y, size = 5, ctype = "red", frame = 1)

        if display['pause']:
            if raw_input("Iteration %d completed: continue? (q to abort) " % iterNum) == "q":
                raise "XXX"

    # If there was rejection on the last round, re-solve the equation using only the good stamps
    if iterNum == config.numIter:
        psTrace("pois", 2, ("Failed to converge in %d iterations\n" % (config.numIter)))
	solution = pois.poisSolveEquation(solution, stamps, config);

    # Clear the mask for bad stamp areas
    if display['mask']:
        displayImage(mask, 0, "Stamps mask")

    psLib.psBinaryOp(mask, mask, "&", psLib.psScalarAlloc(~pois.POIS_MASK_STAMP, psLib.PS_TYPE_U8))

    # Convolve the input image
    convImage = pois.poisConvolveImage(refImage, mask, solution, kernelBasisFunctions, config)
    psTrace("pois.time", 1, ("Convolution completed at %f sec\n" % (getTime() - startTime)))

    if testing and config.outFile:
        # Save the convolved image
        convName = "%s.conv" % config.outFile
        convFile = psLib.psFitsAlloc(convName)
        psLib.psFitsWriteImage(convFile, None, convImage, 0, None)
        psTrace("pois", 2, ("Convolved image written to %s\n" % (convName)))

    if display['convolved']:
        displayImage(convImage, 4, "Convolved image", "red")

    # Do the subtraction
    subImage = psLib.psImageAlloc(config.xImage, config.yImage, psLib.PS_TYPE_F32)
    psLib.psBinaryOp(subImage, inImage, "-", convImage)
    psTrace("pois.time", 1, ("Subtraction completed at %f sec\n" % (getTime() - startTime)))

    if testing:
        stats = psLib.psStatsAlloc(psLib.PS_STAT_SAMPLE_MEAN | psLib.PS_STAT_SAMPLE_STDEV)
	psLib.psImageStats(stats, subImage, mask, 0)
	print "Subtracted image statistics: %f %f\n" % (stats.sampleMean, stats.sampleStdev)
	#psFree(stats);

    # Mask out the edges
    X0 = config.xKernel; X1 = mask.numCols - config.xKernel # non-masked area is (X0,Y0) -- (X1-1, Y1-1)
    Y0 = config.yKernel; Y1 = mask.numRows - config.yKernel

    tmp = psLib.psImageSubset(mask, 0, 0, mask.numCols, Y0)
    psLib.psBinaryOp(tmp, tmp, "|", psLib.psScalarAlloc(pois.POIS_MASK_BAD, psLib.PS_TYPE_U8))

    tmp = psLib.psImageSubset(mask, 0, Y1, mask.numCols, mask.numRows)
    psLib.psBinaryOp(tmp, tmp, "|", psLib.psScalarAlloc(pois.POIS_MASK_BAD, psLib.PS_TYPE_U8))

    tmp = psLib.psImageSubset(mask, 0, Y0, X0, Y1)
    psLib.psBinaryOp(tmp, tmp, "|", psLib.psScalarAlloc(pois.POIS_MASK_BAD, psLib.PS_TYPE_U8))

    tmp = psLib.psImageSubset(mask, X1, Y0, mask.numCols, Y1)
    psLib.psBinaryOp(tmp, tmp, "|", psLib.psScalarAlloc(pois.POIS_MASK_BAD, psLib.PS_TYPE_U8))
    del tmp

    pois.poisImageSetValInMask(subImage, subImage, mask, 0, pois.POIS_MASK_BAD)

    if False:
        for y in range(config.yKernel, subImage.numRows - config.yKernel):
            for x in range(config.xKernel, subImage.numCols - config.xKernel):
                if mask.data.U8[y][x]:
                    psLib.subImage.data.F32[y][x] = 0.0


    if display['subtracted']:
        displayImage(subImage, 3, "Subtracted image", "red")

    if config.outFile:
        subFile = psLib.psFitsAlloc(config.outFile)
        psLib.psFitsWriteImage(subFile, None, subImage, 0, None)
        psTrace("pois", 2, ("Subtracted image written to %s\n" % (config.outFile)))

    psTrace("pois.time", 1, ("I/O completed at %f sec\n" % (getTime() - startTime)))

if not run.__doc__:
    run.__doc__ = _run.__doc__
