#!/usr/bin/env python
#
# Copyright (C) 2015  Kipp Cannon, Chad Hanna
#
# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the
# Free Software Foundation; either version 2 of the License, or (at your
# option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General
# Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.

## @file
# A program to generate offline style web pages for an online analysis

"""Meta program to generate offline style summary pages from online runs."""


import sys, os, subprocess, re, time, glob
from optparse import OptionParser
from glue.text_progress_bar import ProgressBar
from pylal.datatypes import LIGOTimeGPS
import lal

def now():
	return LIGOTimeGPS(lal.UTCToGPS(time.gmtime()), 0)

parser = OptionParser()
parser.add_option("--directory", default = ".", help = "directory to start looking for results")
parser.add_option("--injection-file", default = ".", help = "The injection xml file that corresponds to the low latency injections")
parser.add_option("--web-dir", help = "set the output path to write the ''offline'' style web page to")
options, filenames = parser.parse_args()


# FIXME should be more clever than this
# Match 5 digit directories
dir_pattern = re.compile('[0-9]{5}')

# FIXME we really have to change this hacky convention for injections to start at 1000
non_inj_pattern = re.compile('.*-0[0-9]{3}_LLOID-.*.xml.gz')
inj_pattern = re.compile('.*-1[0-9]{3}_LLOID-.*.xml.gz')

noninj_files_to_merge = []
inj_files_to_merge = []

# FIXME assume that the correct low latency cluster file is in the working
# directory. Perhaps this should be a command line argument.
cluster_file = os.path.join(options.directory, "ll_simplify_and_cluster.sql")

result_dirs = sorted([d for d in os.listdir(options.directory) if dir_pattern.match(d)])
for n, d in enumerate(result_dirs):
	print >> sys.stderr, "processing directory %d of %d: %s" % (n+1, len(result_dirs), d)

	# FIXME assume first directory contains segment files (true because segments are cumulative)
	if n == 0:
		# FIXME assume the first mass bin contains the "correct" segments
		# Should be roughly true since segments are cumulative and each
		# node should see the same segments if the analysis is running
		# properly.
		seg_files = glob.glob('%s/*-0000_SEGMENTS*.xml.gz' % os.path.join(options.directory, d))

	# FIXME don't hard code H1L1
	# FIXME assumes 10 digit GPS
	noninjdb = os.path.join(os.path.join(options.directory, d), 'H1L1-ALL_LLOID-%s00000-100000.sqlite.tmp' % d)
	injdb = os.path.join(os.path.join(options.directory, d), 'H1L1-ALL_LLOID_INJ-%s00000-100000.sqlite.tmp' % d)
	noninjdbfinal = noninjdb.replace(".tmp","")
	injdbfinal = injdb.replace(".tmp","")

	# see if this directory has been processed and is old enough to not have to worry about it any more
	if float(now()) - float("%s00000" % d) > 125000 and os.path.exists(noninjdbfinal) and os.path.exists(injdbfinal):
		print >> sys.stderr, "directory is greater than 125000 seconds old and has already been processed...continuing"
		noninj_files_to_merge.append(noninjdbfinal)
		inj_files_to_merge.append(injdbfinal)
		continue

	# First do non injections
	files = sorted([os.path.join(os.path.join(options.directory, d), xml) for xml in os.listdir(os.path.join(options.directory, d)) if non_inj_pattern.match(xml)])
	progressbar = ProgressBar("Non injection files processed", len(files))
	for f in files:
		subprocess.check_call(["ligolw_sqlite", "--tmp-space", "/dev/shm", "--database", "%s" % noninjdb, "%s" % f])
		subprocess.check_call(["lalapps_run_sqlite", "--tmp-space", "/dev/shm", "--sql-file", cluster_file,  "%s" % noninjdb])
		progressbar.increment()
	del progressbar

	# Then injections
	files = sorted([os.path.join(os.path.join(options.directory, d), xml) for xml in os.listdir(os.path.join(options.directory, d)) if inj_pattern.match(xml)])
	progressbar = ProgressBar("Injection files processed", len(files))
	for f in files:
		subprocess.check_call(["ligolw_sqlite", "--tmp-space", "/dev/shm", "--database", "%s" % injdb, "%s" % f])
		subprocess.check_call(["lalapps_run_sqlite", "--tmp-space", "/dev/shm", "--sql-file", cluster_file,  "%s" % injdb])
		progressbar.increment()
	del progressbar


	# rename files
	for db,dbf in ((noninjdb, noninjdbfinal), (injdb, injdbfinal)):
		if os.path.exists(db):
			os.rename(db, dbf)

	if os.path.exists(noninjdbfinal) and os.path.exists(injdbfinal):
		noninj_files_to_merge.append(noninjdbfinal)
		inj_files_to_merge.append(injdbfinal)


# FIXME only add *new* files

#noninjdb = os.path.join(options.directory, 'H1L1-ALL_LLOID-%s00000-%d.sqlite.tmp' % (result_dirs[0], (int(result_dirs[-1]) - int(result_dirs[0])) * 100000))
#injdb = os.path.join(options.directory, 'H1L1-ALL_LLOID_INJ-%s00000-%d.sqlite.tmp' % (result_dirs[0], (int(result_dirs[-1]) - int(result_dirs[0])) * 100000))
noninjdb = os.path.join(options.directory, 'H1L1-ALL_LLOID-0-2000000000.sqlite.tmp')
injdb = os.path.join(options.directory, 'H1L1-ALL_LLOID_INJ-0-2000000000.sqlite.tmp')

if os.path.exists(noninjdb):
	os.remove(noninjdb)
if os.path.exists(injdb):
	os.remove(injdb)

progressbar = ProgressBar("Merge noninjection files", len(noninj_files_to_merge) + len(seg_files))
for f in noninj_files_to_merge + seg_files:
	subprocess.check_call(["ligolw_sqlite", "--tmp-space", "/dev/shm", "--database", "%s" % noninjdb, "%s" % f])
	subprocess.check_call(["lalapps_run_sqlite", "--tmp-space", "/dev/shm", "--sql-file", cluster_file,  "%s" % noninjdb])
	progressbar.increment()
del progressbar

progressbar = ProgressBar("Merge injection files", len(inj_files_to_merge) + len(seg_files))
for f in inj_files_to_merge + seg_files:
	subprocess.check_call(["ligolw_sqlite", "--tmp-space", "/dev/shm", "--database", "%s" % injdb, "%s" % f])
	subprocess.check_call(["lalapps_run_sqlite", "--tmp-space", "/dev/shm", "--sql-file", cluster_file,  "%s" % injdb])
	progressbar.increment()
del progressbar

# Find injections
progressbar = ProgressBar("Find injections", 4)
subprocess.check_call(["ligolw_sqlite", "--tmp-space", "/dev/shm", "--database", "%s" % injdb, "%s" % options.injection_file])
progressbar.increment()
subprocess.check_call(["ligolw_sqlite", "--tmp-space", "/dev/shm", "--database", "%s" % injdb, "--extract", "%s.xml" % injdb])
progressbar.increment()
subprocess.check_call(["ligolw_inspinjfind", "%s.xml" % injdb])
progressbar.increment()
subprocess.check_call(["ligolw_sqlite", "--tmp-space", "/dev/shm", "--database", "%s" % injdb, "--replace", "%s.xml" % injdb])
progressbar.increment()


# Make plots and such
if os.path.exists(os.path.join(options.directory, "plots")):
	os.rename(os.path.join(options.directory, "plots"), "%s.%s" % (os.path.join(options.directory, "plots"), str(now())))
os.mkdir(os.path.join(options.directory, "plots"))

pattern = re.compile("(?P<id>[0-9]{4})_prior.xml.gz")
for d in os.listdir(options.directory):
	m = pattern.match(d)
	if m:
		subprocess.check_call(["gstlal_inspiral_plot_background", "--output-dir", os.path.join(options.directory, "plots"), "--user-tag", m.group("id"), "--verbose", d])
subprocess.check_call(["gstlal_inspiral_plot_background", "--output-dir", os.path.join(options.directory, "plots"), "--user-tag", "ALL", "--database", noninjdb, "--verbose", "marginalized_likelihood.xml.gz"])
subprocess.check_call(["gstlal_inspiral_plotsummary", "--tmp-space", "/dev/shm", "--segments-name", "statevectorsegments", "--user-tag", "ALL_LLOID_COMBINED", "--output-dir", "%s" % os.path.join(options.directory, "plots"), "%s" % noninjdb, "%s" % injdb])

subprocess.check_call(["gstlal_inspiral_plot_sensitivity", "--output-dir", os.path.join(options.directory, "plots"), "--bin-by-chirp-mass", "--tmp-space", "/dev/shm",  "--zero-lag-database", noninjdb, "--dist-bins", "200", "--bin-by-total-mass",  "--user-tag", "ALL_LLOID_COMBINED", "--include-play",  "--bin-by-mass-ratio",  "--bin-by-mass1-mass2",  "--data-segments-name", "statevectorsegments", injdb])

subprocess.check_call(["gstlal_inspiral_summary_page", "--open-box", "--output-user-tag", "ALL_LLOID_COMBINED",  "--glob-path", "%s" % os.path.join(options.directory, "plots"), "--webserver-dir", options.web_dir, "--title", "gstlal-online"])


# copy the working files back
os.rename(noninjdb, noninjdb.replace(".tmp",""))
os.rename(injdb, injdb.replace(".tmp",""))
