
# Snakemake pipeline to build a mitochondrial genome phylogenetic tree.


import pandas as pd

# file main_table_melt.txt contains gene names, accessions numbers, and name of the files to save the sequences
samples_df = pd.read_table('main_table_melt.txt').set_index("filename", drop=False)
SAMPLES = list(samples_df['filename'])
GENES   = list(samples_df['gene'])


rule all:
    input:
        "concatenated_alignment/concatenated_alignment_params.nex.con.tre"
        


##download_sequences: rule to download sequences
rule download_sequences:
    conda: "env.yml"
    input: "main_table_melt.txt"
    output: "sequences/{sample}.fasta" 
    params:
        # dynamically grab the accession ID and taxID from the "accession" and "taxID" columns in the data frame
        accession = lambda wildcards: samples_df.loc[wildcards.sample, "accession"],
        taxID = lambda wildcards: samples_df.loc[wildcards.sample, "taxID"]
    shell: """
        # fetch from NCBI
        efetch -db protein -id {params.accession} -format fasta > {output}
        
        # trick here, replacing accession by taxID as sequence name. This is useful downstream.
        sed 's/^>{params.accession}/>{params.taxID}/g' {output} > auxfile
        mv auxfile {output}
        """



##combine_sequences: group sequences by gene
rule combine_sequences:
   conda: "env.yml"
   input:
      expand("sequences/{sample}.fasta", sample=SAMPLES)
   output:
      "sequences_grouped/{gene}.fasta"
   shell: """
      cat sequences/*"_"{wildcards.gene}"_"* > {output}
   """


##align_sequences: rule to align sequences. Will generate one alignment per gene
rule align_sequences:
   conda: "env.yml"
   input:
      "sequences_grouped/{gene}.fasta"
   output:
      "sequences_aligned/{gene}_aligned.fasta"
   shell: """
      mafft --localpair --maxiterate 1000 {input} > {output}
   """



##trim_alignments: now trim the alignments
rule trim_alignments:
   conda: "env.yml"
   input:
      "sequences_aligned/{gene}_aligned.fasta"
   output:
      "sequences_aligned_trimmed/{gene}_trimmed.fasta"
   shell: """
      trimal -nogaps -in {input} > {output}
   """



##concatenate_alignments: condatenate the alignments into a superalignment
rule concatenate_alignments:
   conda: "env.yml"
   input:
      expand("sequences_aligned_trimmed/{gene}_trimmed.fasta", gene=GENES)
   output:
      "concatenated_alignment/concatenated_alignment.fasta",
      directory("concatenated_alignment")
   shell: """
      pxcat -s sequences_aligned_trimmed/*.fasta -p concatenated_alignment/pxcat_parts.txt -o {output}
   """



##convert_fasta_nexus: rule to convert alignment in fasta to nexus
rule convert_fasta_nexus:
   conda: "env.yml"
   input:
      "concatenated_alignment/concatenated_alignment.fasta"
   output:
      "concatenated_alignment/concatenated_alignment.nex"
   shell: """
      pxs2nex -s{input} -o {output}
   """



##add_mrbayes_params: rule to add mrbayes parameters at the end of the nexus file
rule add_mrbayes_params:
   conda: "env.yml"
   input:
      "concatenated_alignment/concatenated_alignment.nex"
   output:
      "concatenated_alignment/concatenated_alignment_params.nex"
   shell: """
      cat {input} > {output}
      echo "
begin mrbayes;
    Prset Aamodelpr = mixed;
    outgroup 1174673;
    mcmc nchains = 4 ngen = 500000 samplefreq=200 printfreq=200;
    sumt relburnin = yes burninfrac = 0.25;
    sump relburnin = yes burninfrac = 0.25;
    sump;
    sumt;
end;
" >> {output}
   """
 

##call_mrbayes:
rule call_mrbayes:
   conda: "env.yml"
   input:
      "concatenated_alignment/concatenated_alignment_params.nex"
   output:
      "concatenated_alignment/concatenated_alignment_params.nex.con.tre"
   shell: """
      mb {input}
   """
