#!/home/conda/feedstock_root/build_artifacts/bld/rattler-build_pycbc_1768936173/host_env_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehol/bin/python

"""
Program for running a faithfulness comparisons workflow analysis between two approximants
and generate files containing the match between them and plots.
"""

import pycbc.workflow as wf
import argparse

from pycbc import add_common_pycbc_options, init_logging
from pycbc.workflow.plotting import PlotExecutable

def make_faithsim_plot(workflow, analysis_time, input_file, out_dir, tags=None):
    tags = [] if tags is None else tags
    secs = workflow.cp.get_subsections("pycbc_faithsim_plots")
    files = wf.FileList([])
    for tag in secs:
        node = PlotExecutable(
            workflow.cp,
            "pycbc_faithsim_plots",
            ifos=["X1"],
            out_dir=out_dir,
            tags=[tag] + tags,
        ).create_node()
        node.add_input_opt("--input-file", input_file)
        node.new_output_file_opt(analysis_time, ".png", "--output-plot")
        workflow += node
        files += node.output_files

    return files


class CreateInjectionExecutable(wf.Executable):
    """Class for running create_inj_file"""

    current_retention_level = wf.Executable.ALL_TRIGGERS

    def create_node(self, analysis_time, extra_tags=None):
        if extra_tags is None:
            extra_tags = []

        node = wf.Executable.create_node(self)

        node.new_output_file_opt(
            analysis_time, ".xml", "--output", tags=self.tags + extra_tags
        )
        return node


class FaithsimExecutable(wf.Executable):
    """Class for running pycbc_faithsim"""

    current_retention_level = wf.Executable.ALL_TRIGGERS

    def create_node(self, analysis_time, param_file, extra_tags=None):
        if extra_tags is None:
            extra_tags = []
        node = wf.Executable.create_node(self)

        node.add_input_opt("--param-file", param_file)

        node.new_output_file_opt(
            analysis_time, ".dat", "--match-file", tags=self.tags + extra_tags
        )
        return node


class CollectResultsExecutable(wf.Executable):
    """Class for collecting the results of the faithsim script"""

    current_retention_level = wf.Executable.ALL_TRIGGERS

    def create_node(self, analysis_time, match_list, bank_list, extra_tags=None):
        if extra_tags is None:
            extra_tags = []
        node = wf.Executable.create_node(self)

        for match in match_list:
            node.add_input_opt("--match-inputs", match)
        for bank in bank_list:
            node.add_input_opt("--bank-inputs", bank)

        node.new_output_file_opt(
            analysis_time, ".dat", "--output", tags=self.tags + extra_tags
        )

        return node


parser = argparse.ArgumentParser(description=__doc__)
add_common_pycbc_options(parser)
wf.add_workflow_command_line_group(parser)
wf.add_workflow_settings_cli(parser)
args = parser.parse_args()

init_logging(args.verbose)

workflow = wf.Workflow(args)

num_banks = workflow.cp.get("splitbank", "num_banks")

injections_exe = CreateInjectionExecutable(
    workflow.cp,
    "lalapps_inspinj",
    ifos=["X1"],
    out_dir=workflow.output_dir + "injections",
    tags=["inj"],
)

inj_node = injections_exe.create_node(workflow.analysis_time)
workflow += inj_node

inj = inj_node.output_files[0]

split_exe = wf.PycbcSplitBankXmlExecutable(
    workflow.cp, "pycbc_splitbank", num_banks=num_banks, out_dir=workflow.output_dir +"bank"
)
splitbank_node = split_exe.create_node(inj)

workflow += spltbank_node

faithsim_exe = FaithsimExecutable(
    workflow.cp, "pycbc_faithsim", ifos=["X1"], out_dir=workflow.output_dir +"match"
)

collect_exe = CollectResultsExecutable(
    workflow.cp,
    "pycbc_faithsim_collect_results",
    ifos=["X1"],
    out_dir=workflow.output_dir +"collect_results",
)

faithsim_files = wf.FileList([])

for i, bank in enumerate(splitbank_node.output_files):
    faithsim_node = faithsim_exe.create_node(
        workflow.analysis_time,
        param_file=bank,
        extra_tags=[f"match{i}"],
    )
    workflow += faithsim_node
    faithsim_files.append(faithsim_node.output_file)

collect_node = collect_exe.create_node(
    workflow.analysis_time,
    faithsim_files,
    spltbank_node.output_files,
    extra_tags=["result"],
)
workflow += collect_node
collect_results = collect_node.output_files[0]

make_faithsim_plot(
    workflow,
    workflow.analysis_time,
    collect_results,
    out_dir=workflow.output_dir +"plots",
    tags=None,
)

workflow.save()
