#! /usr/bin/env python
#
#

import numpy as np
#from pylal.series import make_psd_xmldoc
#import pylal.series
import lal.series
import lal

class Object(object):
    pass

import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--fname-psd-ascii",default=None)
parser.add_argument("--prefix-out",default="",type=str)
parser.add_argument("--conventional-postfix",action='store_true')
parser.add_argument("--ifo",default="H1")
parser.add_argument("--asd",action='store_true', help="If specified, the input is an ASD. Square it.")
opts=  parser.parse_args()

dat = np.loadtxt(opts.fname_psd_ascii)
if opts.asd:
    dat[:,1] = dat[:,1]**2


##
## Zero pad as needed
##
if dat[0,0] > 0:
    f0 = dat[0,0]
    n_now = len(dat)
    df = dat[1,0] - dat[0,0]
    n_to_add = int(f0/df)
    print f0, df, n_to_add
    dat_new = np.zeros( (len(dat)+n_to_add,2))
    dat_new[n_to_add:,0] = dat[:,0]
    dat_new[n_to_add:,1] = dat[:,1]
    dat_new[0:n_to_add,0] = np.arange(n_to_add)*df
    dat_new[0:n_to_add,1] = dat[0,1]  # Fill it with a LARGE value
    dat = dat_new 


# dat_wrapper = Object()
# dat_wrapper.epoch = 0
# dat_wrapper.f0 = 0.0
# dat_wrapper.sampleUnits = "s strain^2"
# dat_wrapper.deltaF = dat[1,0] - dat[0,0]
# dat_wrapper.data = dat[:,1]
# dat_wrapper.name = opts.ifo

# convert structure to appropriate 
df = dat[1,0]-dat[0,0]
epoch = lal.LIGOTimeGPS(0.0)
psd_s = lal.CreateREAL8FrequencySeries(name=opts.ifo, epoch=epoch, f0=dat[0,0], deltaF=df, sampleUnits="s", length=len(dat))
psd_s.data.data=dat[:,1] # copy in data

psd_dict = {opts.ifo: psd_s}

xmldoc = lal.series.make_psd_xmldoc(psd_dict)
# Manually set the name
xmldoc.childNodes[0].attributes._attrs  = {"Name": "psd"}

from glue.ligolw import utils
fname_out = opts.ifo+opts.prefix_out+"_fromascii_psd.xml.gz"
if opts.conventional_postfix:
    fname_out = opts.ifo+"-psd.xml.gz"
utils.write_filename(xmldoc, fname_out, gz=True)

np.savetxt(opts.ifo+opts.prefix_out+"_fmax.dat",[np.max(dat[:,0]) - 2]) # add a 2Hz buffer 
