#!/usr/bin/env python
# Program: upolyfold.py
# 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
from optparse import OptionParser
import numpy as np

import os
from os.path import exists

from uw.pulsar.phasedata import PhaseData
from uw.pulsar.toabinner import TOABinner,UniformLogBinner
from uw.pulsar.polyco import Polyco
from uw.pulsar.toagen import UnbinnedTOAGenerator
from uw.pulsar.lctemplate import LCTemplate,prim_io

desc="""Read an FT1 file containing GEOCENTERED events (typically generated by gtbary) and fold according to a polyco.dat or .par file.
Note that if you give it a .par file it will STILL use the most recent polyco_new.dat UNLESS you specify -r|--recalc-polycos!
This means it is convenient to ALWAYS give it the .par file name and just use -r when you want to recalc."""

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("-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("-p","--plot",type="string",default=None,help="Set to base name of files to generate likelihood surface plots for each TOA.")
parser.add_option("","--wmin",type="float",default=0.0,help="Minimum weight value to include.")
parser.add_option("","--wmax",type="float",default=1.0,help="Maximum weight value to include.")
parser.add_option("-t","--template",type="string",default=None,help="File name containing template (LCTemplate compatible)")
parser.add_option("-w","--weights",type="string",default=None,help="Specify column of FT1 file to use as weights for TOA computation")
parser.add_option("-r","--recalc-polycos",action="store_true",default=False,help="Recompute polycos using Tempo2")
parser.add_option("-u","--uniform_sigma",action="store_true",default=False,help="Instead of fixed time bins, make a TOA each time the H-test gets to a sufficient sigma.")
parser.add_option("-a","--addphase",action="store_true",default=False,help="Add PULSE_PHASE column to FT1 file using poly cos.")
parser.add_option("","--blind",action="store_true",default=False,help="Force blind search for TOAs rather than tracking.")
parser.add_option("-o","--output",type="string",default=None,help="File for output of .tim file.  Otherwise output to STDOUT.")

## Parse arguments
(options,args) = parser.parse_args()
if len(args) != 2:
    parser.error("Both FT1_FILENAME and POLYCO_FILENAME arguments are required.")

#if options.gauss is None and options.unbinned:
#    print >>sys.stderr, "ERROR! Unbinned likelihood requires a gaussian template."

ft1name    = args[0]
polyconame = args[1]

# To be safe, if the par file is newer than polyco_new.dat, force remake the polycos, even if
# the user forgot to specify -r
if polyconame.endswith(".par"):
    if not exists("polyco_new.dat"):
        print "WARNING: No polyco_new.dat!  Regenerating...."
        options.recalc_polycos = True
    else:
        if os.stat(polyconame).st_mtime > os.stat("polyco_new.dat").st_mtime:
            print "WARNING: par file provided is newer than polyco_new.dat! Regenerating polyco_new.dat!"
            options.recalc_polycos = True

poly = Polyco(polyconame,recalc_polycos=options.recalc_polycos)
data = PhaseData(ft1name,poly,use_weights=(options.weights is not None),we_col_name=options.weights,wmin=options.wmin,emin=options.emin,emax=options.emax)
if options.addphase: data.write_phase()

primitives,norms = prim_io(options.template)
template = LCTemplate(primitives,norms)
if options.weights is not None:
    logl = np.log(1+data.weights*(template(data.ph)-1)).sum()
else: logl = np.log(template(data.ph)).sum()
print 'Total log likelihood:  %.2f'%(logl)

tg = UnbinnedTOAGenerator(data,poly,template,plot_stem=options.plot,good_ephemeris=(not options.blind))

if options.uniform_sigma:
    binner = UniformLogBinner(options.ntoa,data,template)
else:
    binner = TOABinner(options.ntoa,data)
    
toas,err_toas,tim_strings = tg.get_toas(binner)

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