#! /usr/bin/env python
import logging
import os
import shutil
from six.moves 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, gpu=False,
                 accounting_group=None):
        CondorDAGJob.__init__(self, "vanilla", executable)

        if gpu:
            CondorJob.__init__(self, "vanilla", executable, 2)
        # These are all python jobs so need to pull in the env
        self.add_condor_cmd('getenv', '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(os.path.basename(executable) + '.sub')

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

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

class BanksimNode(CondorDAGNode):
    def __init__(self, job, inj_file, tmplt_file, match_file, gpu=True,
                 gpu_postscript=False, inj_per_job=None):
        CondorDAGNode.__init__(self, job)

        self.add_file_opt("signal-file", inj_file)
        self.add_file_opt("template-file", tmplt_file)

        if gpu:
            self.add_var_opt("processing-scheme", 'cuda')

        if gpu and gpu_postscript:
            self.set_retry(5)
            mf = match_file+".$(Process)"
            mf1 = match_file+".0"
            mf2 = match_file+".1"
            self.add_file_opt("match-file", match_file+".$(Process)",
                              file_is_output_file=True)
            self.job().__queue = 2

            # Needed to satisfy the requirements for both running on atlas and spice
            job.add_condor_cmd('+WantsGPU', 'true')
            job.add_condor_cmd('+WantGPU', 'true')
            job.add_condor_cmd(
                'Requirements',
                '(GPU_PRESENT =?= true) || (HasGPU =?= "gtx580")')

            self.set_post_script(gpu_postscript)
            self.add_post_script_arg(mf1)
            self.add_post_script_arg(mf2)
            self.add_post_script_arg(".0001")
            self.add_post_script_arg(match_file)
            self.add_post_script_arg(str(inj_per_job))
        else:
            self.add_file_opt("match-file", match_file, file_is_output_file=True)
        
class CombineNode(CondorDAGNode):
    def __init__(self, job, inj_num):
        CondorDAGNode.__init__(self, job)
        
        self.add_var_opt("inj-num", inj_num)
        
        outf = "match/match" + str(inj_num) + ".dat"
        
        self.add_file_opt("output-file", outf)    

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 mc_min_max_from_sorted_file(fname):
    from ligo.lw.utils import load_filename
    from ligo.lw.table import Table
    from pycbc.pnutils import mass1_mass2_to_mchirp_eta
    from pycbc.io.ligolw import LIGOLWContentHandler

    doc = load_filename(fname, False, contenthandler=LIGOLWContentHandler)
    try:
        t = Table.get_table(doc, "sngl_inspiral")
    except:
        t = Table.get_table(doc, "sim_inspiral")
    mc_max, et = mass1_mass2_to_mchirp_eta(t[0].mass1, t[0].mass2)
    mc_min, et = mass1_mass2_to_mchirp_eta(t[-1].mass1, t[-1].mass2)
    return mc_min, mc_max
    
        
bf_mchirps = {}
sf_mchirps = {}
def check_outside_mchirp(bf, sf, w):
    if bf not in bf_mchirps:
        bf_mchirps[bf] = mc_min_max_from_sorted_file(bf)    
    if sf not in sf_mchirps:
        sf_mchirps[sf] = mc_min_max_from_sorted_file(sf) 
         
    mc_min, mc_max = bf_mchirps[bf]
    mc2_min, mc2_max =  sf_mchirps[sf]
    
    if (mc_min  <= mc2_max * (1+w) ) and (mc_max * (1+w) >= mc2_min):
        return False
    else:
        return True

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", "banksim")
bank_file = confs.get("workflow", "bank-file")
injections_per_job = confs.get("workflow", "injections-per-job")
templates_per_job = confs.get("workflow", "templates-per-job")

log_path = confs.get("workflow", 'log-path')

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

mchirp_window = None
if confs.has_option("banksim", "mchirp-window"):
    if ',' in confs.get("banksim", "mchirp-window"):
         mchirp_window = max([float(x) for x in confs.get("banksim", "mchirp-window").split(",")])
    else:
         mchirp_window = float(confs.get("banksim", "mchirp-window"))

gpu = False
try:
    gpu = confs.get("workflow", "use-gpus")
    if gpu is not None:
        gpu = True
except:
    pass

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('bank')
mkdir('match')
mkdir('injection')
mkdir('match-part')
mkdir('log')
mkdir('plots')

logging.info("Copying scripts")
shutil.copy(banksim_prog, 'scripts/pycbc_banksim')
os.chmod('scripts/pycbc_banksim', 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', bank_file,
                 '-o', 'bank/bank',
                 '--sort-mchirp'])

logging.info("Splitting injection file")
subprocess.call(['pycbc_splitbank',
                 '--templates-per-bank', str(injections_per_job),
                 '-t', "inj.xml",
                 '-o', 'injection/injection',
                 '--sort-mchirp'])

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

logging.info("Creating DAG")
f = open("banksim.dag", "w")

do_count = 0
skip_count = 0

dag = CondorDAG(logfile)
dag.set_dag_file("banksim")

bsjob = BaseJob("log", "scripts/pycbc_banksim", confs, "banksim", gpu=gpu,
                accounting_group=accounting_group)
cjob = BaseJob("log", "scripts/pycbc_banksim_match_combine", None, None,
               accounting_group=accounting_group)
rjob = BaseJob("log", "scripts/pycbc_banksim_collect_results", None, None,
               accounting_group=accounting_group)
pjob = BaseJob("log", "scripts/pycbc_banksim_plots", None, None,
               accounting_group=accounting_group)
rnode = CondorDAGNode(rjob)
pnode = CondorDAGNode(pjob)

for inj_num in range(num_injs):
    num = str(inj_num)
    combine_has_jobs = False
    cnode = CombineNode(cjob, inj_num)
    for bank_num in range(num_banks):
        if mchirp_window is not None:
            bank_part = "bank/bank" + str(bank_num) + ".xml"
            sim_part =  "injection/injection" + str(inj_num) + ".xml"
            if check_outside_mchirp(bank_part, sim_part, mchirp_window):
                skip_count += 1
                continue
            else:
                do_count += 1
        part_num = str(bank_num)
        mfn = 'match-part/match' + num +'part' + part_num + '.dat'
        sn = 'injection/injection' + num + '.xml'
        bn = 'bank/bank' + part_num + '.xml'
        bsnode = BanksimNode(bsjob, sn, bn, mfn, gpu=gpu,
                             gpu_postscript="scripts/diff_match.sh",
                             inj_per_job=injections_per_job)
        cnode.add_parent(bsnode)
        dag.add_node(bsnode)
        combine_has_jobs = True
    if combine_has_jobs:
        rnode.add_parent(cnode)
        dag.add_node(cnode)      
dag.add_node(rnode)
pnode.add_parent(rnode)
dag.add_node(pnode)

logging.info("DO : %d SKIP %d" %(do_count, skip_count))
f.close()

f = open("scripts/pycbc_banksim_match_combine", "w")
f.write("""#!/usr/bin/env python
from os.path import isfile
from optparse import OptionParser
from numpy import *
import glob
parser = OptionParser()

parser.add_option('--inj-num',help="index of the injection set for the match files",type=int)
parser.add_option('-o','--output-file',help="output file with the maximized values")
options, argv_frame_files = parser.parse_args()

fils = glob.glob("match-part/match"+str(options.inj_num)+"part*.dat")

dtypef={'names': ('match', 'bank', 'bank_i', 'sim', 'sim_i', 'sigmasq'), 'formats': ('f8', 'S256', 'i4', 'S256', 'i4', 'f8')}

matches=[]
maxmatch = []
for fil in fils:
    matches.append(loadtxt(fil, dtype=dtypef))
   
indices = array(matches, dtype=dtypef)['match'].argmax(0)
for i, j in enumerate(indices):
    maxmatch.append(matches[j][i])
    
maxmatch=array(maxmatch, dtype =dtypef)
savetxt(options.output_file, maxmatch,fmt=('%5.5f', '%s', '%i', '%s', '%i', '%5.5f'), delimiter=' ')
""")
os.chmod('scripts/pycbc_banksim_match_combine', 0o0777)

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

from pycbc.io.ligolw import LIGOLWContentHandler

fils = glob.glob("match/match*.dat")

dtypem={'names': ('match', 'bank', 'bank_i', 'sim', 'sim_i', 'sigmasq'), 'formats': ('f8', 'S256', 'i4', 'S256', 'i4', 'f8')}

# Collect the results
res = None
for fil in fils:
    if res is not None:
        res = append(res, loadtxt(fil, dtype=dtypem))
    else:
        res = loadtxt(fil, dtype=dtypem)
 
btables = {}
itables = {}     

f = open("results.dat", "w")
for row in res: 
    outstr = ""
    if row['bank'] not in btables:
        indoc = utils.load_filename(row['bank'], False,
                                    contenthandler=LIGOLWContentHandler)
        btables[row['bank']] = table.get_table(indoc, "sngl_inspiral") 

    if row['sim'] not in itables:
        indoc = utils.load_filename(row['sim'], False,
                                    contenthandler=LIGOLWContentHandler)
        itables[row['sim']] = table.get_table(indoc, "sim_inspiral") 
    
    bt = btables[row['bank']][row['bank_i']]     
    it = itables[row['sim']][row['sim_i']]
 
    outstr += str(row['match']) + " "
    outstr += str(bt.mass1) + " "
    outstr += str(bt.mass2) + " "
    outstr += str(bt.spin1x) + " "
    outstr += str(bt.spin1y) + " "
    outstr += str(bt.spin1z) + " "
    outstr += str(bt.spin2x) + " "
    outstr += str(bt.spin2y) + " " 
    outstr += str(bt.spin2z) + " "
    
    outstr += str(it.mass1) + " "
    outstr += str(it.mass2) + " "
    outstr += str(it.spin1x) + " "
    outstr += str(it.spin1y) + " "
    outstr += str(it.spin1z) + " "
    outstr += str(it.spin2x) + " "
    outstr += str(it.spin2y) + " " 
    outstr += str(it.spin2z) + " "
    
    outstr += str(it.coa_phase) + " "
    outstr += str(it.inclination) + " "
    outstr += str(it.latitude) + " " 
    outstr += str(it.longitude) + " "
    outstr += str(it.polarization) + " "
    
    outstr += str(row['sigmasq']) + " "
                
    outstr += "\\n"
    
    f.write(outstr)
""")
os.chmod('scripts/pycbc_banksim_collect_results', 0o0777)

if gpu:
    f = open("cconfig", "w")
    f.write("""
    DAGMAN_PROHIBIT_MULTI_JOBS = False
    """)

    f = open("scripts/diff_match.sh", "w")
    f.write("""#!/bin/bash

    len=`cat $1 | wc -l`
    len2=`cat $2 | wc -l`

    if [ $len -eq $len2 ] && [ $len -ne 0 ] ; then
       echo "correct length"
    else
       echo "wrong length file"
        exit 1
    fi  

    function fuzzy_diff {
       echo  " ($3>($1-$2)) && ($3>($2-$1)) " | bc  
    }

    exec 3<$1
    exec 4<$2

    while IFS= read -r line1 <&3
    IFS= read -r line2 <&4
    do
        line1=`echo "$line1" | cut --delimiter=' ' -f 1`
        line2=`echo "$line2" | cut --delimiter=' ' -f 1` 

        if ! [[ "$line1" =~ ^[0-9]+([.][0-9]+)?$ ]] ; then
           exec >&2; echo "error: Not a number"; exit 1
        fi
        
        if ! [[ "$line2" =~ ^[0-9]+([.][0-9]+)?$ ]] ; then
           exec >&2; echo "error: Not a number"; exit 1
        fi

        ok=`fuzzy_diff $line1 $line2 $3`

        if  [ $ok -eq 0 ] ; then
           echo "Files do not match"
           exit 1
        fi 
       
    done


    cp $1 $4
    cp $1.found $4.found
    echo "The files are close enough"

    exit 0
    """)
    os.chmod('scripts/diff_match.sh', 0o0777)
    
logging.info("Creating submit script")
f = open("submit.sh","w")
if gpu:
    f.write("""#!/bin/bash
    condor_submit_dag -config cconfig banksim.dag
    """)
else:
    f.write("""#!/bin/bash
    condor_submit_dag banksim.dag
    """)
os.chmod('submit.sh', 0o0777)

f = open("partial_results.sh", "w")
f.write("""#!/bin/bash
scripts/pycbc_banksim_collect_results
""")
os.chmod('partial_results.sh', 0o0777)

dag.write_sub_files()
dag.write_script()
dag.write_dag()

f = open("scripts/pycbc_banksim_plots", "w")
f.write("""#!/usr/bin/env python
import matplotlib
matplotlib.use('agg')
from matplotlib import pyplot
from pycbc import pnutils
import numpy
import pylab

goldenratio = 2 / (1 + 5**.5)
#matplotlib.rcParams.update({
#        "font.size": 8.0,
#        "axes.titlesize": 8.0,
#        "axes.labelsize": 8.0,
#        "xtick.labelsize": 8.0,
#        "ytick.labelsize": 8.0,
#        "legend.fontsize": 8.0,
#	"figure.figsize": (3.3,3.3*goldenratio),
#        "figure.dpi": 200,
#	"subplots.left": 0.2, 
#	"subplots.right": 0.75, 
#	"subplots.bottom": 0.15,
#	"subplots.top": 0.75,
#        "savefig.dpi": 600,
#        "text.verticalalignment": "center",
#})

res = numpy.loadtxt("results.dat")
match = res[:,0]

tmass1 = res[:,1]
tmass2 = res[:,2]
tspin1x = res[:,3] 
tspin1y = res[:,4]
tspin1z = res[:,5]
tspin2x = res[:,6]
tspin2y = res[:,7]
tspin2z = res[:,8]
tmchirp, teta = pnutils.mass1_mass2_to_mchirp_eta(tmass1, tmass2)

imass1 = res[:,9]
imass2 = res[:,10]
ispin1x = res[:,11]
ispin1y = res[:,12]
ispin1z = res[:,13]
ispin2x = res[:,14]
ispin2y = res[:,15]
ispin2z = res[:,16]
imchirp, ieta = pnutils.mass1_mass2_to_mchirp_eta(imass1, imass2)

coa_phase = res[:,17]
inclination = res[:,18]
latitude  = res[:,19]
longitude = res[:,20]
polarization = res[:,21]

sigmasq = res[:,22]

q = numpy.maximum(imass1/imass2, imass2/imass1)
s1m = (ispin1x**2+ispin1y**2+ispin1z**2)**0.5
s2m = (ispin2x**2+ispin2y**2+ispin2z**2)**0.5

def mhist(c1, name, cum=False, normed=True, log=False, bins=100, xl="", yl=""):
    pylab.figure()
    pylab.xlabel(xl)
    pylab.ylabel(yl)
    if log:
        pylab.yscale('log')
    pylab.hist(c1, bins=bins, normed=normed, histtype='step', cumulative=cum)
    pylab.savefig(name)

def mplot(c1, c2, c, name, xl="", yl="", vmin=None, vmax=None):
    pylab.figure()
    pylab.axes((0.15, 0.15, 0.8, 0.8))
    pylab.scatter(c1, c2, c=c, linewidth=0, s=1, vmin=vmin, vmax=vmax)
    pylab.colorbar()
    pylab.xlim(min(c1), max(c1))
    pylab.ylim(min(c2), max(c2))
    pylab.xlabel(xl)
    pylab.ylabel(yl)
    pylab.savefig(name)

mhist(imchirp-tmchirp, "plots/hist-mchirp-diff.png")
mhist((imchirp-tmchirp)/imchirp, "plots/hist-mchirp-reldiff.png")
mhist(match, "plots/hist-match.png")
mhist(match, "plots/hist-match-cum.png", cum=1, log=True, bins=10000, xl = "Match", yl="Fraction of injections < Match")
    
pylab.figure(102)
pylab.ylabel('Fraction of Injections')
pylab.xlabel('Fitting factor')
pylab.yscale('log') 
pylab.xlim(0.95, 1.0)
pylab.ylim(1e-4, 1)
hBins = pylab.arange(0.,1.,0.0005,dtype=float)
n, bins,patches=pylab.hist(match,cumulative=1,bins=hBins,normed=True)
pylab.grid()
pylab.savefig("plots/cum_hist.png")
    
mplot(imass1, imass2, match, "plots/m1-m2-match.png")
mplot(tmass1, tmass2, match, "plots/tm1-tm2-match.png")
mplot(q, s1m, match, "plots/q-s1m-match.png")
mplot(q, s2m, match, "plots/q-s2m-match.png")
mplot(q, ispin1z, match, "plots/q-s1z-match.png")  
mplot(q, ispin2z, match, "plots/q-s2z-match.png", "Mass Ratio", "Spin2z") 
mplot(q, ispin2z, match, "plots/q-s2z-match97.png", "Mass Ratio", "Spin2z", vmin=0.97)
mplot(q, ispin2z, match, "plots/q-s2z-match90.png", "Mass Ratio", "Spin2z", vmin=0.90)
mplot(inclination, match, match, "plots/inc-match.png")   

mplot(imass1, imass2, imchirp-tmchirp, "plots/m1-m2-mchirpdiff.png")
mplot(q, ispin1z, imchirp-tmchirp, "plots/q-s1z-mchirpdiff.png", "Mass Ratio", "Spin1z")
mplot(q, ispin2z, imchirp-tmchirp, "plots/q-s2z-mchirpdiff.png", "Mass Ratio", "Spin2z")

mplot(imass1, imass2, (imchirp-tmchirp)/imchirp, "plots/m1-m2-mchirpreldiff.png")
mplot(q, ispin1z, (imchirp-tmchirp)/imchirp, "plots/q-s1z-mchirpreldiff.png")
mplot(q, ispin2z, (imchirp-tmchirp)/imchirp, "plots/q-s2z-mchirpreldiff.png")

""")
os.chmod("scripts/pycbc_banksim_plots", 0o0777)

logging.info("Done")