#! /usr/bin/env python
import os
import logging
import shutil
import configparser as ConfigParser
import subprocess
import glob
import tempfile
from optparse import OptionParser
from 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)

        self.add_condor_cmd('request_disk', 1024)

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', 0o0777)

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")

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.gz'
    
    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 numpy import *
from ligo.lw import utils, table, lsctables
import glob
from pycbc.io.ligolw import LIGOLWContentHandler


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.gz"
            if bname not in btables:
                indoc = utils.load_filename(bname, False, contenthandler=LIGOLWContentHandler)
                btables[bname] = lsctables.SimInspiralTable.get_table(indoc)
            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', 0o0777)
    
logging.info("Creating submit script")
f = open("submit.sh", 'w')
f.write("""#!/bin/bash
condor_submit_dag faithsim.dag
""")
os.chmod('submit.sh', 0o0777)

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 numpy import *
from ligo.lw 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', 0o0777)

logging.info('Done')