#! /usr/bin/env python
import os
import logging
import shutil
import ConfigParser
import subprocess
import glob
import tempfile
from optparse import OptionParser
from pycbc_glue.pipeline import CondorDAGJob, CondorDAGNode, CondorDAG, CondorJob


class BaseJob(CondorDAGJob, CondorJob):
    def __init__(self, log_dir, executable, cp, section, accounting_group=None):
        CondorDAGJob.__init__(self, "vanilla", executable)

        # These are all python jobs so need to pull in the env
        self.add_condor_cmd('getenv', 'True')
        self.add_condor_cmd('+Allow_OGrid', 'True')
        log_base = os.path.join(
            log_dir, os.path.basename(executable) + '-$(cluster)-$(process)')
        self.set_stderr_file(log_base + '.err')
        self.set_stdout_file(log_base + '.out')
        self.set_sub_file(section + ".sub")   

        if cp is not None: 
            self.add_ini_opts(cp, section)

        if accounting_group:
            self.add_condor_cmd('accounting_group', accounting_group)

class FaithsimNode(CondorDAGNode):
    def __init__(self, job, tmplt_file, match_file, inj_per_job=None):
        CondorDAGNode.__init__(self, job)    
        self.add_file_opt("param-file", tmplt_file)
        self.add_file_opt("match-file", match_file, file_is_output_file=True) 

def get_ini_opts(confs, section):
    op_str = ""
    for opt in confs.options(section):
        val = confs.get(section, opt)
        op_str += "--" + opt + " " + val + " \\" + "\n"
    return op_str
    
def mkdir(dir_name):
    try :
        os.mkdir(dir_name)
    except OSError:
        pass

def matches_in_list(slist, match):
    import re
    matches = []
    for st in slist:
        if st.startswith(match):
            matches.append(st)
    return matches

parser = OptionParser()
parser.add_option('--config', type=str)          
(options, args) = parser.parse_args() 

if options.config is None:
    raise ValueError("Config file is required")

logging.basicConfig(format='%(asctime)s %(message)s', level=logging.INFO)

confs = ConfigParser.ConfigParser()
confs.read(options.config)

banksim_prog = confs.get("executables", "faithsim")
templates_per_job = confs.get("workflow", "templates-per-job")

try:
    log_path = confs.get("workflow", 'log-path')
except:
    log_path = './'

tempfile.tempdir = log_path
tempfile.template='faithsim.dag.log.'
logfile = tempfile.mktemp()

try:
    accounting_group = confs.get('workflow', 'accounting-group')
except:
    accounting_group = None
    logging.warn('Warning: accounting-group not specified, LDG clusters may'
                 ' reject this workflow!')

logging.info("Making workspace directories")
mkdir('scripts')
mkdir('match')
mkdir('bank')
mkdir('log')
mkdir('plots')

logging.info("Copying scripts")
shutil.copy(banksim_prog, 'scripts/pycbc_faithsim')
os.chmod('scripts/pycbc_faithsim', 0777)

logging.info("Creating injection file")
inj_str = "lalapps_inspinj " + get_ini_opts(confs, "inspinj") + "--output inj.xml"
os.system(inj_str)

logging.info("Splitting template bank")
subprocess.call(['pycbc_splitbank',
                 '--templates-per-bank', str(templates_per_job),
                 '-t', 'inj.xml',
                 '-o', 'bank/bank'])

num_banks = len(glob.glob("bank/bank*"))

logging.info("Creating DAG")

dag = CondorDAG(logfile)
dag.set_dag_file("faithsim")
dag.set_dax_file("faithsim")

fs_secs = matches_in_list(confs.sections(), 'faithsim')
fsjobs = []
for sec in fs_secs:
    fsjobs.append(BaseJob("log", "scripts/pycbc_faithsim", confs, sec,
                          accounting_group=accounting_group))

rjob = BaseJob("log", "scripts/pycbc_faithsim_collect_results", None,
               'collect_results', accounting_group=accounting_group)
rnode = CondorDAGNode(rjob)
pjob = BaseJob("log", "scripts/pycbc_faithsim_plots", None, 'faithsim_plots',
               accounting_group=accounting_group)
pnode = CondorDAGNode(pjob)

for inj_num in range(num_banks):
    bn = 'bank/bank' + str(inj_num) + '.xml'
    
    for fsjob, sec in zip(fsjobs, fs_secs):
        sec_sub = str(sec[len('faithsim'):])
        mf = 'match/match' + sec_sub + '-' +  str(inj_num) + '.dat'
        fsnode = FaithsimNode(fsjob, bn, mf, inj_per_job=templates_per_job)
        dag.add_node(fsnode)
        rnode.add_parent(fsnode) 
dag.add_node(rnode)
pnode.add_parent(rnode)
dag.add_node(pnode)

f = open("scripts/pycbc_faithsim_collect_results", "w")
f.write("""#!/usr/bin/env python
from os.path import isfile
from numpy import *
from pycbc_glue.ligolw import utils, table, lsctables, ligolw
import glob

class ContentHandler(ligolw.LIGOLWContentHandler):
    pass
lsctables.use_in(ContentHandler)

fils = glob.glob("match/match*.dat")
mfields = ('match', 'overlap', 'time_offset', 'sigma1', 'sigma2')
bfields = ('match', 'overlap', 'time_offset', 'sigma1', 'sigma2', 'mass1',
           'mass2', 'spin1x', 'spin1y', 'spin1z', 'spin2x', 'spin2y',
           'spin2z', 'inclination', 'latitude', 'longitude',
           'polarization', 'coa_phase')
dtypem={'names': mfields, 'formats': ('f8', 'f8', 'f8', 'f8', 'f8')}
dtypeo={'names': bfields,
        'formats': ('f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8',
                    'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8')}
                    
if __name__ == "__main__":
    fils = glob.glob("match/match*.dat")
    tags = []
    for fil in fils:
        tags.append(fil.split('-')[1])
    tags = list(set(tags))
    btables = {}
    for tag in tags:
        parts = glob.glob("match/match-" + tag + "-*.dat")
        data = zeros(0, dtype=dtypeo)   
        for part in parts:
            btag = part.split('-')[2].split('.')[0]
            bname = "bank/bank" + btag + ".xml"
            if bname not in btables:
                indoc = utils.load_filename(bname, False, contenthandler=ContentHandler)
                btables[bname] = table.get_table(indoc, lsctables.SimInspiralTable.tableName) 
            bt = btables[bname]
            try:      
                md = loadtxt(part, dtype=dtypem)
                if md.size == 0:
                    continue  
            except IOError:
                continue
            pdata = zeros(len(bt), dtype=dtypeo)
            
            for field in mfields: 
                pdata[field] = md[field]
                
            for field in bfields:
                if field not in mfields:
                    pdata[field] = bt.get_column(field)
            
            data = append(data, pdata)
        savetxt('result-' + tag + '.dat', data)
""")
os.chmod('scripts/pycbc_faithsim_collect_results', 0777)
    
logging.info("Creating submit script")
f = open("submit.sh", 'w')
f.write("""#!/bin/bash
condor_submit_dag faithsim.dag
""")
os.chmod('submit.sh', 0777)

dag.write_sub_files()
dag.write_script()
dag.write_concrete_dag()

f = open('scripts/pycbc_faithsim_plots', "w")
f.write("""#!/usr/bin/env python
import matplotlib
matplotlib.use('Agg')
import matplotlib.cm
import pylab
from os.path import isfile
from numpy import *
from pycbc_glue.ligolw import utils, table, lsctables
from matplotlib.ticker import MultipleLocator, FormatStrFormatter
import glob

from pycbc import pnutils

bfields = ('match', 'overlap', 'time_offset', 'sigma1', 'sigma2', 'mass1',
           'mass2', 'spin1x', 'spin1y', 'spin1z', 'spin2x', 'spin2y',
           'spin2z', 'inclination', 'latitude', 'longitude',
           'polarization', 'coa_phase')
dtypeo={'names': bfields,
        'formats': ('f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8',
                    'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8')}
                    
def basic_scatter(out_name, xname, yname, title, xval, yval,  
                  cval=None, cname="", vmin=0, vmax=1, xmin=None, 
                  ymin=None, majorL=None, minorL=None):
    cmap = matplotlib.cm.jet
    cmap.set_under(color='gray')
    fig = pylab.figure(num=None, figsize=(10, 5))
    pylab.scatter(xval, yval, c=cval, linewidths=0, s=3, vmin=vmin, vmax=vmax, cmap=cmap, alpha=0.7)
    if cval is not None:
        bar = pylab.colorbar()
        bar.set_label(cname)
        
    pylab.xlabel(xname)
    pylab.ylabel(yname)
    
    if xmin is None:
        xmin = min(xval)
        
    if ymin is None:
        ymin = min(yval)
        
    pylab.xlim(xmin, max(xval))
    pylab.ylim(ymin, max(yval))
    
    ax = fig.gca()
    if majorL:
        ax.xaxis.set_major_locator(MultipleLocator(majorL))
        ax.yaxis.set_major_locator(MultipleLocator(majorL))

    if minorL:
        ax.xaxis.set_minor_locator(MultipleLocator(minorL))
        ax.yaxis.set_minor_locator(MultipleLocator(minorL))
    
    fol = 'plots/' 
    
    pylab.grid()
    pylab.title(title)
    pylab.savefig(fol + 'THUMB-' + out_name, dpi=50)
    pylab.savefig(fol + out_name, dpi=500)

fils = glob.glob("result*.dat")
for fil in fils:
    tag = fil.split('-')[1].split('.')[0]
    data = loadtxt(fil, dtype=dtypeo)

    mmass1 = maximum(data['mass1'], data['mass2'])
    mmass2 = minimum(data['mass1'], data['mass2'])
    mchirp, eta = pnutils.mass1_mass2_to_mchirp_eta(mmass1, mmass2)

    M = data['mass1'] + data['mass2']    
    q = maximum(data['mass1'] / data['mass2'], data['mass2'] / data['mass1'])
    s1 = (data['spin1x']**2 + data['spin1y']**2 + data['spin1z']**2)**0.5
    s2 = (data['spin2x']**2 + data['spin2y']**2 + data['spin2z']**2)**0.5
    
    s81 = data['sigma1'] / 8
    s82 = data['sigma2'] / 8
    
    pname = tag + '-'

    red_data = data[data['match'] > -0.5]
    
    combs =  [('mass1', 'mass2', 'match'),
             ]            
    for v1, v2, v3 in combs:
        name = pname + 'scatter' + v1 + '-' + v2 + '-' + v3
        title = tag
        basic_scatter(name, v1, v2, title, data[v1], data[v2], data[v3], v3, xmin=0, ymin=0)
        
    v1 = 'mass1'
    v2 = 'mass2'
    v3 = 'time_offset'
    name = pname + 'scatter' + v1 + '-' + v2 + '-' + v3
    title = tag
    basic_scatter(name, v1, v2, title, red_data[v1], red_data[v2], red_data[v3], v3, vmin=None, vmax=None)

    v1 = 'mass ratio'
    v1d = q
    v2 = 'spin 1 magnitude' 
    v2d = s1
    v3 = 'match'
    v3d = data['match']
    name = pname + 'scatter' + v1 + '-' + v2 + '-' + v3
    basic_scatter(name, v1, v2, tag, v1d, v2d, v3d, v3)


    v1 = 'mass ratio'
    v1d = q
    v2 = 'spin 1 z' 
    v2d = data['spin1z']
    v3 = 'match'
    v3d = data['match']
    name = pname + 'scatter' + v1 + '-' + v2 + '-' + v3
    basic_scatter(name, v1, v2, tag, v1d, v2d, v3d, v3)
    
    v1 = 'mass ratio'
    v1d = q
    v2 = 'spin 2 magnitude' 
    v2d = s2
    v3 = 'match'
    v3d = data['match']
    name = pname + 'scatter' + v1 + '-' + v2 + '-' + v3
    basic_scatter(name, v1, v2, tag, v1d, v2d, v3d, v3)


    v1 = 'mass ratio'
    v1d = q
    v2 = 'spin 2 z' 
    v2d = data['spin2z']
    v3 = 'match'
    v3d = data['match']
    name = pname + 'scatter' + v1 + '-' + v2 + '-' + v3
    basic_scatter(name, v1, v2, tag, v1d, v2d, v3d, v3)

    v1 = 'mass ratio'
    v1d = q
    v2 = 'total mass' 
    v2d = M
    v3 = 'match'
    v3d = data['match']
    name = pname + 'scatter' + v1 + '-' + v2 + '-' + v3
    basic_scatter(name, v1, v2, tag, v1d, v2d, v3d, v3)

    v1 = 'mass 1'
    v1d = mmass1
    v2 = 'mass 2' 
    v2d = mmass2
    v3 = 'match'
    v3d = data['match']
    name = pname + 'scatter' + v1 + '-' + v2 + '-' + v3
    basic_scatter(name, v1, v2, tag, v1d, v2d, v3d, v3, xmin=0, ymin=0, majorL=5, minorL=1)
    
    v1 = 'Coalescence Phase'
    v1d = data['coa_phase']
    v2 = 'Match' 
    v2d = data['match']
    v3 = 'Mchirp'
    v3d = mchirp
    name = pname + 'scatter' + v1 + '-' + v2 + '-' + v3
    basic_scatter(name, v1, v2, tag, v1d, v2d, v3d, v3, vmin=None, vmax=None, ymin=0)
    
    v1 = 'Inclination'
    v1d = data['inclination']
    v2 = 'Match' 
    v2d = data['match']
    v3 = 'Mchirp'
    v3d = mchirp
    name = pname + 'scatter' + v1 + '-' + v2 + '-' + v3
    basic_scatter(name, v1, v2, tag, v1d, v2d, v3d, v3, vmin=None, vmax=None)
    
    v1 = 'Inclination'
    v1d = data['inclination']
    v2 = 'Sigma2' 
    v2d = data['sigma2']
    v3 = 'Mchirp'
    v3d = mchirp
    name = pname + 'scatter' + v1 + '-' + v2 + '-' + v3
    basic_scatter(name, v1, v2, tag, v1d, v2d, v3d, v3, vmin=None, vmax=None)
         
    v1 = 'Inclination'
    v1d = data['inclination']
    v2 = 'Sigma1' 
    v2d = data['sigma1']
    v3 = 'Mchirp'
    v3d = mchirp
    name = pname + 'scatter' + v1 + '-' + v2 + '-' + v3
    basic_scatter(name, v1, v2, tag, v1d, v2d, v3d, v3, vmin=None, vmax=None)     
 """)
os.chmod('scripts/pycbc_faithsim_plots', 0777)

logging.info('Done')