"""
1/23/2008: Fixed typo preventing regridding when no axis rearrangment was
           needed.
"""

import taskutil
import os, shutil

# List for converting integers to Stokes parameters from Stokes.h.  Obviously
# the order should not be changed!
Stokes_list = ['Undefined',
               'I',  
               'Q',  
               'U',  
               'V',
               'RR', 
               'RL', 
               'LR', 
               'LL',
               'XX', 
               'XY', 
               'YX', 
               'YY',
               'RX', 
               'LX', 
               'LY',
               'XR', 
               'XL', 
               'YR', 
               'YL',
               'PP',          # General
               'PQ',          # quasiorthogonal
               'QP',          # polarization
               'QQ',          # correlations
               'RCircular',   # Single 'dish'
               'LCircular', 
               'Linear',
               'Ptotal',      # Polarized intensity ((Q^2+U^2+V^2)^(1/2))
               'Plinear',     # Linearly Polarized intensity ((Q^2+U^2)^(1/2))
               'PFtotal',     # Polarization Fraction (Ptotal/I)
               'PFlinear',    # Linear Polarization Fraction (Plinear/I)
               'Pangle']      # Linear Polarization Angle (0.5 arctan(U/Q)) (in radians)

def thorough_regrid(inim, template, outim, keep_inim_Stokes=False):
    """
    CASA uses a different order for the frequency and Stokes axes from the rest
    of the world.  This function swaps the freq and Stokes axes of a newly
    loaded FITS file (inim) into CASA's convention (template), makes the RA
    nonnegative, and puts the result in outim.

    The output image header will have the same Stokes parameter as the template
    image, unless keep_inim_Stokes is True.  No actual Stokes conversion is
    done (i.e. RR can get called I without being averaged with LL).  Even if
    keep_inim_Stokes is True, the output image will only have a Stokes axis if
    the template does.  If keep_inim_Stokes is True and template has a Stokes
    axis, but inim does not, outim will default with a warning to I.
    """
    myf         = taskutil.get_global_namespace()
    ia          = myf['ia']
    regridimage = myf['regridimage']
#    iatool = myf['iatool']

    ia.open(template)
    tsumm = ia.summary()['header']
    ia.close()

    taxesl = tsumm['axisnames'].tolist()
    tStokesind = index_or_minus1(taxesl, 'Stokes')
    tfreqind   = index_or_minus1(taxesl, 'Frequency')
        
    ia.open(inim)
    insumm = ia.summary()['header']
    ia.close()

    # Check degenerate axes
    inaxesl = insumm['axisnames'].tolist()
    inStokesind = index_or_minus1(inaxesl, 'Stokes')
    infreqind   = index_or_minus1(inaxesl, 'Frequency')
    tmpim = outim + '_tmp'
    if(inStokesind != tStokesind or infreqind != tfreqind):
        # inim's degenerate axes will have to come off for now if present.
        if keep_inim_Stokes:
            if(inStokesind >= 0):                        # Store the Stokes parameter.
                outStokes = Stokes_list[int(insumm['refval'][inStokesind])]
            elif tStokesind >= 0:
                print """
                Warning in thorough_regrid(inim = %s,
                                           template = %s,
                                           outim = %s,
                                           keep_inim_Stokes = True):
                    The template image has a Stokes axis (%s) but the input
                    image does not.  The output image will use I since
                    keep_inim_Stokes = True!
                    """ % (inim, template, outim,
                           Stokes_list[tsumm['refval'][inStokesind]])
                outStokes = 'I'
        if(infreqind >= 0):                          # Store the frequency.
            infreq     = insumm['refval'][infreqind]
            inreffreq  = insumm['refpix'][infreqind]
            innchan    = insumm['shape'][infreqind]
            infreqincr = insumm['incr'][infreqind]
        ia.newimagefromimage(inim, tmpim, dropdeg = True)
        if tStokesind >= 0:
            ia.open(tmpim)
            hasStokes = outim + '_has_Stokes'
            if keep_inim_Stokes == False:
                outStokes = Stokes_list[int(tsumm['refval'][tStokesind])]
            ib = ia.adddegaxes(outfile=hasStokes, stokes=outStokes, overwrite=True)
            ib.done()
            ia.done()
            shutil.rmtree(tmpim)
            os.rename(hasStokes, tmpim)
        if tfreqind >= 0:
            ia.open(tmpim)
            hasFreq = outim + '_has_Freq'
            ib = ia.adddegaxes(outfile=hasFreq, spectral = True, overwrite=True)
            ib.done()
            ia.done()
            shutil.rmtree(tmpim)
            os.rename(hasFreq, tmpim)
    else:
        shutil.copytree(inim, tmpim, True)

    regridimage(tmpim, template, outim)
    shutil.rmtree(tmpim)

def index_or_minus1(lis, item):
    """
    Returns the index of item in list lis, or -1 if it is not found.
    """
    try:
        ind = lis.index(item)
    except ValueError:
        ind = -1
    return ind