#! /usr/bin/env python
#
#  GOAL
#    Simple pipeline: just CIP + fitting (and puff) iteration jobs.  Here CIP acts as ILE would in normal pipeline. Input from individual events needed.
#
#  
#    - Inputs
#         args_marg_event.txt   # CIP arguments, arguments to compute the marginalized likelihood of each event
#         args_eos_post.txt    # arguments to compute the posterior from marginal likeliyhood
#    - data formats
#         Common format for input (with likelihoods) and output

import argparse
import sys
import os
import shutil
import numpy as np
import RIFT.lalsimutils as lalsimutils
import lalsimulation as lalsim
import lal
import functools
import itertools
import json

# Backend-neutral pipeline namespace (htcondor/glue/slurm) provided by dag_utils_generic
from RIFT.misc.dag_utils_generic import pipeline

import RIFT.misc.dag_utils_generic as dag_utils
from RIFT.misc.dag_utils_generic import mkdir
from RIFT.misc.dag_utils_generic import which

parser = argparse.ArgumentParser()
parser.add_argument("--working-directory",default="./")
parser.add_argument("--input-grid",default=None,help="EOS parameters ")
parser.add_argument("--event-file",default=None,action='append',help="Specify the events here.  Will be copied into 'event-X.net' ")
parser.add_argument("--marg-event-args",default=None,help="filename of args_marg_event.txt file  which holds CIP arguments.  Should NOT conflict with arguments auto-set by this DAG ... in particular, i/o arguments will be modified.   ")
parser.add_argument("--marg-event-exe",default=None,help="filename of CIP or equivalent executable. Will default to `which util_ConstructIntrinsicPosterior_GenericCoordinates` in low-level code")
parser.add_argument("--marg-event-exe-list-file",default=None,help="filename of a file containing a LIST of exe ")
parser.add_argument("--marg-event-args-list-file",default=None,help="filename of a file with one line of arguments per event.  ")
parser.add_argument("--marg-event-nchunk-list-file",default=None,help="filename of a file with one number per event, indicating the chunk size of its marginalization job (how many points to evaluate per job).  Job must accept --using-eos-index and --n-events-to-analyze arguments, like ILE  ")
parser.add_argument("--request-memory-marg",default=16384,type=int,help="Memory request for condor (in Mb) for fitting jobs.")
parser.add_argument("--puff-exe",default=None,help="similar to PUFF, but cannot be the same executable since format different")
parser.add_argument("--puff-args",default=None,help=" argumetns")
parser.add_argument("--puff-max-it",default=np.inf,type=int,help="Maximum iteration number that puffball is applied.  default to infinity ")
parser.add_argument("--tracer-only-marg", action='store_true',
                    help="DEPRECATED no-op. Earlier versions used this flag to skip MARG_* on intermediate iterations; that was wrong -- MARG should run every iteration so the tracer has fresh all.marg_net to consume. Use --puff-input-source marg_net to suppress the redundant MARG_PUFF lane instead.")
parser.add_argument("--tracer-final-marg-iterations", type=int, default=1,
                    help="DEPRECATED no-op. See --tracer-only-marg.")
parser.add_argument("--puff-input-source", default="posterior", choices=("posterior", "marg_net"),
                    help="Which file feeds the puff/tracer node. 'posterior' (default, legacy) reads grid-{k+1}.dat (EOS_POST posterior samples) and writes grid_puff-{k+1}.dat which MARG_PUFF then evaluates. 'marg_net' wires the puff/tracer node to consume all.marg_net (the cumulative ILE-evaluated (lambda, lnL, sigma) table) and write grid-{k+1}.dat directly, suppressing the MARG_PUFF lane entirely; use with --puff-exe util_*TracerUpdate.py.")
parser.add_argument("--eos-post-args",default=None,help="filename of args_ile.txt file  which holds ILE arguments.  Should NOT conflict with arguments auto-set by this DAG ... in particular, i/o arguments will be modified")
parser.add_argument("--eos-post-exe",default=None,help="filename of ILE or equivalent executable. Will default to `which integrate_likelihood_extrinsic` in low-level code")
parser.add_argument("--eos-post-retries",default=0,type=int,help="Number of retry attempts for ILE jobs. (These can fail)")
parser.add_argument("--eos-post-explode-jobs",default=2,type=int,help="Number of MARG jobs to use in posterior generation.")
parser.add_argument("--eos-post-explode-jobs-last",default=2,type=int,help="Number of MARG jobs to use in posterior generation on the final stage (if desired).")
parser.add_argument("--test-args",default=None,help="filename of args_test.txt, which holds test arguments.  Note i/o arguments will be modified, so should NOT specify the samples files or the output file, just the test to be performed and any related arguments")
parser.add_argument("--test-exe",default='convergence_test_samples.py',help="convergence test samples")
parser.add_argument("--general-retries",default=0,type=int,help="Number of retry attempts for internal jobs (convert, CIP, ...). (These can fail, albeit more rarely, usually due to filesystem problems)")
parser.add_argument("--general-request-disk",default="10M",type=str,help="Request disk passed to condor. Must be done for all jobs now")
parser.add_argument("--transfer-file-list",default=None,help="File containing list of *input* filenames to transfer, one name per file. Copied into transfer_files for condor directly.  If provided, also enables attempts to deduce files that need to be transferred for the pipeline to operate, as needed for OSG, etc")
parser.add_argument("--use-singularity",action='store_true',help="Attempts to use a singularity image in SINGULARITY_RIFT_IMAGE")
parser.add_argument("--use-osg",action='store_true',help="Attempts to set up an OSG workflow.  Must submit from an osg allowed submit machine")
parser.add_argument("--use-singularity-local",action='store_true',help="Uses the singularity container even for local universe jobs. Note we are NOT doing file transfer, just starting an image")
parser.add_argument('--n-samples-per-job',default=2000,type=int,help="Number of samples generated each iteration; also, number of marg jobs run each iteration. Should increase with dimension of problem")
parser.add_argument('--neff-threshold',default=800,type=int,help="Number of samples generated each iteration")
parser.add_argument("--condor-local-nonworker",action='store_true',help="Uses local universe for non-worker condor jobs. Important to run in non-NFS location, as other jobs don't have file transfer set up.  Required for public OSG.")
parser.add_argument("--condor-local-nonworker-igwn-prefix",action='store_true', help="Adds some prefix text to start up cvmfs igwn environment, so local jobs have access to standard RIFT operators. Required for public OSG.")
parser.add_argument("--condor-nogrid-nonworker",action='store_true',help="Uses local flocking for non-worker condor jobs. Important to run in non-NFS location, as other jobs don't have file transfer set up.")
parser.add_argument('--n-iterations',default=3,type=int,help="Number of iterations to perform")
parser.add_argument("--start-iteration",default=0,type=int,help="starting iteration. If >0, does not copy over --input-grid. DOES rewrite sub files.  This allows you to change the arguments provided (e.g., use more iterations or settings at late times). Note this overwrites the .sub files ")
parser.add_argument("--use-full-submit-paths",action='store_true',help="DAG created has full paths to submit files generated. Note this is implemented on a per-file/as-needed basis, mainly to facilitate using this dag as an external subdag")

opts=  parser.parse_args()
print(opts)
if not(opts.input_grid):
    raise Exception(" --input-grid FNAME required ")

main_worker_universe='vanilla'
local_worker_universe="vanilla"
desired_sites_name = "none"
no_worker_grid=False
if opts.condor_local_nonworker:
    local_worker_universe="local"
if opts.condor_nogrid_nonworker:
    no_worker_grid=True

extra_text = ''
if opts.condor_local_nonworker_igwn_prefix:
    extra_text += "source /cvmfs/software.igwn.org/conda/etc/profile.d/conda.sh \n"
    extra_text +=  "conda activate igwn \n"


working_dir_inside_local = working_dir_inside = opts.working_directory
out_dir_inside_marg = opts.working_directory
out_dir_inside_marg+= "/iteration_$(macroiteration)_marg/event_$(macroid)/"
out_dir_inside_post = opts.working_directory
if opts.use_singularity or opts.use_osg:
    working_dir_inside = "./" # all files on the remote machine are in the current directory
    out_dir_inside_marg = "./"
    out_dir_inside_post = "./"

singularity_image = None
if opts.use_singularity:
    print(" === USING SINGULARITY === ")
    singularity_image = os.environ["SINGULARITY_RIFT_IMAGE"]  # must be present to use singularity
    # SINGULARITY IMAGES ARE ON CVMFS, SO WE CAN AVOID THE SINGULARITY EXEC CALL
    # hardcoding a fiducial copy of lalapps_path2cache; beware about the executable name change
    os.environ['LALAPPS_PATH2CACHE'] = "/cvmfs/oasis.opensciencegrid.org/ligo/sw/conda/envs/igwn-py39/bin/lal_path2cache" #"singularity exec {singularity_image} lalapps_path2cache".format(singularity_image=singularity_image)
    print(singularity_image)


###
### Event files : copy, give standard name (naming convention assumes CIP for now)
###

indx=0
for event in opts.event_file:
    shutil.copyfile(event,opts.working_directory+"/event-{}.net".format(indx))  # put in working directory !
    indx+=1
n_events = len(opts.event_file)



###
### Process args
###

# Load args.txt. Remove first item.  Store
marg_event_args_list=None
marg_event_exe_list=None
marg_event_nchunk_list=None
if opts.marg_event_args:
    with open(opts.marg_event_args) as f:
        marg_event_args_list = f.readlines()
    marg_event_args = ' '.join( [x.replace('\\','') for x in marg_event_args_list] )
    marg_event_args = ' '.join(marg_event_args.split(' ')[1:])
    # Some argument protection for later
    marg_event_args = marg_event_args.replace('[', ' \'[')
    marg_event_args = marg_event_args.replace(']', ']\'')
    marg_event_args=marg_event_args.rstrip()
    print("MARG", marg_event_args)
    marg_event_args_list =None
elif opts.marg_event_args_list_file:
    with open(opts.marg_event_args_list_file) as f:
        marg_event_args_list = f.readlines()
    with open(opts.marg_event_exe_list_file) as f:
        marg_event_exe_list = f.readlines()
    marg_event_nchunk_list = np.ones(len(marg_event_exe_list)) # default: one point per evaluation
    if opts.marg_event_nchunk_list_file and os.path.isfile(opts.marg_event_nchunk_list_file):        
        marg_event_nchunk_list = np.loadtxt(opts.marg_event_nchunk_list_file,dtype=int) # load in 1d array
        if len(marg_event_exe_list) ==1 or not(isinstance(marg_event_nchunk_list, np.ndarray)):
            marg_event_nchunk_list = [marg_event_nchunk_list] # we have a single element
        assert len(marg_event_nchunk_list) == len(marg_event_exe_list)
        
    assert len(marg_event_exe_list) == len(marg_event_args_list)  # fail if not true: not matching MARG args to exe
    assert len(marg_event_exe_list) == len(opts.event_file)  # fail if not true: not matching MARG to event list
    for indx in np.arange(len(marg_event_args_list)):
        exe_here = marg_event_exe_list[indx]
        exe_here.rstrip()
        marg_event_exe_list[indx] = exe_here
        line = marg_event_args_list[indx]
        line = line.rstrip()  # remove CR, whitespace
        marg_event_args_list[indx]=line
        print("MARG {} {} ".format(indx, exe_here), line)

marg_exe_master = None
if opts.marg_event_exe:
    marg_exe_master = "{}".format(marg_exe) # force clpy, so NOT BY REFERENCE

if opts.marg_event_args:
    with open('command-single.sh','w') as f:
        f.write('util_ConstructIntrinsicPosterior_GenericCoordinates.py {} '.format(marg_event_args))


puff_args=None
puff_cadence = None
puff_max_it = opts.puff_max_it
# Gate on --puff-exe: if the user did not provide a puff executable, we do not
# write puff/MARG_PUFF nodes at all. Holding --puff-args without --puff-exe is
# an obvious mistake; refuse to proceed silently.
if opts.puff_args and not opts.puff_exe:
    sys.exit(" --puff-args given without --puff-exe: refusing to write puff nodes without a puff executable. Pass --puff-exe or drop --puff-args.")
if opts.puff_args and opts.puff_exe:
    puff_cadence = 1
    # Load args.txt. Remove first item.  Store
    with open(opts.puff_args) as f:
        puff_args_list = f.readlines()
    puff_args = ' '.join( [x.replace('\\','') for x in puff_args_list] )
    puff_args = ' '.join(puff_args.split(' ')[1:])
    # Some argument protection for later
    puff_args = puff_args.replace('[', ' \'[')
    puff_args = puff_args.replace(']', ']\'')
    puff_args=puff_args.rstrip()
    print("PUFF", puff_args)
    print("PUFF CADENCE", puff_cadence)

# Tracer-aware puff input plumbing. In 'marg_net' mode the puff/tracer node
# consumes all.marg_net (the cumulative ILE-evaluated table emitted by UNIFY)
# and writes grid-{k+1}.dat directly. EOS_POST output is redirected to
# posterior-{k+1}.dat for diagnostics, and the MARG_PUFF lane is suppressed --
# the tracer's output IS the next iteration's grid.
puff_uses_marg_net = (opts.puff_input_source == "marg_net") and bool(puff_args)
if puff_uses_marg_net:
    print("PUFF INPUT SOURCE: marg_net (tracer-aware grid update)")


test_args=None
test_node_list = []
if opts.test_args:
    # Load args.txt. Remove first item.  Store
    with open(opts.test_args) as f:
        test_args_list = f.readlines()
    test_args = ' '.join( [x.replace('\\','') for x in test_args_list] )
    test_args = ' '.join(test_args.split(' ')[1:])
# Some argument protection for later
    test_args = test_args.replace('[', ' \'[')
    test_args = test_args.replace(']', ']\'')
    test_args=test_args.rstrip()
    print("CONVERGE", test_args)

    
# Load args.txt. Remove first item.  Store
if True: # not (opts.eos_post_args is None):
    with open(opts.eos_post_args) as f:
        eos_post_args_list = f.readlines()
    eos_post_args = ' '.join( [x.replace('\\','') for x in eos_post_args_list] )
    eos_post_args = ' '.join(eos_post_args.split(' ')[1:])
    # Some argument protection for later.  Make sure spaces in place to avoid breaking compound arguments like g0:[a,b] for ranges.
    # intended to make sure any spaces in arguments like [a, b] are encoded as '[a, b]' 
#    eos_post_args = eos_post_args.replace(' [', '  \'[')
#    eos_post_args = eos_post_args.replace('] ', ']\' ')
    eos_post_args=eos_post_args.rstrip()
    eos_post_args += ' --no-plots '  
    print("EOS_POST", eos_post_args)

eospost_exe = which('util_ConstructEOSPosterior.py')
if opts.eos_post_exe:
    eospost_exe = "{}".format(opts.eos_post_exe) # force clpy, so NOT BY REFERENCE
if opts.use_osg:
    # Correctly use EOS_POST in container location
    path_split = eospost_exe.split("/")
    singularity_base_exe_path = os.environ['SINGULARITY_BASE_EXE_DIR']
    eospost_exe =singularity_base_exe_path + path_split[-1]


# Copy seed grid into place as grid-0.dat
it_start = opts.start_iteration
n_initial = opts.n_samples_per_job
if (it_start == 0):
    shutil.copyfile(opts.input_grid,"grid-0.dat")  # put in working directory !
    n_initial = len(np.loadtxt("grid-0.dat"))

transfer_file_names = []
if not (opts.transfer_file_list is None):
    transfer_file_names=[]
    with open(opts.transfer_file_list) as f:
        for  line in f.readlines():
            transfer_file_names.append(line.rstrip())
    print(" Input files to transfer to job working directory (note!)", transfer_file_names)
transfer_file_names_post = list(transfer_file_names) + ['../all.marg_net'] # make a copy



###
### DAG generation
###


dag = pipeline.CondorDAG(log=os.getcwd())


# Make directories for all iterations
for it_indx in np.arange(it_start,opts.n_iterations+1):
    ile_dir = opts.working_directory+"/iteration_"+str(it_indx)+"_marg"
    cip_dir = opts.working_directory+"/iteration_"+str(it_indx)+"_post"
    consolidate_dir = opts.working_directory+"/iteration_"+str(it_indx)+"_con"
    mkdir(ile_dir); mkdir(ile_dir+"/logs")
    for event_indx, _event in enumerate(opts.event_file):
        mkdir(ile_dir + "/event_{}".format(event_indx))
    mkdir(cip_dir);  mkdir(cip_dir+"/logs")
    mkdir(consolidate_dir); mkdir(consolidate_dir+"/logs")

    if opts.test_args:
        test_dir = opts.working_directory+"/iteration_"+str(it_indx)+"_test"
        mkdir(test_dir); mkdir(test_dir+'/logs')


transfer_file_names_puff = list(transfer_file_names)  # allocate a copy. Actually fine if the same thing, just a little redundant.
if opts.use_singularity or opts.use_osg:
    transfer_file_names_post.append("../grid-$(macroiteration).dat")
    transfer_file_names.append(opts.working_directory+"/grid-$(macroiteration).dat")
    transfer_file_names_puff.append(opts.working_directory+"/grid_puff-$(macroiteration).dat")


# ++++
# Create workflow tasks
# ++++
cip_job_list =[]
cip_puff_job_list =[]
if opts.marg_event_args:
##   MARG job (CIP is default)
#   - issue is that transfer files depend on event
#   - all output goes into the same iter*_marg directory
#   - join script will merge all together
    cip_job, cip_job_name = dag_utils.write_CIP_sub(tag='MARG',log_dir=None,arg_str=marg_event_args,using_eos='file:'+working_dir_inside+"/grid-$(macroiteration).dat",using_eos_index='$(macroevent)',request_memory=opts.request_memory_marg,input_net=working_dir_inside+'/event-$(macroid).net',output='MARG-$(macroid)-$(macroevent)',out_dir=out_dir_inside_marg,exe=marg_exe_master,universe=main_worker_universe,no_grid=not(opts.use_osg),use_osg=opts.use_osg,use_singularity=opts.use_osg and opts.use_singularity,singularity_image=singularity_image,use_simple_osg_requirements=opts.use_osg,transfer_files=['../../event-$(macroid).net']+transfer_file_names)
    # Modify: set 'initialdir'
    cip_job.add_condor_cmd("initialdir",opts.working_directory+"/iteration_$(macroiteration)_marg/event_$(macroid)")
    # Modify output argument: change logs and working directory to be subdirectory for the run
    cip_job.set_log_file(opts.working_directory+"/iteration_$(macroiteration)_marg/logs/marg-$(macroevent)-$(macroid)-$(cluster)-$(process).log")
    cip_job.set_stderr_file(opts.working_directory+"/iteration_$(macroiteration)_marg/logs/marg-$(macroevent)-$(macroid)-$(cluster)-$(process).err")
    cip_job.set_stdout_file(opts.working_directory+"/iteration_$(macroiteration)_marg/logs/marg-$(macroevent)-$(macroid)-$(cluster)-$(process).out")
    cip_job.add_condor_cmd('request_disk',opts.general_request_disk)
    if opts.use_full_submit_paths:
        fname = opts.working_directory+"/"+cip_job.get_sub_file()
        cip_job.set_sub_file(fname)
    cip_job.write_sub_file()

    # PUFF variant -- only emit when MARG_PUFF will actually be added to the
    # DAG (legacy posterior-mode wiring). In tracer mode (puff_uses_marg_net)
    # the tracer writes grid-{k+1}.dat directly and MARG_PUFF is redundant.
    if puff_args and not puff_uses_marg_net:
        cipPuff_job, cipPuff_job_name = dag_utils.write_CIP_sub(tag='MARG_PUFF',log_dir=None,arg_str=marg_event_args,using_eos='file:'+working_dir_inside+"/grid_puff-$(macroiteration).dat",using_eos_index='$(macroevent)',request_memory=opts.request_memory_marg,input_net=working_dir_inside+'/event-$(macroid).net',output='MARG_puff-$(macroid)-$(macroevent)',out_dir=out_dir_inside_marg,exe=marg_exe_master,universe=main_worker_universe,no_grid=not(opts.use_osg),use_osg=opts.use_osg,use_singularity=opts.use_osg and opts.use_singularity,singularity_image=singularity_image,use_simple_osg_requirements=opts.use_osg,transfer_files=['../../event-$(macroid).net']+transfer_file_names_puff)
        # Modify: set 'initialdir'
        cipPuff_job.add_condor_cmd("initialdir",opts.working_directory+"/iteration_$(macroiteration)_marg/event_$(macroid)")
        # Modify output argument: change logs and working directory to be subdirectory for the run
        cipPuff_job.set_log_file(opts.working_directory+"/iteration_$(macroiteration)_marg/logs/margPuff-$(macroevent)-$(cluster)-$(process).log")
        cipPuff_job.set_stderr_file(opts.working_directory+"/iteration_$(macroiteration)_marg/logs/margPuff-$(macroevent)-$(cluster)-$(process).err")
        cipPuff_job.set_stdout_file(opts.working_directory+"/iteration_$(macroiteration)_marg/logs/margPuff-$(macroevent)-$(cluster)-$(process).out")
        cipPuff_job.add_condor_cmd('request_disk',opts.general_request_disk)
        if opts.use_full_submit_paths:
            fname = opts.working_directory+"/"+cipPuff_job.get_sub_file()
            cipPuff_job.set_sub_file(fname)
        cipPuff_job.write_sub_file()

else:
    for indx in np.arange(len(marg_event_args_list)):
        n_chunk_here = int(marg_event_nchunk_list[indx])
        cip_job, cip_job_name = dag_utils.write_CIP_sub(tag='MARG_{}'.format(indx),log_dir=None,exe=marg_event_exe_list[indx],arg_str=marg_event_args_list[indx],using_eos='file:'+working_dir_inside+"/grid-$(macroiteration).dat",using_eos_index='$(macroevent)',n_events_to_analyze=n_chunk_here,request_memory=opts.request_memory_marg,input_net=working_dir_inside+'/event-$(macroid).net',output='MARG-$(macroid)-$(macroevent)',out_dir=out_dir_inside_marg,universe=main_worker_universe,no_grid=not(opts.use_osg),use_osg=opts.use_osg,use_singularity=opts.use_osg and opts.use_singularity,singularity_image=singularity_image,use_simple_osg_requirements=opts.use_osg,transfer_files=['../../event-$(macroid).net']+transfer_file_names)
        # Modify: set 'initialdir'
        cip_job.add_condor_cmd("initialdir",opts.working_directory+"/iteration_$(macroiteration)_marg/event_$(macroid)")
        # Modify output argument: change logs and working directory to be subdirectory for the run
        cip_job.set_log_file(opts.working_directory+"/iteration_$(macroiteration)_marg/logs/marg-$(macroevent)-$(macroid)-$(cluster)-$(process).log")
        cip_job.set_stderr_file(opts.working_directory+"/iteration_$(macroiteration)_marg/logs/marg-$(macroevent)-$(macroid)-$(cluster)-$(process).err")
        cip_job.set_stdout_file(opts.working_directory+"/iteration_$(macroiteration)_marg/logs/marg-$(macroevent)-$(macroid)-$(cluster)-$(process).out")
        cip_job.add_condor_cmd('request_disk',opts.general_request_disk)
        if opts.use_full_submit_paths:
            fname = opts.working_directory+"/"+cip_job.get_sub_file()
            cip_job.set_sub_file(fname)
        cip_job.write_sub_file()
        cip_job_list.append(cip_job)

        # PUFF variant -- only emit when MARG_PUFF will actually be added to
        # the DAG (legacy posterior-mode wiring). In tracer mode the tracer
        # writes grid-{k+1}.dat directly and MARG_PUFF is redundant.
        if puff_args and not puff_uses_marg_net:
            cipPuff_job, cipPuff_job_name = dag_utils.write_CIP_sub(tag='MARG_PUFF_{}'.format(indx),log_dir=None,exe=marg_event_exe_list[indx],arg_str=marg_event_args_list[indx],using_eos='file:'+working_dir_inside+"/grid_puff-$(macroiteration).dat",using_eos_index='$(macroevent)',n_events_to_analyze=n_chunk_here,request_memory=opts.request_memory_marg,input_net=working_dir_inside+'/event-$(macroid).net',output='MARG_puff-$(macroid)-$(macroevent)',out_dir=out_dir_inside_marg,universe=main_worker_universe,no_grid=not(opts.use_osg),use_osg=opts.use_osg,use_singularity=opts.use_osg and opts.use_singularity,singularity_image=singularity_image,use_simple_osg_requirements=opts.use_osg,transfer_files=['../../event-$(macroid).net','../../grid_puff-$(macroiteration).dat']+transfer_file_names_puff)
            # Modify: set 'initialdir'
            cipPuff_job.add_condor_cmd("initialdir",opts.working_directory+"/iteration_$(macroiteration)_marg/event_$(macroid)")
            # Modify output argument: change logs and working directory to be subdirectory for the run
            cipPuff_job.set_log_file(opts.working_directory+"/iteration_$(macroiteration)_marg/logs/margPuff-$(macroevent)-$(macroid)-$(cluster)-$(process).log")
            cipPuff_job.set_stderr_file(opts.working_directory+"/iteration_$(macroiteration)_marg/logs/margPuff-$(macroevent)-$(macroid)-$(cluster)-$(process).err")
            cipPuff_job.set_stdout_file(opts.working_directory+"/iteration_$(macroiteration)_marg/logs/margPuff-$(macroevent)-$(macroid)-$(cluster)-$(process).out")
            cipPuff_job.add_condor_cmd('request_disk',opts.general_request_disk)
            if opts.use_full_submit_paths:
                fname = opts.working_directory+"/"+cipPuff_job.get_sub_file()
                cipPuff_job.set_sub_file(fname)
            cipPuff_job.write_sub_file()
            cip_puff_job_list.append(cipPuff_job)



##   CON job 1 : consolidate single events
#  - joins together all data files from a single run
#  - matches duplicates, averaging likelihoods
#  - if multiple events present, works fine,but will AVERAGE and not ADD, so not as significant as needed. Change?
#  - skip for now, just have one event at firsty ==>? SONOT QUTE CORRECT because likelihoods are averaged, which reduces net significance.  FIXME: two-stage process

# If we make the job inside the container (as usual), the paths will be correct, so no need to override this
with open("con_marg.sh",'w') as f:
    relevant_path = dag_utils.which('util_HyperCombine.py')
    if opts.condor_local_nonworker_igwn_prefix:
        relevant_path = 'util_HyperCombine.py'  # we are in igwn environment
    f.write("""#! /bin/bash
 {}
{} {}/iteration_$1_marg/event_$2/MARG*[0-9]+annotation.dat  > {}/iteration_$1_marg/consolidated_$1_$2.net_marg
""".format(extra_text,relevant_path,opts.working_directory,opts.working_directory) )
    os.system("chmod a+x con_marg.sh")

con_marg_job, con_marg_job_name = dag_utils.write_convert_sub(exe=opts.working_directory+"/con_marg.sh",tag='CON',log_dir=None,arg_str='',file_input="$(macroiteration)  $(macroevent) ", out_dir=opts.working_directory,universe=local_worker_universe,no_grid=no_worker_grid)
con_marg_job.add_condor_cmd("initialdir",opts.working_directory)
con_marg_job.set_log_file(opts.working_directory+"/iteration_$(macroiteration)_con/logs/con-$(macroevent).log")
con_marg_job.set_stdout_file(opts.working_directory+"/iteration_$(macroiteration)_con/logs/con-$(macroevent).out")
con_marg_job.set_stderr_file(opts.working_directory+"/iteration_$(macroiteration)_con/logs/con-$(macroevent).err")
con_marg_job.add_condor_cmd('request_disk',opts.general_request_disk)
if opts.use_full_submit_paths:
    fname = opts.working_directory+"/"+con_marg_job.get_sub_file()
    con_marg_job.set_sub_file(fname)
if opts.use_osg:
    con_marg_job.add_condor_cmd("+DESIRED_SITES",'"{}"'.format(desired_sites_name))
    con_marg_job.add_condor_cmd("+flock_local",'true')
if opts.use_singularity_local:
    con_marg.add_condor_cmd("+SingularityImage", '"' + singularity_image + '"')
con_marg_job.write_sub_file()


##   CON job 2 : join multiple events together into a single overall result
#  - joins together all data files from a single run
#  - matches duplicates, averaging likelihoods
#  - if multiple events present, works fine,but will AVERAGE and not ADD, so not as significant as needed. Change?
#  - skip for now, just have one event at firsty ==>? SONOT QUTE CORRECT because likelihoods are averaged, which reduces net significance.  FIXME: two-stage process
with open("con_prod.sh",'w') as f:
    relevant_path = dag_utils.which('util_HyperCombine.py')
    if opts.condor_local_nonworker_igwn_prefix:
        relevant_path = 'util_HyperCombine.py'  # we are in igwn environment
    f.write("""#! /bin/bash
 {}
{} --combination product {}/iteration_$1_marg/consolidated_*_*.net_marg  > {}/consolidated_$1.net_marg
""".format(extra_text,relevant_path,opts.working_directory,opts.working_directory) )
    os.system("chmod a+x con_prod.sh")

con_prod_job, con_prod_job_name = dag_utils.write_convert_sub(exe=opts.working_directory+"/con_prod.sh",tag='CON_PROD',log_dir=None,arg_str='',file_input="$(macroiteration)  ", out_dir=opts.working_directory,universe=local_worker_universe,no_grid=no_worker_grid)
con_prod_job.add_condor_cmd("initialdir",opts.working_directory)
con_prod_job.set_log_file(opts.working_directory+"/iteration_$(macroiteration)_con/logs/con_prod-$(macroevent)-$(process).log")
con_prod_job.set_stdout_file(opts.working_directory+"/iteration_$(macroiteration)_con/logs/con_prod-$(macroevent)-$(process).out")
con_prod_job.set_stderr_file(opts.working_directory+"/iteration_$(macroiteration)_con/logs/con_prod-$(macroevent)-$(process).err")
con_prod_job.add_condor_cmd('request_disk',opts.general_request_disk)
if opts.use_full_submit_paths:
    fname = opts.working_directory+"/"+con_prod_job.get_sub_file()
    con_prod_job.set_sub_file(fname)
if opts.use_osg:
    con_prod_job.add_condor_cmd("+DESIRED_SITES",'"{}"'.format(desired_sites_name))
    con_prod_job.add_condor_cmd("+flock_local",'true')
if opts.use_singularity_local:
    con_marg.add_condor_cmd("+SingularityImage", '"' + singularity_image + '"')
con_prod_job.write_sub_file()
        

if opts.test_args:
    # In tracer mode (puff_uses_marg_net) the file grid-{k}.dat is a placement
    # grid (lnL columns zero), and the true iteration posterior lives in
    # posterior-{k}.dat. So the convergence test must compare posterior
    # samples, not placement grids.
    if puff_uses_marg_net:
        _test_curr = opts.working_directory + '/posterior-$(macroiteration).dat'
        _test_prev = opts.working_directory + '/posterior-$(macroiterationlast).dat'
    else:
        _test_curr = opts.working_directory + '/grid-$(macroiteration).dat'
        _test_prev = opts.working_directory + '/grid-$(macroiterationlast).dat'
    # Resolve --test-exe to a full path matching how eospost_exe / the
    # patched puff_exe are handled: bare names get resolved via `which`;
    # absolute paths pass through. Condor on some test machines does not
    # search PATH for local-universe jobs.
    #
    # The script as shipped is named ``convergence_test_samples.py``
    # (setup.py installs everything in bin/ verbatim), but upstream
    # layers historically default to the bare ``convergence_test_samples``
    # name. If the bare lookup fails, try the .py-suffixed form too so
    # those defaults keep working without forcing every caller to update.
    test_exe_resolved = opts.test_exe
    if test_exe_resolved and not os.path.isabs(test_exe_resolved):
        _r = dag_utils.which(test_exe_resolved)
        if not _r and not test_exe_resolved.endswith(".py"):
            _r = dag_utils.which(test_exe_resolved + ".py")
        if _r:
            test_exe_resolved = os.path.abspath(_r)
        else:
            print(
                " WARNING: --test-exe {!r} did not resolve via PATH; test.sub "
                "will use the bare name (condor must find it itself).".format(
                    opts.test_exe
                )
            )
    test_job, test_job_name = dag_utils.write_test_sub(tag='test',log_dir=None,arg_str=test_args,samples_files=[_test_curr, _test_prev] ,out_dir=opts.working_directory,exe=test_exe_resolved,universe=local_worker_universe,no_grid=no_worker_grid)
    # Modify: set 'initialdir'
    test_job.add_condor_cmd("initialdir",opts.working_directory+"/iteration_$(macroiteration)_test")
    # Modify output argument: change logs and working directory to be subdirectory for the run
    test_job.set_log_file(opts.working_directory+"/iteration_$(macroiteration)_test/logs/test-$(cluster)-$(process).log")
    test_job.set_stderr_file(opts.working_directory+"/iteration_$(macroiteration)_test/logs/test-$(cluster)-$(process).err")
    test_job.set_stdout_file(opts.working_directory+"/iteration_$(macroiteration)_test/logs/test-$(cluster)-$(process).out")
    test_job.add_condor_cmd('request_disk',opts.general_request_disk)
    if opts.use_full_submit_paths:
        fname = opts.working_directory+"/"+test_job.get_sub_file()
        test_job.set_sub_file(fname)
#    if opts.condor_containerize_nonworker and singularity_image:
#        test_job.add_condor_cmd("+SingularityImage", '"' + singularity_image + '"')
    test_job.write_sub_file()


##   unify job
#  - joins together all files into one
with open("unify.sh",'w') as f:
    f.write("""#! /bin/bash
ls -1 *.net_marg | head -n 1 | xargs head -n 1  > all.marg_net
grep -h ^[[:space:]][0-9] *.net_marg | sort -nr >> all.marg_net
grep -h ^[[:space:]][-] *.net_marg | grep -v "inf" | sort -nr >> all.marg_net
""")
    os.system("chmod a+x unify.sh")
unify_marg_job, unify_marg_job_name = dag_utils.write_convert_sub(exe=opts.working_directory+"/unify.sh",tag='UNIFY',log_dir=None,arg_str='',file_input='',out_dir=opts.working_directory,universe=local_worker_universe,no_grid=no_worker_grid)
unify_marg_job.add_condor_cmd("initialdir",opts.working_directory)
unify_marg_job.set_log_file(opts.working_directory+"/iteration_$(macroiteration)_con/logs/unify-$(process).log")
unify_marg_job.set_stdout_file(opts.working_directory+"/iteration_$(macroiteration)_con/logs/unify-$(process).out")
unify_marg_job.set_stderr_file(opts.working_directory+"/iteration_$(macroiteration)_con/logs/unify-$(process).err")
unify_marg_job.add_condor_cmd('request_disk',opts.general_request_disk)
if opts.use_full_submit_paths:
    fname = opts.working_directory+"/"+unify_marg_job.get_sub_file()
    unify_marg_job.set_sub_file(fname)
if opts.use_osg:
    unify_marg_job.add_condor_cmd("+DESIRED_SITES",'"{}"'.format(desired_sites_name))
    unify_marg_job.add_condor_cmd("+flock_local",'true')
if opts.use_singularity_local:
    con_marg.add_condor_cmd("+SingularityImage", '"' + singularity_image + '"')
unify_marg_job.write_sub_file()


##   EOS_POST job  (workers)
#  - main hyperparameter integration executable
eospost_outfile = 'iteration_$(macroiteration)_post/output-$(macroiterationnext)-$(cluster)'
if opts.use_osg:
    eospost_outfile = 'output-$(macroiterationnext)-$(cluster)'
eospost_marg_job, eospost_marg_job_name = dag_utils.write_hyperpost_sub(exe=eospost_exe,tag='EOS_POST_worker',arg_str=eos_post_args,out_dir=out_dir_inside_post,output=eospost_outfile,input_net=working_dir_inside + "/all.marg_net",transfer_files=transfer_file_names_post)
eospost_marg_job.add_condor_cmd("initialdir",opts.working_directory+"/iteration_$(macroiteration)_post")
if opts.use_singularity:
    eospost_marg_job.add_condor_cmd("+SingularityImage", '"' + singularity_image + '"')
    eospost_marg_job.add_condor_cmd("transfer_executable", 'False') # don't copy it
eospost_marg_job.set_log_file(opts.working_directory+"/iteration_$(macroiteration)_post/logs/post-$(cluster)-$(process).log")
eospost_marg_job.set_stderr_file(opts.working_directory+"/iteration_$(macroiteration)_post/logs/post-$(cluster)-$(process).err")
eospost_marg_job.set_stdout_file(opts.working_directory+"/iteration_$(macroiteration)_post/logs/post-$(cluster)-$(process).out")
eospost_marg_job.add_condor_cmd('request_disk',opts.general_request_disk)
if opts.use_full_submit_paths:
    fname = opts.working_directory+"/"+eospost_marg_job.get_sub_file()
    eospost_marg_job.set_sub_file(fname)
eospost_marg_job.write_sub_file()


# worker join job
with open("join_post.sh",'w') as f:
    f.write("""#! /bin/bash
head -n 1 grid-0.dat > $2
cat $1/output*.dat  | sort -r | uniq | shuf >> $2
""")
    os.system("chmod a+x join_post.sh")
# In tracer mode the tracer node owns grid-{k+1}.dat; EOS_POST output goes to a
# parallel posterior-{k+1}.dat file (diagnostic only). In legacy mode EOS_POST
# still produces grid-{k+1}.dat directly.
_join_post_output_path = opts.working_directory + (
    "/posterior-$(macroiterationnext).dat" if puff_uses_marg_net
    else "/grid-$(macroiterationnext).dat")
join_post_job, join_post_job_name = dag_utils.write_convert_sub(exe=opts.working_directory+"/join_post.sh",tag='JOIN_POST',log_dir=None,arg_str='',file_input=opts.working_directory+"/iteration_$(macroiteration)_post/    "+_join_post_output_path, out_dir=opts.working_directory,universe=local_worker_universe,no_grid=no_worker_grid)
join_post_job.add_condor_cmd("initialdir",opts.working_directory)
join_post_job.set_log_file(opts.working_directory+"/iteration_$(macroiteration)_con/logs/unify-$(macroevent).log")
join_post_job.set_stdout_file(opts.working_directory+"/iteration_$(macroiteration)_con/logs/unify-$(macroevent).out")
join_post_job.set_stderr_file(opts.working_directory+"/iteration_$(macroiteration)_con/logs/unify-$(macroevent).err")
join_post_job.add_condor_cmd('request_disk',opts.general_request_disk)
if opts.use_full_submit_paths:
    fname = opts.working_directory+"/"+join_post_job.get_sub_file()
    join_post_job.set_sub_file(fname)
join_post_job.write_sub_file()
if opts.use_osg:
    join_post_job.add_condor_cmd("+DESIRED_SITES",'"{}"'.format(desired_sites_name))
    join_post_job.add_condor_cmd("+flock_local",'true')


# PUFF job
#   ... pending
if puff_args and puff_cadence:
    # Resolve --puff-exe to a full path matching how eospost_exe is handled
    # above: bare names (e.g. "util_HyperparameterTracerUpdate.py") get
    # resolved via `which`; if the user already passed an absolute path,
    # leave it alone. Condor on some test machines does not search PATH
    # when launching local-universe jobs and fails to start PUFF otherwise.
    exe_here = opts.puff_exe
    if exe_here and not os.path.isabs(exe_here):
        resolved = dag_utils.which(exe_here)
        if resolved:
            exe_here = os.path.abspath(resolved)
    if opts.use_osg:
        relevant_path = dag_utils.which(opts.puff_exe)
        if opts.condor_local_nonworker_igwn_prefix:
            relevant_path = relevant_path.split("/")[-1]
            exe_here = opts.working_directory + "/local_puff.sh"
        with open("local_puff.sh", 'w') as f:
            f.write("""#! /bin/bash
{}
{} $@
""".format(extra_text,relevant_path))
            os.system("chmod a+x local_puff.sh")
    # Tracer mode: consume all.marg_net (ILE-evaluated table) and write
    # grid-{k+1}.dat directly. Legacy mode: consume grid-{k+1}.dat (posterior
    # samples from JOIN_POST) and write grid_puff-{k+1}.dat for MARG_PUFF.
    if puff_uses_marg_net:
        _puff_input_net = opts.working_directory + '/all.marg_net'
        _puff_output    = opts.working_directory + '/grid-$(macroiterationnext).dat'
    else:
        _puff_input_net = opts.working_directory + '/grid-$(macroiterationnext).dat'
        _puff_output    = opts.working_directory + '/grid_puff-$(macroiterationnext).dat'
    puff_job, puff_job_name = dag_utils.write_puff_sub(tag='PUFF',log_dir=None,arg_str=puff_args,request_memory=opts.request_memory_marg,input_net=_puff_input_net,output=_puff_output,out_dir=opts.working_directory,exe=exe_here,universe=local_worker_universe,no_grid=no_worker_grid)
    # Modify: set 'initialdir' to CIP WORKING DIR 
    puff_job.add_condor_cmd("initialdir",opts.working_directory+"/iteration_$(macroiteration)_post")
    # Modify output argument: change logs and working directory to be subdirectory for the run
    puff_job.set_log_file(opts.working_directory+"/iteration_$(macroiteration)_post/logs/puff-$(cluster)-$(process).log")
    puff_job.set_stderr_file(opts.working_directory+"/iteration_$(macroiteration)_post/logs/puff-$(cluster)-$(process).err")
    puff_job.set_stdout_file(opts.working_directory+"/iteration_$(macroiteration)_post/logs/puff-$(cluster)-$(process).out")
    puff_job.add_condor_cmd('request_disk',opts.general_request_disk)
    if opts.use_full_submit_paths:
        fname = opts.working_directory+"/"+puff_job.get_sub_file()
        puff_job.set_sub_file(fname)
    puff_job.write_sub_file()
    if opts.use_osg:
        puff_job.add_condor_cmd("+DESIRED_SITES",'"{}"'.format(desired_sites_name))
        puff_job.add_condor_cmd("+flock_local",'true')


# ++++
# Create workflow 
# ++++

# Create workflow
#   - Create grid node as needed
#   - Loop over iterations
#      - if iteration0, use seed grid (should already be copied in place)
#      - if not iteration 0, grid should be in place (from previous stage)
#      - Loop over events, make ILE node per event
#      - create consolidate job, make it depend on all events in that iteration
#      - create fit job, make it depend on consolidate job
#    

parent_fit_node = None
last_node=None


unify_node_list = []
for it in np.arange(it_start,opts.n_iterations):
    print(it, opts.n_iterations)
    consolidate_now = None
    fit_node_now = None
    marg_nodes_now =[]
    # Create consolidate job 2
    con_prod_node = pipeline.CondorDAGNode(con_prod_job)
    con_prod_node.add_macro("macroiteration",it)
    con_prod_node.set_retry(opts.general_retries)
    # Create unify job
    unify_node = pipeline.CondorDAGNode(unify_marg_job)
    unify_node.add_macro("macroiteration",it)
    unify_node.set_retry(opts.general_retries)
    if not(it ==0): # and not(opts.first_iteration_jumpstart):  # don't require first composite to be nonempty if we are running a jumpstart!
        unify_node_list.append(unify_node)
    
    # Create one node per job.  (Will change on per event basis)
    n_group = 1
    n_jobs_this_time = opts.n_samples_per_job # hardcode for now?
    if it ==it_start:
        n_jobs_this_time = n_initial
    indx_max = int((1.0*n_jobs_this_time)/n_group)
    if indx_max*n_jobs_this_time < n_group:
        indx_max+=1

    if True: #not(it==it_start): # and opts.first_iteration_jumpstart):   # if on first iteration ,don't do this for jumpstart
        for event in np.arange(n_events):
            n_chunk_here = marg_event_nchunk_list[event]
            indx_max_event = indx_max/n_chunk_here
            # Create consolidate job 1 for that event
            con_node = pipeline.CondorDAGNode(con_marg_job)
            con_node.add_macro("macroiteration",it)
            con_node.add_macro("macroevent",int(event))
            con_node.set_retry(opts.general_retries)
            for row in np.arange(indx_max_event): #np.arange(n_jobs_this_time):
                # Add task per MARG operation. MARG runs every iteration; the
                # tracer-vs-legacy choice only affects the MARG_PUFF lane below.
                if opts.marg_event_args:
                    cip_node = pipeline.CondorDAGNode(cip_job)
                else:
                    cip_node = pipeline.CondorDAGNode(cip_job_list[event])
                cip_node.set_retry(opts.general_retries)
                cip_node.add_macro("macroid",event)
                cip_node.add_macro("macroevent", int(row*n_chunk_here))
                cip_node.add_macro("macroiteration", it)
                if not(parent_fit_node is None):
                    cip_node.add_parent(parent_fit_node)
                con_node.add_parent(cip_node) # consolidate depends on all of the individual jobs
                dag.add_node(cip_node)
                # if PUFF is done for this event, also add MARG_PUFF nodes.  Note grid size assumed identical, so don't need another loop over 'id'/row
                # In tracer mode (puff_uses_marg_net) the tracer writes grid-{k+1}.dat
                # directly -- MARG already evaluates that grid, so no MARG_PUFF lane.
                if puff_args and puff_cadence and not puff_uses_marg_net:
                    if it>it_start and it <= puff_max_it  and (it-1)%puff_cadence ==0:  # we made a puffball last iteration, so run it through ILE now
                        if opts.marg_event_args:
                            cipPuff_node = pipeline.CondorDAGNode(cipPuff_job)
                        else:
                            cipPuff_node = pipeline.CondorDAGNode(cip_puff_job_list[event])
                        cipPuff_node.set_retry(opts.general_retries)
                        cipPuff_node.add_macro("macroid",int(event))
                        cipPuff_node.add_macro("macroevent", int(row*n_chunk_here))
                        cipPuff_node.add_macro("macroiteration", it)
                        if not(parent_fit_node is None):
                            cipPuff_node.add_parent(parent_fit_node)
                        con_node.add_parent(cipPuff_node) # consolidate depends on all of the individual jobs
                        dag.add_node(cipPuff_node)
                        
            # Add node to combine event into consolidated_*_*.net_marg in subdirectory, at end of for loop
            dag.add_node(con_node)
            # Add node to combine all events.   con_prod requires all con nodes done
            con_prod_node.add_parent(con_node)
    unify_node.add_parent(con_prod_node)

    # add con job, then unify
    dag.add_node(con_prod_node)
    dag.add_node(unify_node)
    parent_fit_node=unify_node

    # Create EOS_POST nodes. Explode out workers
    # Create fit node, which depends on consolidate node
    print(" Exploding workers out ")
        # Create job to consolidate worker outputs
    join_node =pipeline.CondorDAGNode(join_post_job)
    join_node.add_macro("macroiteration", it)
    join_node.add_macro("macroiterationnext", it+1)
    join_node.set_category("join_post")
    join_node.set_retry(opts.general_retries)
        
        # Create exploded worker job nodes
    n_explode = opts.eos_post_explode_jobs
    # if we are on the last iteration and we want to use more exploded jobs now, explode more jobs
    if opts.eos_post_explode_jobs_last and (it == opts.n_iterations-1): # last iteration
        print("   Last iteration explode size ", opts.eos_post_explode_jobs_last)
        n_explode = opts.eos_post_explode_jobs_last
    for indx in np.arange(n_explode):
                worker_node =pipeline.CondorDAGNode(eospost_marg_job)
                worker_node.add_macro("macroiteration", it)
                worker_node.add_macro("macroiterationnext", it+1)
                worker_node.set_category("EOS_post_worker")
                worker_node.add_parent(parent_fit_node)  # only fit if we have results from the previous iteration
                worker_node.set_retry(opts.general_retries)
                join_node.add_parent(worker_node)  # make sure to add worker node as parent
                dag.add_node(worker_node)
    dag.add_node(join_node)
    parent_fit_node=join_node

    # add test here, using output of this most recent iteration
    if opts.test_args and it>0:
        # Cannot run test on first iteration
        test_node = pipeline.CondorDAGNode(test_job)
        test_node.add_macro("macroiteration", it+1)   # test the NEWLY-PRODUCED iteration against the old
        test_node.add_macro("macroiterationlast", it)
        test_node.add_parent(parent_fit_node)
        test_node.set_category("CONVERGE")
        dag.add_node(test_node)
        test_node_list.append(test_node)

        parent_fit_node=test_node


    # Check if puffball being created, and if so create it *immediately* after the above join  creates the final 'grid-X.dat' file
    if puff_args and puff_cadence:
      if it > -1 and it <= puff_max_it and it%puff_cadence ==0:
        print(" Puffball for iteration ", it)
        puff_node = pipeline.CondorDAGNode(puff_job)
        puff_node.add_macro("macroiteration", it)
        puff_node.add_macro("macroiterationnext", it+1)
        puff_node.set_category("PUFF")
        puff_node.set_retry(opts.general_retries)
        if puff_uses_marg_net:
            # Tracer reads all.marg_net (output of UNIFY). Run in parallel with
            # EOS_POST -- the tracer's own grid output drives the next iteration,
            # EOS_POST is now a downstream-free posterior diagnostic.
            puff_node.add_parent(unify_node)
        else:
            if not (parent_fit_node is None):
                puff_node.add_parent(parent_fit_node)  # only fit if we have results from the previous iteration
        dag.add_node(puff_node)

        parent_fit_node = puff_node




dag_name="marginalize_hyperparameters"
dag.set_dag_file(dag_name)
dag.write_concrete_dag()
