#!/usr/bin/env python
#
# Copyright (C) 2020 Prayush Kumar
# 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 3 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.

import os
import logging
import shutil
from six.moves import configparser as ConfigParser
import subprocess
import glob
import tempfile
from optparse import OptionParser
from glue.pipeline import CondorDAGNode, CondorDAG
from gwnrtools.workflow.condor import BaseJob, FaithsimNode


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):
    matches = []
    for st in slist:
        if st.startswith(match):
            matches.append(st)
    return matches


parser = OptionParser()
parser.add_option('--config', type=str)
parser.add_option("--skip-creating-injections", action="store_true",
                  help="Skip calling lalapps_inspinj and assume injections already exist",
                  default=False)
(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")
inspinj_prog = confs.get("executables", "inspinj")
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/gwnrtools_faithsim')
os.chmod('scripts/gwnrtools_faithsim', 0o0777)

logging.info("Creating injection file")
inj_str = inspinj_prog + " " + \
    get_ini_opts(confs, "inspinj") + "--output-prefix inj"
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/gwnrtools_faithsim", confs, sec,
                          accounting_group=accounting_group))

rjob = BaseJob("log", "scripts/gwnrtools_faithsim_collect_results", None,
               'collect_results', accounting_group=accounting_group)
rnode = CondorDAGNode(rjob)
pjob = BaseJob("log", "scripts/gwnrtools_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/gwnrtools_faithsim_collect_results", "w")
f.write("""#!/usr/bin/env python
from os.path import isfile
from numpy import *
from 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/gwnrtools_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/gwnrtools_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 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/gwnrtools_faithsim_plots', 0o0777)

logging.info('Done')
