#!/usr/bin/env python
#
# $Id: oicopy,v 1.4 2004/01/05 13:05:53 jsy1001 Exp $

"""Copy from one OI Exchange Format FITS file to another.

Usage: oicopy [--target_id n] [--arrname array] [--overwrite] [--flag vis/vis2/t3amp/t3] infile outfile

-t/--target_id:  only copies records for specified TARGET_ID

-a/--arrname:    only copies tables with specified ARRNAME

-o/--overwrite:  overwrite output file

-f/--flag:       flag all data of specified type (vis/vis2/t3amp/t3)

Other switches will be added later.

"""

__revision__ = '$Revision: 1.4 $'

import sys, getopt, os, string
import oi_fits, pyfits, ieee


def readfile(fileName):
    """Read FITS file.

    Returns OI_FITS instance, or None on failure.

    """
    hlist = pyfits.open(fileName)
    oi = oi_fits.OI_FITS()
    try:
        oi.FromHDUList(hlist)
    except oi_fits.OI_FITSConformError:
        hlist.close()
        return
    hlist.close()
    return oi


def writefile(oi, fileName):
    """Write FITS file from OI_FITS instance."""
    oi.ToHDUList().writeto(fileName)


def filtertargets(oi, target_ids):
    """Remove records with target_ids not in specified list.

    Removes empty data tables afterwards.

    """
    for id in oi.target.records.keys():
            if id not in target_ids:
                del oi.target.records[id]
    if len(oi.target.records) == 0:
        print '%s: no match' % target_ids
        return
    for list in (oi.vis, oi.vis2, oi.t3):
        for table in list[:]: # [:] makes a copy
            for rec in table.records[:]:
                if rec.target_id not in target_ids:
                    table.records.remove(rec)
            if len(table.records) == 0: list.remove(table)
    return 1


def filterarray(oi, arrname):
    """Remove OI_ARRAY/data tables with arrname != specified value.

    Removes redundant OI_WAVELENGTH tables. Leaves instance unchanged and
    returns false if no OI_ARRAY table with specified ARRNAME.

    """
    try:
        oi.arrays = {arrname: oi.arrays[arrname]}
    except KeyError:
        print '%s: no match' % arrname
        return
    insnames = []
    for list in (oi.vis, oi.vis2, oi.t3):
        for table in list[:]: # [:] makes a copy
            if table.arrname == arrname:
                insnames.append(table.insname)
            else:
                list.remove(table)
    for ins in oi.wavelengths.keys():
        if ins not in insnames: del oi.wavelengths[ins]
    return 1


def flag(oi, type):
    """Flag all data of specified type.

    type argument should be string 'vis', 'vis2', 't3amp', or 't3'
    (the Exchange Format does not accommodate flagging of closure
    phases but not the corresponding triple amplitudes). NULL
    amplitudes are assigned if 't3amp' passed.

    """
    
    if type[:2] == 't3':
        tabtype = 't3'
    else:
        tabtype = type
    try:
        list = oi.__dict__[tabtype]
    except KeyError:
        print 'Invalid data type:', type
        return
    for table in list:
        for rec in table.records:
            if type == 't3amp':
                rec.t3amp = [ieee.NAN]*table.nwave
            else:
                rec.flag = [1]*table.nwave
    return 1

def print_revision():
        print '%s v%s' % (sys.argv[0], string.split(__revision__)[1])


def _main():
    # Set defaults
    target_id = None
    arrname = None
    overwrite = 0
    flagtype = None
    # Parse command-line arguments
    try:
        opts, args = getopt.getopt(sys.argv[1:], "ht:a:f:o",
                                   ["target_id=", "arrname=", "flag=",
                                    "overwrite", "help"])
    except getopt.GetoptError:
        print_revision()
        print __doc__
        sys.exit(2)
    if len(args) != 2:
        print_revision()
        print __doc__
        sys.exit(2)
    for o, a in opts:
        if o in ("-h", "--help"):
            print_revision()
            print __doc__
            sys.exit()
        if o in ("-t", "--target_id"):
            target_id = int(a)
        if o in ("-a", "--arrname"):
            arrname = a
        if o in ("-o", "--overwrite"):
            overwrite = 1
        if o in ("-f", "--flag"):
            flagtype = a
    inputName, outputName = args
    # Perform copy...
    # ...read input file
    oi = readfile(inputName)
    # ...filter
    if arrname is not None:
        print 'Selecting ARRNAME', arrname
        if not filterarray(oi, arrname):
            sys.exit(1)
    if target_id is not None:
        print 'Selecting TARGET_ID', target_id
        if not filtertargets(oi, [target_id]):
            sys.exit(1)
    if flagtype is not None:
        print 'Flagging', flagtype
        if not flag(oi, flagtype):
            sys.exit(1)
    # ...write output file
    if os.path.exists(outputName):
        if overwrite:
            os.remove(outputName)
        else:
            print outputName, 'exists. Aborting'
            sys.exit(1)
    writefile(oi, outputName)


if __name__ == '__main__':
    _main()