#!/usr/bin/env python

import argparse
import glob
import pyfits

import dvo


def main(filename):
    astrom = dvo.readAstrom(filename)
    photom = dvo.readPhotom(filename, astrom)
    selection = dvo.select(astrom, photom)
    num = selection.sum()

    schema = astrom.columns + photom.columns
    out = pyfits.BinTableHDU.from_columns(schema, nrows=num)
    for col in astrom.columns:
        out.data.field(col.name)[:] = astrom.data.field(col.name)[selection]
    for col in photom.columns:
        out.data.field(col.name)[:] = photom.data.field(col.name)[selection]

    name = filename.replace(".cpt", ".extract").replace(".cps", ".extract")
    out.writeto(name, clobber=True)
    print "Wrote %d rows to %s" % (num, name)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("filename", nargs="+", help="Name of file(s) to read, or glob")
    args = parser.parse_args()
    for ff in sum([glob.glob(fn) for fn in args.filename], []):
        main(ff)
