"""
This workflow takes bed files listing regions + where they occur, a ref genome, a list of matched sample reads and pacbio assemblies and initial prgs to start from and runs:
    - validate input variant regions are present in truth assemblies
    - {vg, gramtools} prg index building
    - {vg, gramtools} mapping/genotyping
    - vcf merging per condition + validate input var regions present in called samples
    - edit distance of called regions to truth assemblies

The workflow has some rules similar/identical to pacb_ilmn_validation workflow, but is kept distinct as it does not assess gramtools perf against other standard var callers, but against other genome graph tool(s).
"""
WORKFLOW = "tb_bigdel"
VALIDATION_WORKFLOW = "pacb_ilmn_validation"


configfile: "analysis/configs/common.yaml"
configfile: f"analysis/configs/tb_bigdel.yaml"


container: config["container"]


# Include convenience functions
include: f"../{VALIDATION_WORKFLOW}/utils.py"
include: "../common_utils.py"


GMTOOLS_COMMIT = get_gmtools_commit()

# _____Set up output paths______#
output_gram_build = Path(f'{config["gram_builds_dir"]}/{WORKFLOW}/{GMTOOLS_COMMIT}')
output_vg_build = Path(f'{config["vg_builds_dir"]}/{WORKFLOW}')
output_base = Path(f'{config["output_dir"]}/{WORKFLOW}/')

output_bowtie_indexes = output_base / "bowtie_indexes"
output_regions = output_base / "input_regions"
output_minimap2 = output_regions / "minimap2"
output_genotyped = output_base / "genotyped"
output_prg_desc = output_base / "prg_descriptions"
output_alignments = output_base / "alignments" / GMTOOLS_COMMIT
output_ref_index = output_base / "alignments" / "ref_index"
output_ref_alignments = output_base / "alignments" / "ilmn_alignments_ref"
output_plots = output_base / f"plots/{GMTOOLS_COMMIT}"
output_varifier = output_base / "varifier"

conditions = [f"gramtools_{GMTOOLS_COMMIT}", "vg", "graphtyper2", "baseline_ref"]
output_gmtools = output_genotyped / conditions[0]
output_vg_mapped = output_genotyped / conditions[1] / "mapped_reads"
output_graphtyper = output_genotyped / conditions[2]
output_graphtyper_original = output_graphtyper / "original"
output_baseline_ref = output_genotyped / conditions[3]

mk_output_dirs(dir())

SAMPLES = get_samples(config["sample_tsv"])
ref_name = Path(config["starting_prg"]["fasta_ref"]).stem

with open(config["beds"]["with_flank"]) as fin:
    num_regions = len(fin.readlines())

def get_vcf_to_analyse_varifier(wildcards):
    if wildcards.condition == "input_variants":
        return f'{config["input_vcf_dir"]}/{wildcards.sample}.vcf.gz'
    else:
        return f"{output_genotyped}/{wildcards.condition}/{wildcards.sample}.vcf.gz",


localrules:
    all,
    tb_baseline_ref_genotype,


# Include rule sets
include: "validate_input_dels.smk"
include: "genotype.smk"


filters = {
    "callsunfiltered": "",
    "callsmindepth5": "-i 'FORMAT/DP>5'",
    "callsfilterpass": '-i \'FILTER=="PASS" || FORMAT/FT=="PASS"\'',
}


rule all:
    input:
        genotyped=expand(
            f"{output_genotyped}/{{condition}}_{{ftype}}",
            condition=conditions[0:3],
            ftype=["merged.vcf.gz", "validated.tsv"],
        ),
        validated=f"{output_regions}/minimap2_validated.tsv",
        variant_counts=expand(
            f"{output_prg_desc}/{{condition}}_variants.tsv", condition=conditions[0:2]
        ),
        plots_mapped_regions=expand(
            f"{output_plots}/{{mapper}}/{{calls}}/upset_plot_{GMTOOLS_COMMIT}_unfiltered.pdf",
            mapper=["bowtie2", "minimap2"],
            calls=filters.keys(),
        ),
        varifier_plots=expand(
                f"{output_plots}/varifier/{{filtering}}/genome_graph_impact.pdf",
                filtering=["nofilter","filterpass"],
        )

rule index_pacb_assemblies:
    input:
        assembly=get_assembly,
    output:
        expand(
            f"{output_bowtie_indexes}/{{sample}}.{{ext}}",
            ext=bowtie2_idx_extensions,
            allow_missing=True,
        ),
    shell:
        f"bowtie2-build {{input.assembly[0]}} {output_bowtie_indexes}/{{wildcards.sample}}"


rule index_reference:
    output:
        ref_idx=expand(
            f"{output_ref_index}/{ref_name}.{{ext}}", ext=bowtie2_idx_extensions
        ),
    shell:
        f'bowtie2-build {config["starting_prg"]["fasta_ref"]} {output_ref_index}/{ref_name}'


rule tb_map_reads_to_ref:
    """Used by graphtyper genotyper and for manually validating some input dels"""
    input:
        reads_files=get_reads,
        ref_idx=rules.index_reference.output.ref_idx,
    output:
        bam=f"{output_ref_alignments}/{{sample}}.bam",
        bai=f"{output_ref_alignments}/{{sample}}.bam.bai",
    shadow:
        "shallow"
    threads: 10
    params:
        idx_prefix=f"{output_ref_index}/{ref_name}",
    shell:
        """
        bowtie2 -x {params.idx_prefix} -1 {input.reads_files[0]} -2 {input.reads_files[1]} -S tmp.sam -p {threads}
        samtools sort tmp.sam -O BAM -o {output.bam}
        samtools index {output.bam}
        """


rule describe_vcf_prg:
    input:
        input_regions=config["beds"]["merged"],
        vg_vcf=f"{output_genotyped}/{conditions[1]}_merged.vcf.gz",
    output:
        desc=f"{output_prg_desc}/{conditions[1]}_variants.tsv",
    params:
        script=f'{config["scripts"]}/{WORKFLOW}/count_prg_variants.py',
    shell:
        "python3 {params.script} {input.vg_vcf} {input.input_regions} {output.desc}"


rule describe_jvcf_prg:
    input:
        input_regions=config["beds"]["merged"],
        gram_jvcf=(
            f"{output_genotyped}/{conditions[0]}/{SAMPLES[0]}/genotype/genotyped.json"
        ),
    output:
        desc=f"{output_prg_desc}/{conditions[0]}_variants.tsv",
    params:
        script=f'{config["scripts"]}/{WORKFLOW}/count_prg_variants.py',
    shell:
        "python3 {params.script} {input.gram_jvcf} {input.input_regions} {output.desc}"


rule tb_produce_gene_portions:
    input:
        result_vcf=f"{output_genotyped}/{{condition}}/{{sample}}.vcf.gz",
        var_regions=config["beds"]["with_flank"],
        fasta_ref=config["starting_prg"]["fasta_ref"],
    output:
        gene_portions=(
            f"{output_alignments}/gene_portions/{{calls}}/{{condition}}_{{sample}}.fa"
        ),
    params:
        idx_prefix=f"{output_bowtie_indexes}/{{sample}}",
        applied_filter=lambda wildcards: filters[wildcards.calls],
    shadow:
        "shallow"
    shell:
        """
        mkdir -p $(dirname {output.gene_portions}) 
        > {output.gene_portions}
        bcftools filter {params.applied_filter} {input.result_vcf} -Oz -o used.vcf.gz
        bcftools index used.vcf.gz
        IFS="\n"; for gene_line in $(cat {input.var_regions})
        do
            IFS="\t"; elems=($gene_line)    
            adjusted_start=$((${{elems[1]}} + 1))
            reg="${{elems[0]}}:${{adjusted_start}}-${{elems[2]}}"
            gene_name=${{elems[3]}}
            samtools faidx {input.fasta_ref} $reg | bcftools consensus -s {wildcards.sample} used.vcf.gz |
                sed 's/>.*/>'${{gene_name}}'/' >> {output.gene_portions}
        done
        """


rule tb_align_gene_portions_bowtie2:
    input:
        truth_ref_idx=expand(
            f"{output_bowtie_indexes}/{{sample}}.{{ext}}",
            ext=bowtie2_idx_extensions,
            allow_missing=True,
        ),
        gene_portions=rules.tb_produce_gene_portions.output,
    output:
        alignment_file=(
            f"{output_alignments}/bowtie2/{{calls}}/{{condition}}_{{sample}}.sam"
        ),
    params:
        idx_prefix=f"{output_bowtie_indexes}/{{sample}}",
    resources:
        mem_mb=4000,
    shell:
        """
        mkdir -p $(dirname {output.alignment_file})
        bowtie2 -x {params.idx_prefix} -U {input.gene_portions} -f > {output.alignment_file}
        """


rule tb_align_gene_portions_minimap2:
    input:
        assembly=get_assembly,
        gene_portions=rules.tb_produce_gene_portions.output,
    output:
        ori_alignment_file=f"{output_alignments}/minimap2/{{calls}}/original/{{condition}}_{{sample}}.sam",
        alignment_file=(
            f"{output_alignments}/minimap2/{{calls}}/{{condition}}_{{sample}}.sam"
        ),
    resources:
        mem_mb=4000,
    shell:
        """
        mkdir -p $(dirname {output.ori_alignment_file})
        minimap2 -a {input.assembly} {input.gene_portions} > {output.ori_alignment_file}
        """
        f"""
        python3 {config["scripts"]}/{WORKFLOW}/add_NW_eddist.py {{output.ori_alignment_file}} {{input.assembly}} {{output.alignment_file}} --num_seqs {num_regions}
        """


rule tb_process_alignments:
    input:
        alignments=expand(
            f"{output_alignments}/{{mapper}}/{{calls}}/{{condition}}_{{sample}}.sam",
            condition=conditions,
            sample=SAMPLES,
            allow_missing=True,
        ),
        var_regions=config["beds"]["with_flank"],
        mask_bed=config["beds"]["ilmn_asm_mask"],
    output:
        stats=f"{output_plots}/{{mapper}}/{{calls}}_stats.tsv",
    shadow:
        "shallow"
    shell:
        f"""
        python3 {config["scripts"]}/{VALIDATION_WORKFLOW}/process_alignments.py {output_alignments}/{{wildcards.mapper}}/{{wildcards.calls}} {{input.var_regions}} {{output.stats}}
        """


rule tb_plot_alignments:
    input:
        stats=f"{output_plots}/{{mapper}}/{{calls}}_stats.tsv",
    output:
        expand(
            f"{output_plots}/{{mapper}}/{{calls}}/{{plot_name}}_plot_{GMTOOLS_COMMIT}_unfiltered.pdf",
            plot_name=["ecdf", "upset"],
            allow_missing=True,
        ),
    shell:
        f"""
        mkdir -p {output_plots}/{{wildcards.mapper}}/{{wildcards.calls}}
        Rscript {config["scripts"]}/{WORKFLOW}/plot_alignments.R {{input.stats}} {output_plots}/{{wildcards.mapper}}/{{wildcards.calls}} {GMTOOLS_COMMIT}
        """

rule tb_run_varifier:
    """Breaks down precision/recall by variant type and size"""
    input:
        result_vcf=get_vcf_to_analyse_varifier,
        assembly=get_assembly,
        fasta_ref=config["starting_prg"]["fasta_ref"],
        var_regions=config["beds"]["with_flank"],
    output:
        varifier_json=f"{output_varifier}/{{filtering}}/{{sample}}/{{condition}}/summary_stats.json",
        varifier_stats=f"{output_varifier}/{{filtering}}/{{sample}}/stats_{{condition}}.tsv",
    shadow:
        "shallow"
    resources:
        mem_mb=20000
    params:
        jvcf_converter=f'{config["scripts"]}/{WORKFLOW}/jvcf_to_vcf.py',
        postprocess_vcf=f'{config["scripts"]}/{WORKFLOW}/postprocess_vcf.py',
        varifier_stats=f'{config["scripts"]}/{WORKFLOW}/get_varifier_perf_per_record.py',
        varifier_run=f"{output_varifier}/{{filtering}}/{{sample}}/{{condition}}",
    shell:
        """
        input_vcf={input.result_vcf}
        # Convert gramtools jvcf to vcf
        if [[ {wildcards.condition} =~ gramtools.+ ]]; then
            jvcf=$(dirname {input.result_vcf})/{wildcards.sample}/genotype/genotyped.json
            python3 {params.jvcf_converter} $jvcf input.vcf
            input_vcf=input.vcf
        fi

        python3 {params.postprocess_vcf} {input.fasta_ref} $input_vcf used_vcf.vcf --remove_ref_and_null
        gzip -dc {input.assembly} > assembly.fa
        filtering=""
        if [[ {wildcards.filtering} == "filterpass" ]]; then
            filtering="--filter_pass .,PASS"
        fi
        varifier vcf_eval assembly.fa {input.fasta_ref} used_vcf.vcf {params.varifier_run} --force $filtering
        python3 {params.varifier_stats} --tool_name {wildcards.condition} --sample_name {wildcards.sample} --region_file {input.var_regions} {params.varifier_run} {output.varifier_stats}
        """


rule tb_concat_varifier_stats:
    input:
        expand(f"{output_varifier}/{{filtering}}/{{sample}}/stats_{{condition}}.tsv",
                sample=SAMPLES,
                condition=conditions[0:3] + ["input_variants"], allow_missing=True)
    output:
        f"{output_plots}/varifier/{{filtering}}/stats_all.tsv",
    shell:
        f"""
        python3 {config["scripts"]}/{WORKFLOW}/get_varifier_perf_per_record.py --header_only > {{output}}
        cat {{input}} >> {{output}}
        """

rule tb_plot_varifier_stats:
    input:
        varifier_stats=rules.tb_concat_varifier_stats.output,
        var_regions=config["beds"]["with_flank"],
    output:
        f"{output_plots}/varifier/{{filtering}}/genome_graph_impact.pdf"
    params:
        script=f'{config["scripts"]}/{WORKFLOW}/plot_varifier.R',
        output_dir=f"{output_plots}/varifier/{{filtering}}"
    shell:
        """
        Rscript {params.script} {input.varifier_stats} {params.output_dir} {input.var_regions}
        """

