import glob
import taskutil

def summarize(imlist, collist = 'bunit, bmaj, bmin, bpa'):
    """
    For each image in imlist, prints the values of collist, and returns the
    output of ia.summary() for each image as a dictionary keyed by the image
    names.

    Except for the special values listed in the default of collist, its
    columns should be keys of ia.summary()'s output.

    imlist and collist can both be either single valued strings, comma
    separated strings, or lists of strings, i.e.

    summarize('robust0_nonoise.image',
              ['bmaj', 'bmin']);

    summarize(['robust0_nonoise.image', 'robust0_nonoise.residual'],
              ['bunit']);

    and
    
    summarize('robust0_nonoise.im*, robust0_nonoise.residual',
              'bmaj, bmin');

    are all valid.

    N.B.
    1. The above lines all end in ; to *skip the return value* in ipython.
    2. The items in imlist can have wildcards.

    Sample output:
    Image                       bunit     bmaj (mas)  bmin (mas)  bpa (deg)  
    input672GHz_50pc.image      Jy/pixel  
    robust0_nonoise.image       Jy/beam   28.8220365  25.5090762    97.6345  
    casa_robust0intclean.image            15.7047883  14.4392923   101.9130    
    """
    myf = taskutil.get_global_namespace()
    ia  = myf['ia']
    qa  = myf['qa']

    wcimlist = str_or_list_to_list(imlist)
    collist  = str_or_list_to_list(collist)

    imlist = []
    for fpat in wcimlist:              # Do filename matching.
        imlist += glob.glob(fpat)
    
    summary = {}
    maximw  = 0
    for im in imlist:
        ia.open(im)
        summary[im] = ia.summary()['header']
        ia.close()
        imw = len(im)
        if imw > maximw:
            maximw = imw

    if 'bmaj' in collist:
        fwhmunit = None
        for im in imlist:
            try:
                fwhmunit = summary[im]['restoringbeam']['restoringbeam']['major']['unit']
            except KeyError, e:
                if e == 'restoringbeam':
                    pass
            if fwhmunit:
                break
        if fwhmunit == 'arcsec':
            fwhmlab = '"'
        elif fwhmunit == 'arcmin':
            fwhmlab = "'"
        else:
            fwhmlab = fwhmunit

        # Check if each bmaj is < 1".
        usemas = True
        onearcsec = qa.quantity(1.0, 'arcsec')
        for im in imlist:
            try:
                if qa.ge(summary[im]['restoringbeam']['restoringbeam']['major'], onearcsec):
                    usemas = False
                break
            except KeyError, e:
                if e == 'restoringbeam':
                    pass            
        if usemas:
            fwhmunit = 'mas'
            fwhmlab  = 'mas'
    if 'bpa' in collist:
        paunit = ''          # None has no length.
        for im in imlist:
            try:
                paunit = summary[im]['restoringbeam']['restoringbeam']['positionangle']['unit']
                break
            except KeyError, e:
                if e == 'restoringbeam':
                    pass

    imfmt = "%%-%ds " % maximw
    print imfmt % "Image",
    colwidths = []
    colfmts   = []
    colunits  = []
    for col in collist:
        colw = len(col)
        colunit = None
        if col in ['bmaj', 'bmin']:
            colw += len(fwhmlab) + 3
            colunit = fwhmlab
            colfmt  = "%%#%d.%dg " % (colw, colw - 1)
        elif col == 'bpa':
            colw += len(paunit) + 3
            colunit = paunit
            colfmt  = "%% %d.%df " % (colw, colw - 5)
        elif col == 'bunit':
            for im in imlist:
                bunitl = len(summary[im]['unit'])
                if(bunitl > colw):
                    colw = bunitl
            colfmt = "%%-%ds " % colw
        else:
            colfmt = "%%-%ds " % colw            
        colwidths.append(colw)
        colunits.append(colunit)
        colfmts.append(colfmt)
        format = "%%-%ds " % colw
        if colunit:
            print format % (col + " ("  + colunit + ")"),
        else:
            print format % col,
    print ""

    for im in imlist:
        print imfmt % im,
        for col, colw, colunit, colfmt in zip(collist, colwidths, colunits, colfmts):
            if col == 'bunit':
                print colfmt % summary[im]['unit'],
            elif col == 'bmaj':
                try:
                    bmaj = qa.convert(summary[im]['restoringbeam']['restoringbeam']['major'],
                                      colunit)
                    print colfmt % bmaj['value'],
                except KeyError, e:
                    if e == 'restoringbeam':
                        blankfmt = "%%%ds " % colw
                        print blankfmt % ' ',
            elif col == 'bmin':
                try:
                    bmin = qa.convert(summary[im]['restoringbeam']['restoringbeam']['minor'],
                                      colunit)
                    print colfmt % bmin['value'],
                except KeyError, e:
                    if e == 'restoringbeam':
                        blankfmt = "%%%ds " % colw
                        print blankfmt % ' ',
            elif col == 'bpa':
                try:
                    bpa = qa.convert(summary[im]['restoringbeam']['restoringbeam']['positionangle'],
                                     colunit)
                    print colfmt % bpa['value'],
                except KeyError, e:
                    if e == 'restoringbeam':
                        blankfmt = "%%%ds " % colw
                        print blankfmt % ' ',
            else:
                print colfmt % summary[im][col],
        print ""

#    print "|col|\t|colw|\t|colunit|\t|colfmt|"
#    for col, colw, colunit, colfmt in zip(collist, colwidths, colunits, colfmts):
#        print "|%s|\t|%s|\t|%s|\t|%s|" % (col, colw, colunit, colfmt)
    
    return summary
        
def str_or_list_to_list(strorlist):
    """
    Returns a list made from splitting strorlist at commas and/or whitespace
    """
    if isinstance(strorlist, str):
        decommaed = strorlist.replace(',', ' ')
        return decommaed.split()
    return strorlist