
myconfig = "config.yml"
configfile: myconfig
singularity: "../ldcpipe.sif" 

rundir = config["dirname"]
source = config["sources"]
lbatch = ["%d:%d"%(i, config["nbatch"]) for i in range(config["nbatch"])]

rule all:
    input:
        expand(rundir+"/{source}-y.h5", source=source),

rule source_config:
    output: ["param_dgb.yml", 'param_igb.yml']
        
        
rule source_selection:
    input:
        cfg = "param_{source}.yml",
    output:
        rundir+"/{source}.h5"
    shell:
        """
        source_selection -c {input.cfg} -o {output}
        """

rule orbits:
    input: myconfig
    output: rundir+"/orbits.h5"
    shell:
       """
       build_orbits -o {output} -c {input}
       """
        
rule arm_projection:
    input:
        cfg = "config.yml",
	src = rundir+"/{source}.h5",
        orbits = rules.orbits.output,
    output: temp(rundir+"/{source}--{batch}-y.h5")
    log: "log/{source}--{batch}-y.log"
    threads: 1
    params:
        cfg = "param_{source}.yml",
        checkpoint = config["checkpoint"]
    shell:
        """
        export MKL_NUM_THREADS=1
        export OPENBLAS_NUM_THREADS=1
        arm_projection -i {input.src} -c {input.cfg} -o {output} --select {wildcards.batch} --log {log} --source-config {params.cfg} --checkpoint {params.checkpoint} --orbits {input.orbits}
        """

rule batch_merge:
    #input: lambda wildcards: [rundir+"/{source}--%d:%d-y.h5"%(i,config["nbatch"]) for i in range(config["nbatch"])]
    input:
        src = expand(rundir+"/{source}--{batch}-y.h5", source=source, batch=lbatch),
        cfg = "config.yml"
    output:
        rundir+"/{source}-y.h5"
    log: "log/{source}-y.log"
    shell:
        """
        strain_combination -i {input.src} -c {input.cfg} -o {output} --log {log}
        """
