###############

import glob as glob
import os
import re
import numpy as np

PRARR=[0,0.01,0.02,0.03,0.04,0.05,0.06,0.07,0.08,0.09,0.1,0.11,0.12,0.13,0.14,0.15,0.16,0.17,0.18,0.19,0.2]
PAARA=[0,0.01,0.02,0.03,0.04,0.05,0.06,0.07,0.08,0.09,0.1,0.11,0.12,0.13,0.14,0.15,0.16,0.17,0.18,0.19,0.2]
PRRRA=[0,0.01,0.02,0.03,0.04,0.05,0.06,0.07,0.08,0.09,0.1,0.11,0.12,0.13,0.14,0.15,0.16,0.17,0.18,0.19,0.2]
PRAAA=[0,0.01,0.02,0.03,0.04,0.05,0.06,0.07,0.08,0.09,0.1,0.11,0.12,0.13,0.14,0.15,0.16,0.17,0.18,0.19,0.2]
PAARR=[0,0.01,0.02,0.03,0.04,0.05,0.06,0.07,0.08,0.09,0.1,0.11,0.12,0.13,0.14,0.15,0.16,0.17,0.18,0.19,0.2]
PRRAA=[0,0.01,0.02,0.03,0.04,0.05,0.06,0.07,0.08,0.09,0.1,0.11,0.12,0.13,0.14,0.15,0.16,0.17,0.18,0.19,0.2]
id="halfSiblings1_g3-b4-i1"

rule all:
    input:
        expand('data/vcferr/{id}.prarr{prarr}_paara{paara}_prrra{prrra}_praaa{praaa}_paarr{paarr}_prraa{prraa}.vcf.gz',prarr=PRARR, paara=0,prrra=0,praaa=0,paarr=0,prraa=0,id=id),
        expand('data/vcferr/{id}.prarr{prarr}_paara{paara}_prrra{prrra}_praaa{praaa}_paarr{paarr}_prraa{prraa}.vcf.gz.tbi',prarr=PRARR,paara=0,prrra=0,praaa=0,paarr=0,prraa=0,id=id),
        expand('data/akt/{id}.prarr{prarr}_paara{paara}_prrra{prrra}_praaa{praaa}_paarr{paarr}_prraa{prraa}.akt',prarr=PRARR,paara=0,prrra=0,praaa=0,paarr=0,prraa=0,id=id),
        expand('data/vcferr/{id}.prarr{prarr}_paara{paara}_prrra{prrra}_praaa{praaa}_paarr{paarr}_prraa{prraa}.vcf.gz',prarr=0,paara=PAARA,prrra=0,praaa=0,paarr=0,prraa=0,id=id),
        expand('data/vcferr/{id}.prarr{prarr}_paara{paara}_prrra{prrra}_praaa{praaa}_paarr{paarr}_prraa{prraa}.vcf.gz.tbi',prarr=0,paara=PAARA,prrra=0,praaa=0,paarr=0,prraa=0,id=id),
        expand('data/akt/{id}.prarr{prarr}_paara{paara}_prrra{prrra}_praaa{praaa}_paarr{paarr}_prraa{prraa}.akt',prarr=0,paara=PAARA,prrra=0,praaa=0,paarr=0,prraa=0,id=id),
        expand('data/vcferr/{id}.prarr{prarr}_paara{paara}_prrra{prrra}_praaa{praaa}_paarr{paarr}_prraa{prraa}.vcf.gz',prarr=0,paara=0,prrra=PRRRA,praaa=0,paarr=0,prraa=0,id=id),
        expand('data/vcferr/{id}.prarr{prarr}_paara{paara}_prrra{prrra}_praaa{praaa}_paarr{paarr}_prraa{prraa}.vcf.gz.tbi',prarr=0,paara=0,prrra=PRRRA,praaa=0,paarr=0,prraa=0,id=id),
        expand('data/akt/{id}.prarr{prarr}_paara{paara}_prrra{prrra}_praaa{praaa}_paarr{paarr}_prraa{prraa}.akt',prarr=0,paara=0,prrra=PRRRA,praaa=0,paarr=0,prraa=0,id=id),
        expand('data/vcferr/{id}.prarr{prarr}_paara{paara}_prrra{prrra}_praaa{praaa}_paarr{paarr}_prraa{prraa}.vcf.gz',prarr=0,paara=0,prrra=0,praaa=PRAAA,paarr=0,prraa=0,id=id),
        expand('data/vcferr/{id}.prarr{prarr}_paara{paara}_prrra{prrra}_praaa{praaa}_paarr{paarr}_prraa{prraa}.vcf.gz.tbi',prarr=0,paara=0,prrra=0,praaa=PRAAA,paarr=0,prraa=0,id=id),
        expand('data/akt/{id}.prarr{prarr}_paara{paara}_prrra{prrra}_praaa{praaa}_paarr{paarr}_prraa{prraa}.akt',prarr=0,paara=0,prrra=0,praaa=PRAAA,paarr=0,prraa=0,id=id),
        expand('data/vcferr/{id}.prarr{prarr}_paara{paara}_prrra{prrra}_praaa{praaa}_paarr{paarr}_prraa{prraa}.vcf.gz',prarr=0,paara=0,prrra=0,praaa=0,paarr=PAARR,prraa=0,id=id),
        expand('data/vcferr/{id}.prarr{prarr}_paara{paara}_prrra{prrra}_praaa{praaa}_paarr{paarr}_prraa{prraa}.vcf.gz.tbi',prarr=0,paara=0,prrra=0,praaa=0,paarr=PAARR,prraa=0,id=id),
        expand('data/akt/{id}.prarr{prarr}_paara{paara}_prrra{prrra}_praaa{praaa}_paarr{paarr}_prraa{prraa}.akt',prarr=0,paara=0,prrra=0,praaa=0,paarr=PAARR,prraa=0,id=id),
        expand('data/vcferr/{id}.prarr{prarr}_paara{paara}_prrra{prrra}_praaa{praaa}_paarr{paarr}_prraa{prraa}.vcf.gz',prarr=0,paara=0,prrra=0,praaa=0,paarr=0,prraa=PRRAA,id=id),
        expand('data/vcferr/{id}.prarr{prarr}_paara{paara}_prrra{prrra}_praaa{praaa}_paarr{paarr}_prraa{prraa}.vcf.gz.tbi',prarr=0,paara=0,prrra=0,praaa=0,paarr=0,prraa=PRRAA,id=id),
        expand('data/akt/{id}.prarr{prarr}_paara{paara}_prrra{prrra}_praaa{praaa}_paarr{paarr}_prraa{prraa}.akt',prarr=0,paara=0,prrra=0,praaa=0,paarr=0,prraa=PRRAA,id=id)

rule vcferr:
    input: 'data/input/simulation.vcf.gz'
    output: 'data/vcferr/{id}.prarr{prarr}_paara{paara}_prrra{prrra}_praaa{praaa}_paarr{paarr}_prraa{prraa}.vcf.gz'
    shell:
        """ vcferr {input} \
            --sample='{wildcards.id}' \
            --p_rarr={wildcards.prarr} \
            --p_aara={wildcards.paara} \
            --p_rrra={wildcards.prrra} \
            --p_raaa={wildcards.praaa} \
            --p_aarr={wildcards.paarr} \
            --p_rraa={wildcards.prraa} \
            --output_vcf={output} """

rule tabix:
    input: 'data/vcferr/{id}.prarr{prarr}_paara{paara}_prrra{prrra}_praaa{praaa}_paarr{paarr}_prraa{prraa}.vcf.gz'
    output: 'data/vcferr/{id}.prarr{prarr}_paara{paara}_prrra{prrra}_praaa{praaa}_paarr{paarr}_prraa{prraa}.vcf.gz.tbi'
    shell:
        """ tabix {input} """

rule akt:
    input: vcf='data/vcferr/{id}.prarr{prarr}_paara{paara}_prrra{prrra}_praaa{praaa}_paarr{paarr}_prraa{prraa}.vcf.gz',
           tbi='data/vcferr/{id}.prarr{prarr}_paara{paara}_prrra{prrra}_praaa{praaa}_paarr{paarr}_prraa{prraa}.vcf.gz.tbi'
    output: 'data/akt/{id}.prarr{prarr}_paara{paara}_prrra{prrra}_praaa{praaa}_paarr{paarr}_prraa{prraa}.akt'
    shell:
        """ akt kin --force -M1 {input.vcf} > {output} """
