#!/usr/bin/env python
from __future__ import print_function

import sys,re
import argparse
import numpy as np
import pandas as pd
from keputils import koiutils as ku


def main():
    parser = argparse.ArgumentParser(description='prints list of KOIs satisfying given query')
    
    parser.add_argument('-f','--file',default=None,help='super-list of KOIs to query.')
    parser.add_argument('-Q','--query',default='',help='query')
    parser.add_argument('--stars',action='store_true',help='print names as star names rather than KOI names')
    parser.add_argument('-I','--info',action='store_true',help='print basic info (P,Rp) for each koi (will just be .01 if --stars option is set)')
    parser.add_argument('-p','--pars',nargs='*',default=None,help='include other parameters in output table')

    args = parser.parse_args()

    if args.file is not None:
        kois = pd.Series(np.loadtxt(args.file,dtype=str)).apply(ku.koiname)
        df = ku.DATA.ix[kois]
    else:
        df = ku.DATA

    subdf = df.query(args.query)

    kois = subdf['kepoi_name']
    if args.stars:
        kois = kois.apply(ku.koistar)
    kois = np.sort(kois.unique())

    if args.info:
        printpars = ['koi_prad','koi_period']
    else:
        printpars = []

    if args.pars is not None:
        printpars += args.pars

    for k in kois:
        sys.stdout.write(k)
        for p in printpars:
            sys.stdout.write(' {}'.format(subdf.ix[k,p]))
        sys.stdout.write('\n')

if __name__=='__main__':
    main()
