"""
12/13/2007: Added logging,
                  inclusion of a component list if used,
                  and
                  the ability to use an image with a different name from
                  the project.
"""

from cfgsummary import *
from radplot import *
from taskinit import *
import taskutil
import os, re, shutil

def calc_fidelity(project = None, image = None, mod_w_complist = None):
    """
    Calculates both the image and uv plane fidelities for the specified
    simulation project as in
    http://www.alma.nrao.edu/memos/html-memos/alma398/memo398.pdf

    If project is given it can be used to find image and the component list (if
    any), assuming the information is in project.almasimmos.  Otherwise image
    and vis must be specified, and project will be extracted from image.

    Because almasimmos does not necessarily use the same modelimage as
    specified in its parameters when ignorecoords is True, the modelimage MUST
    be project.input_model!

    mod_w_complist is a input model image with the input component list added
    to it.  It will be used if specified, and otherwise generated if necessary.
    """
    casalog.origin('calc_fidelity')
    casalog.post("""calc_fidelity(project        = %s,
                                  image          = %s,
                                  mod_w_complist = %s)
                       starting.""" % (project, image, mod_w_complist))

    complist = None
    
    # Handle arguments.
    imextre = re.compile('\.im[^.]*$')
    if isinstance(project, str):
        try:
            siminps = os.popen('grep complist ' + project + '.almasimmos')
            complist = siminps.readline()
            siminps.close()
            complist = complist.split()[2].replace('"', '')
        except:
            casalog.post("Component list not found in %s" % project + '.almasimmos')
            complist = None

        if image == None:
            imroot = project
            if os.path.isdir(project + '.image'):
                image  = imroot + '.image'
            elif complist == None:
                casalog.post("No input image or complist found!", 'ERROR')
                return
        else:
            imroot = imextre.sub('', image)
    elif os.path.isdir(image):
        imext = imextre.search(image)
        if imext:
            # Use imextre instead of imext to get the $.
            project = imextre.sub('', image)
            imroot  = project
        else:
            casalog.post("Unsure how to extract project from \'%s\'" % image, 'ERROR')
            return
    else:
        casalog.post("No input image or project found!", 'ERROR')
        return
    
    myf = taskutil.get_global_namespace()
    ia  = myf['ia']

    ia.open(image)
    ims = ia.summary()

    # Yes, restoringbeam really is needed twice.
    clnbm = ims['header']['restoringbeam']['restoringbeam']

    imstats = ia.statistics() # Includes flux.

    ia.close() # Closes image

    mod = project + '.input_model'

    # If necessary, make a copy of the modelimage with complist included.
    if complist:
        if not mod_w_complist:
            # I would like to replace complist with its value, but that might
            # include /.
            mod_w_complist = mod + '_w_complist'
            try:
                mod_w_complist = im_plus_complist(mod, complist, mod_w_complist, ia)
            except:
                casalog.post("Error adding component list to input model image.",
                             'ERROR')
                return
        elif not os.path.isdir(mod_w_complist):
            casalog.post(mod_w_complist + " is not a valid directory!", 'ERROR')
            return
        else:
            casalog.post(
          "Using existing image " + mod_w_complist + " with included component list.",
            'WARN')
        mod = mod_w_complist
        
    # Convolve mod with clean beam.
    convmod = mod + '_conv_by_%.2gx%.2g' % (clnbm['major']['value'],
                                            clnbm['minor']['value'])
    ia.open(mod)
    if not os.path.isdir(convmod):
        ia.convolve2d(convmod, major = clnbm['major'], minor = clnbm['minor'],
                      pa = clnbm['positionangle'], overwrite = True)  

    # Make difference image.
    diffim = imroot + '.convmod-im'
    ia.imagecalc(diffim, "'%s' - '%s'" % (convmod, image), overwrite = True)
    ia.close()                            # Closes model image.
    
    imfidelity, diffstats   = fidel_im(convmod, diffim, imroot, ia)
    uvfidelity, cfg_summary = fidel_uv(convmod, diffim, project, imroot)

    casalog.post("""calc_fidelity(project        = %s,
                                  image          = %s,
                                  mod_w_complist = %s)
                       done.""" % (project, image, mod_w_complist))
    return imfidelity, diffstats, uvfidelity, cfg_summary

def fidel_im(convmod, diff, imroot, ia = None):
    """
    Calculates and returns the filename of the image fidelity for convmod and
    diff as in http://www.alma.nrao.edu/memos/html-memos/alma398/memo398.pdf

    One difference is that the median |value - median| is used instead of the
    rms to prevent division by 0.  The median |value - median| is close to the
    0.675 * the rms, but is less sensitive to outliers (i.e. catastrophic
    errors confined to a small region of the image).
    """
    # Get rms of diff.
    if not ia:
        myf = taskutil.get_global_namespace()
        ia = myf['ia']
    ia.open(diff)
    diffstats = ia.statistics(robust = True)

    posdiff = imroot + '_posdiff.im'
    ia.imagecalc(posdiff,
                 "max(abs('%s'), %f)" % (diff, 0.7 * diffstats['medabsdevmed']
                                         / 0.675))
    imfidelity = imroot + '_imfidelity.im'
    ia.imagecalc(imfidelity, "abs('%s') / '%s'" % (convmod, posdiff))
    ia.close()
    return imfidelity, diffstats

def fidel_uv(convmod, diff, project, imroot = None, longestbllinl = None):
    """
    Calculates and returns the filename of the uv fidelity for convmod and
    diff as in http://www.alma.nrao.edu/memos/html-memos/alma398/memo398.pdf

    The image will be chopped down to 2*longestbll on a side, where longestbll
    is returned as the longest baseline length.

    If longestbllinl, the longest baseline length in wavelengths, is not
    supplied, it will be looked for in project.almasimmos.
    """
    myf = taskutil.get_global_namespace()
    ia = myf['ia']
    qa = myf['qa']
    mul_scal_qa = myf['mul_scal_qa']

    if not longestbllinl:  # Get the maximum baseline length from the saved inputs.
        inpfile = os.popen("egrep 'antennalist|freq|chan' " + project + '.almasimmos | head -4')
        cfgfile   = inpfile.readline().split()[2].replace('"', '')
        nchan     = int(inpfile.readline().split()[2])
        startfreq = qa.quantity(inpfile.readline().split()[2].replace('"', ''))
        chanwidth = qa.quantity(inpfile.readline().split()[2].replace('"', ''))
        inpfile.close()
        centfreq = qa.add(startfreq, mul_scal_qa(0.5 * nchan, chanwidth))
        centfreq = qa.convert(centfreq, 'GHz')['value']
        conf_summary = cfgsummary(cfgfile, centfreq)
        longestbllinl =  conf_summary['longestbll'] * centfreq / 0.29978
        conf_summary['longestbllinl'] = longestbllinl
    
    convmodamp = uvamp(convmod, longestbllinl)
    diffamp = uvamp(diff, longestbllinl)
    ia.open(convmodamp)
    if not imroot:
        imroot = project
    uvfidelity = imroot + '_uvfidelity.im'
    ia.imagecalc(uvfidelity, "'%s' / '%s'" % (convmodamp, diffamp))
    ia.close()
    return uvfidelity, conf_summary

def uvamp(img, longestbll, outname = None):
    """
    FFTs the image with directory name img and returns the directory name of
    the resulting amplitudes.  That directory name will be outname if supplied,
    and img + '_amp' otherwise.  Each axis of the amplitude image will go from
    -longestbll to longestbll (which should be in world coordinates).
    """
    myf = taskutil.get_global_namespace()
    ia = myf['ia']
    pl = myf['pl']
    rg = myf['rg']
    
    ia.open(img)
    imgamp = img + '_ucamp.tmp'
    ia.fft(amp = imgamp)
    ia.close()
    ia.open(imgamp)
    summary = ia.summary()['header']
    longestbll /= summary['incr'][1]
    longestbll = int(pl.ceil(pl.absolute(longestbll)))
    refpix = summary['refpix'][0:2]
    blc = refpix - [longestbll, longestbll]
    trc = refpix + [longestbll, longestbll]
    if outname == None:
        imgcamp = img + '_amp'
    else:
        imgcamp = outname
    ia.subimage(imgcamp, rg.box(blc.tolist(), trc.tolist()), overwrite=True)
    ia.close()
    shutil.rmtree(imgamp)
    return imgcamp
    
def im_plus_complist(im, complist, im_w_complist = None, ia = None):
    """
    Returns an image equal to im + complist.  The name will be im_w_complist
    if specified, and im + '_w_complist' otherwise.
    """
    casalog.origin('im_plus_complist')
    if ia == None:
        myf = taskutil.get_global_namespace()
        ia  = myf['ia']
    if not im_w_complist:
        # I would like to replace complist with its value, but that might include /.
        im_w_complist = im + '_w_complist'
    if os.path.isdir(im_w_complist):
        casalog.post("Using existing " + im_w_complist, 'WARN')
    else:
        try:
            shutil.copytree(im, im_w_complist)
        except:
            casalog.post("Error '%s' copying %s to %s" % (repr(sys.exc_value), im,
                                                          im_w_complist), 'ERROR')
            raise RuntimeError, sys.exc_value
        try:
            ia.open(im_w_complist)
            cl.open(complist)                        # Thanks, Kumar!
            ia.modify(cl.torecord(), subtract=False)
            ia.close()
            cl.close()
        except:
            casalog.post("Error '%s' forming %s = %s + %s" % (repr(sys.exc_value),
                                                              im_w_complist, im,
                                                              complist), 'ERROR')
            raise RuntimeError, sys.exc_value
    return im_w_complist