"""
From a bed describing variant regions and non-var regions (prefixed by "nonvar_") and a directory containing one msa per var region in the bed, produce a vg prg. 

Assumptions: - the bed is ordered by genome chrom and pos
			 - each input msa has a ref fasta portion named 'ref'

These inputs are generated by the `make_prgs` workflow.

A note about vg construct calls: `-a` and `-S` options, for retaining alt paths and structural vars resp., don't make any difference to the produced graph, so not used.
"""
from pathlib import Path

WORKFLOW = "vg_make_prgs"


configfile: "analysis/configs/common.yaml"
configfile: "analysis/configs/vg_make_prgs.yaml"


container: config["container"]


include: "../common_utils.py"


# Paths
output_base = Path(f'{config["output_dir"]}/{WORKFLOW}/{config["prg_name"]}')
output_prgs = output_base / "prgs"
output_vcfs = output_base / "vcfs"

mk_output_dirs(dir())

# Get list of regions to build prgs of
nonvar_regions, all_regions = [], []
cs_var_regions = ""  # For use by bcftools merge
with open(config["full_bed"]) as f:
    for line in f:
        rows = line.split("\t")
        reg_name = rows[3].strip()
        if reg_name.startswith("nonvar_"):
            nonvar_regions.append(reg_name)
        else:
            cs_var_regions += f"{rows[0]}:{int(rows[1]) + 1}-{rows[2]},"
        all_regions.append(reg_name)


rule all:
    input:
        f"{output_base}/prg_msa.vg",
        f"{output_base}/prg_vcf.vg",


rule make_var_prgs:
    input:
        msa=Path(config["msa_dir"]) / "{region}.msa",
    output:
        prg=f"{output_prgs}/{{region}}.vg",
    resources:
        mem_mb=5000,
    shell:
        """
        vg construct -M {input.msa} > {output.prg}
        """


rule make_nonvar_prgs:
    input:
        full_bed=config["full_bed"],
        ref_genome=config["genome_fasta"],
    output:
        expand(f"{output_prgs}/{{region}}.vg", region=nonvar_regions),
    params:
        outdir=output_prgs,
    shadow:
        "shallow"
    shell:
        """
        IFS="\n"; for nonvar_line in $(grep -w "nonvar_.*" {input.full_bed}) 
        do
            IFS="\t"; elems=($nonvar_line)    
            adj_start=$((${{elems[1]}} + 1))
            prg_name="${{elems[3]}}"
            reg="${{elems[0]}}:${{adj_start}}-${{elems[2]}}"
            samtools faidx {input.ref_genome} $reg | \
                sed 's/>.*/>ref/' > ref_portion.fa
            vg construct -M ref_portion.fa > {params.outdir}/${{prg_name}}.vg
        done
        """


rule concat_prgs:
    input:
        expand(f"{output_prgs}/{{region}}.vg", region=all_regions),
    output:
        f"{output_base}/prg_msa.vg",
    resources:
        mem_mb=5000,
    shell:
        "vg concat {input} > {output}"


rule make_input_vcf:
    input:
        vcf_fofn=config["vcf_list"],
        ref_genome=config["genome_fasta"],
    output:
        merged=f"{output_vcfs}/merged.vcf.gz",
        indexes=expand(
            f"{output_vcfs}/merged.vcf.gz.{{suffix}}", suffix=["csi", "tbi"]
        ),
    params:
        regions=cs_var_regions,
    shadow:
        "shallow"
    shell:
        """
        bcftools merge -f PASS -r {params.regions} -Oz -o tmp.vcf.gz -l {input.vcf_fofn}
        bcftools norm -c x -f {input.ref_genome} -Oz -o {output.merged} tmp.vcf.gz
        bcftools index {output.merged}
        tabix {output.merged} # required for `vg call -v and graphtyper`
        """


rule prg_from_vcf:
    input:
        vcf=rules.make_input_vcf.output.merged,
        ref_genome=config["genome_fasta"],
    output:
        f"{output_base}/prg_vcf.vg",
    resources:
        mem_mb=5000,
    shell:
        "vg construct -r {input.ref_genome} -v {input.vcf} -a > {output}"
