#!/usr/bin/env python
"""
$Header: /nfs/slac/g/glast/ground/cvs/pointlike/python/uw/pulsar/polyfold.py,v 1.1 2011/04/27 18:32:03 kerrm Exp $

Program: polyfold.py
Version: 1.3
Authors: Paul S. Ray <paul.ray@nrl.navy.mil>
          Matthew Kerr <matthew.kerr@gmail.com>
Description:
Reads an FT1 file with geocentered event times and
folds according to a polyco.dat file generated by tempo2 at the geocenter site ("coe")
Then, fits TOAs and outputs them in tempo2 format.
"""
from __future__ import division
import sys, os
from polyco import Polyco,PolycoEntry
from toagen import UnbinnedTOAGenerator,BinnedTOAGenerator,EDFTOAGenerator
from optparse import OptionParser
from lcfitters import LCTemplate
import numpy as np
import astropy.io.fits as pyfits

try:
    import psr_utils
    HAVE_PRESTO = True
except:
    print 'Unable to import presto packages; binned fit will be unavailable.'
    HAVE_PRESTO = False

__version__ = "1.3"


if __name__ == '__main__':

    desc="Read an FT1 file containing GEOCENTERED events (typically generated by gtbary) and fold according to a polyco.dat or .par file"
    parser=OptionParser(usage=" %prog [options] [FT1_FILENAME] [POLYCO_FILENAME/PAR_FILENAME]",
                                           description=desc)
    parser.add_option("-n","--ntoa",type="int",default=10,help="Number of TOAs to produce between TSTART and TSTOP.")
    parser.add_option("-p","--profile",action="store_true",default=False,help="Plot folded profile")
    parser.add_option("-b","--nbins",type="int",default=32,help="Number of bins in each profile.")
    parser.add_option("-e","--emin",type="float",default=0.0,help="Minimum energy to include.")
    parser.add_option("-x","--emax",type="float",default=300000.0,help="Maximum energy to include.")
    parser.add_option("-g","--gauss",type="string",default=None,help="File name containing gaussian template (from pygaussfit.py)")
    parser.add_option("-t","--template",type="string",default=None,help="File name containing template (e.g. .bestprof file)")
    parser.add_option("-u","--unbinned",action="store_true",default=False,help="Calculate TOAs with binned or unbinned method.")
    parser.add_option("-a","--addphase",action="store_true",default=False,help="Add PULSE_PHASE column to FT1 file using poly cos.")
    parser.add_option("-o","--output",type="string",default=None,help="File for output of .tim file.  Otherwise output to STDOUT.")
    parser.add_option("-e","--edf",action="store_true",default=False,help="Use EDF to calculated TOAs.  Must be specified with unbinned option.")
    parser.add_options("-r","--recalc_polycos",action="store_true",default=True,help="Recalc polycos")
    
    ## Parse arguments
    (options,args) = parser.parse_args()
    if len(args) != 2:
        parser.error("Both FT1_FILENAME and POLYCO_FILENAME arguments are required.")

    if (options.profile and (options.ntoa != 1)):
        print >>sys.stderr, "ERROR! ntoa MUST be 1 for profile generation!"

    if options.gauss is not None and options.template is not None:
        print >>sys.stderr, "ERROR! Specify either a gaussian template or a .bestprof type template -- but not both!"    

    #if options.gauss is None and options.unbinned:
    #    print >>sys.stderr, "ERROR! Unbinned likelihood requires a gaussian template."
    
    ft1name    = args[0]
    polyconame = args[1]

    poly = Polyco(polyconame,recalc_polycos=options.recalc_polycos)
    mjds,phases,met_converter,weights = process_ft1(ft1name,poly,options.emin,options.emax,write=options.addphase)

    if not options.unbinned:
        if not HAVE_PRESTO:
            raise ImportError,'Presto tools not found; binned fitting not available!'
        # Read template profile
        if (options.template is not None):
           template = psr_utils.read_profile(options.template, normalize=1)
        else:
           if (options.gauss is not None):
               template = psr_utils.read_gaussfitfile(options.gauss, options.nbins)
           else:
               print >>sys.stderr,  "WARNING: Using single gaussian template!"
               gaussianwidth=0.1 # DEFAULT
               template = psr_utils.gaussian_profile(options.nbins, 0.0, gaussianwidth)
           template = template / template.max()
    else:
        if options.gauss is not None:
            template = LCTemplate(template=options.gauss)
        elif options.template is not None:
            template = LCTemplate(template=options.template)
        elif options.edf:
            print 'EDF option selected; will use EDF for template.'
            template = None  
        else: # default is kernel density estimator
            print 'No template provided -- Please build one with itemplate.py and supply it!'
            sys.exit(1)

    if options.unbinned:
        if options.edf:
            tg = EDFTOAGenerator(mjds,phases,poly,met_converter,template)
        else:
            tg = UnbinnedTOAGenerator(mjds,phases,poly,met_converter,template)
    else:
        tg = BinnedTOAGenerator(mjds,phases,poly,met_converter,template,profile_only = options.profile)

    toas,err_toas,tim_strings = tg.get_toas(options.ntoa)

    if options.output is not None:
        f = open(options.output,'w')
        for s in tim_strings: f.write(s + '\n')
        f.close()
        print 'Wrote TOAs to %s'%(options.output)

    else:
        for s in tim_strings: print s
