shell:
    """
    rm -rf {params.outdir}
    dmesg -T
    spades.py --meta -1 {input.fwd} -2 {input.rev} \
    -o {params.outdir} -t {threads} -m {resources.mem}
    """

shell:
    """
    rm -rf {params.outdir}
    metawrap binning -t {threads} -m {resources.mem} \
    -a {params.assembly} --metabat2 --maxbin2 --concoct \
    -l {params.mincontiglength} -o {params.outdir} {params.reads}
    touch {output}
    """

shell:
    """
    scripts/rename_multifasta_prefix.py -f {input.MAG} -p {params.name} > {output}
    """

shell:
    """
    scripts/rename_multifasta_prefix.py -f {input.MAG} -p {params.name} > {output}
    """

shell:
    """
    prokka {input} --kingdom Bacteria --outdir {params.out_prokka} \
    --prefix {params.prefix} --force --locustag {params.prefix} --cpus {threads}
    """

shell:
    """
    scripts/cmseq.sh -t {threads} -i {input.r1} -n {input.r2} -r {input.MAG} -g {input.prokka} -o {params.name}
    """

shell:
    """
    rm -f {output.cmseq}
    rm -f {output.done}
    for i in {params.r1}; do run=$(basename ${{i}} _1.fastq); scripts/cmseq.sh\
    -t {threads} -i ${{i}} -n ${{i%%_1.fastq}}_2.fastq -r {input.MAG} -g {input.prokka}\
    -o {params.name}_${{run}}; awk 'NR==2' {params.name}_${{run}}.cmseq.csv >> {output.cmseq};done

    touch {output.done}
    """

run:
    if os.path.exists(str(output)):
        os.remove(str(output))
    else:
        print("Aggregating single run strain heterogeneity results")
    path = str(params)
    file1 = open(str(output), "w")
    for filename in glob.glob(os.path.join(path, "*.csv")):
        with open(filename, "r") as f:
            lines = f.readlines()
            if len(lines) > 1:
                lines_sub = lines[1].strip()
                L = filename + "\t" + lines_sub + "\n"
                file1.writelines(L)
    file1.close()

run:
    if os.path.exists(str(output)):
        os.remove(str(output))
    else:
        print("Aggregating coassembly strain heterogeneity results")
    path = str(params)
    file1 = open(str(output), "w")
    for filename in glob.glob(os.path.join(path, "*SRR*.csv")):
        with open(filename, "r") as f:
            lines = f.readlines()
            if len(lines) > 1:
                lines_sub = lines[1].strip()
                L = filename + "\t" + lines_sub + "\n"
                file1.writelines(L)
    file1.close()

shell:
    """
    cat {input}>{output}         
    """

shell:
    """
    cat {input}>{output}
    """

shell:
    """
    Rscript scripts/plotting/plot_cmseq.R {input.sr} {input.coas}
    """

shell:
    """
    cat {input.r1} > {output.r1}
    cat {input.r2} > {output.r2}
    """

shell:
    """
    repair.sh in={input.fwd} in2={input.rev} out={output.fwd} out2={output.rev} repair
    rm -f {input.fwd}
    rm -f  {input.rev}
    """

shell:
    """
    rm -f {params.fwd}
    rm -f {params.rev}
    rm -f {output.fwd}
    rm -f {output.rev}
    gzip {input.fwd}
    gzip {input.rev}
    mv {params.fwd} {output.fwd}
    mv {params.rev} {output.rev} 
    """

run:
    outfile = str(output)
    if os.path.exists(outfile):
        os.remove(outfile)
    with open(outfile, "w") as outf:
        outf.writelines("Run\tReadcount\n")
        for run in input:
            readcount = int(linecount(run))
            line = run + "\t" + str(readcount)
            outf.writelines(line + "\n")

shell:
    """
    sort -gk3 {input.mashdist}|sed -n 1p >{output}
    """

shell:
    """
    while read col1 col2 rem
      do
        echo 'dnadiff ${{col1}} ${{col2}} -p ${{col1%%.fasta}}_${{col2%%.fa}}_'
        dnadiff ${{col1}} ${{col2}} -p {params.outdir}/{params.bins}
      done < {input.bestmash}
    """

run:
    outfile = str(output)
    f = open(input.dnadiff)
    data = f.read()
    first_line = data.split("\n", 1)[0]
    a = first_line.split(" ")
    ref = a[0]
    quer = a[1]
    with open(outfile, "w") as outf:
        path_dna = input.dnadiff
        base = os.path.basename(path_dna)
        base = base.split(".report")[0]
        with open(path_dna) as f:
            for line in f:
                if "TotalBases" in line:
                    cols = line.split()
                    lenref = int(cols[1])
                    lenquer = int(cols[2])
                if "AlignedBases" in line:
                    cols = line.split()
                    aliref = cols[1].split("(")[-1].split("%")[0]
                    alique = cols[2].split("(")[-1].split("%")[0]
                if "AvgIdentity" in line:
                    cols = line.split()
                    ident = float(cols[1])
        line = "%s\t%s\t%i\t%.2f\t%i\t%.2f\t%.2f" % (ref, quer, lenref, float(aliref), lenquer, float(alique), float(ident))
        outf.writelines(line + "\n")

shell:
    """
    cat {input}>{output}
    """

shell:
    """
    cat {input}>{output}
    """

shell:
    """
    cat {input}>{output}
    """

shell:
    """
    Rscript scripts/plotting/dnadiff_plot.R {input.checkm_sr} {input.checkm_coas} {input.summ}
    """

shell:
    """
    awk 'FNR!=1' {input.checkm} {input.checkm_coas}>{params.checkm}
    echo -e "genome,completeness,contamination,strain_heterogeneity" | cat - {params.checkm} > {output}
    rm {params.checkm}
    """

shell:
    """
    rm -rf {params.outdir}
    mkdir -p {params.indir}
    rsync {params.singlerun_dRep}/* {params.indir}
    rsync {params.all_coas}/* {params.indir}
    dRep dereplicate -p {threads} {params.outdir} -g {params.indir}/*.fa \
    -pa 0.9 -sa 0.95 -nc 0.30 -cm larger --genomeInfo {input.checkm} -comp 50 -con 5
    """

shell:
    """
    Rscript scripts/plotting/plot_gtdb.R {input}
    """

run:
    with open(str(output), "w") as outf:
        for i in input:
            run = str(i).split("singlerun/")[1]
            run = run.split("/mapreads/")[0]
            dict = eval(open(str(i), "r").read())
            supplementary = dict["pass_qc"]["supplementary"]
            secondary = dict["pass_qc"]["secondary"]
            mapped = dict["pass_qc"]["mapped"]
            total = dict["pass_qc"]["total"]
            perc = (mapped - supplementary - secondary) 
            outf.write(str(run) + "\t" + str(perc) + "\n")
    outf.close()

shell:
    """
    cat {params.indir}/*fa>{output}
    bwa index {output}
    """

shell:
    """
    cat {params.indir}/*.fa>{output}
    bwa index {output}
    """

shell:
    """
    bwa mem -t {threads} {input.catalogue} {input.fwd} {input.rev} \
    | samtools view -bS - | samtools sort -@ {threads} -o {params.alignedsorted} -
    samtools index {params.alignedsorted}
    samtools flagstat {params.alignedsorted} > {output.flagstat}
    rm {params.alignedsorted}
    """

shell:
    """
    bwa mem -t {threads} {input.catalogue} {input.fwd} {input.rev}\
    | samtools view -bS - | samtools sort -@ {threads} -o {params.alignedsorted} -
    samtools index {params.alignedsorted}
    samtools flagstat {params.alignedsorted} > {output.flagstat}
    rm {params.alignedsorted}
    """

shell:
    """
    crimson flagstat {input}>{output}
    """

shell:
    """
    crimson flagstat {input}>{output}
    """

run:
    with open(str(output), "w") as outf:
        for i in input:
            base = os.path.basename(str(i))
            run = os.path.splitext(base)[0]
            dict = eval(open(str(i), "r").read())
            supplementary = dict["pass_qc"]["supplementary"]
            secondary = dict["pass_qc"]["secondary"]
            mapped = dict["pass_qc"]["mapped"]
            total = dict["pass_qc"]["total"]
            perc = (mapped - supplementary - secondary) 
            outf.write(str(run) + "\t" + str(perc) + "\n")
    outf.close()

run:
    with open(str(output), "w") as outf:
        for i in input:
            base = os.path.basename(str(i))
            run = os.path.splitext(base)[0]
            dict = eval(open(str(i), "r").read())
            supplementary = dict["pass_qc"]["supplementary"]
            secondary = dict["pass_qc"]["secondary"]
            mapped = dict["pass_qc"]["mapped"]
            total = dict["pass_qc"]["total"]
            perc = (mapped - supplementary - secondary) 
            outf.write(str(run) + "\t" + str(perc) + "\n")
    outf.close()

shell:
    """
    Rscript scripts/plotting/plot_framework.R {input.readcounts} {input.flagstat} {input.mapreads_sr} {input.mapreads_coas} {params.summary}
    """

shell:
    """
    fastqc {input.rev} --outdir {params.outdir}
    """

shell:
    """
    fastqc {input.fwd} --outdir {params.outdir}
    """

shell:
    """
    fastqc {input.rev} --outdir {params.outdir}
    """

shell:
    """
    multiqc --force {params.indir} -o {params.outdir}
    mv {params.outfile} {output}
    """

run:
    outfile = str(output)
    if os.path.exists(outfile):
        os.remove(outfile)
    with open(outfile, "w") as outf:
        outf.writelines("Run\tReadcount\n")
        for run in input:
            readcount = int(linecount(run))
            line = run + "\t" + str(readcount)
            outf.writelines(line + "\n")

shell:
    """
    rm -rf {params.folder}
    mkdir -p {params.folder}
    cp {params.refined_folder}/*.fa {params.folder}
    """

shell:
    """
    mkdir -p {params.dir1}
    mkdir -p {params.dir2}
    cp {input} {params.out1}
    cp {input} {params.out2}
    """

shell:
    """
    touch {output}
    """

shell:
    """
    checkm data setRoot ~/checkm_database/
    rm -rf {params.outdir}
    checkm lineage_wf -t {threads} -x {params.ext} --tab_table -f {output} {params.indir} {params.outdir}
    """

shell:
    """
    sed -i '1d' {input}
    cut -f1,12,13,14 {input} | tr '\t' ','>{params.checkm}
    sed 's/,/.fa,/' {params.checkm}>{params.checkm2}
    echo -e "genome,completeness,contamination,strain_heterogeneity" | cat - {params.checkm2} > {output}
    rm {params.checkm}
    rm {params.checkm2}
    """

shell:
    """
    checkm data setRoot ~/checkm_database/ 
    rm -rf {params.outdir}
    metawrap bin_refinement -o {params.outdir} -t {threads} -A {params.metabat} -B {params.maxbin} -C {params.concoct} -c 50 -x 10
    touch {output}
    """

shell:
    """
    rm -rf {params.folder}
    mkdir -p {params.folder}
    cp {params.refined_folder}/*.fa {params.folder}
    """

shell:
    """
    mkdir -p {params.dir1}
    mkdir -p {params.dir2}
    cp {input} {params.out1}
    cp {input} {params.out2}
    """

shell:
    """
    touch {output}
    """

shell:
    """
    checkm data setRoot ~/checkm_database/
    rm -rf {params.outdir}
    checkm lineage_wf -t {threads} -x {params.ext} --tab_table -f {output} {params.indir} {params.outdir}
    """

shell:
    """
    sed -i '1d' {input}
    cut -f1,12,13,14 {input} | tr '\t' ','>{params.checkm}
    sed 's/,/.fa,/' {params.checkm}>{params.checkm2}
    echo -e "genome,completeness,contamination,strain_heterogeneity" | cat - {params.checkm2} > {output}
    rm {params.checkm}
    rm {params.checkm2}
    """

shell:
    """
    Rscript scripts/plotting/plot_checkm_mags.R {input.sr} {input.coas}
    """

shell:
    """
    echo {threads}
    prefetch {params.run} && vdb-validate {params.run} && parallel-fastq-dump --threads {threads} \
    --outdir {params.outdir} --skip-technical --split-3 --sra-id {params.run} --gzip
    """

usage()
{
cat << EOF
usage: $0 options

CMseq workflow to infer strain heterogeneity of MAG.

'bsub' with '-M 5000'

OPTIONS:
   -t      Number of threads [REQUIRED]
   -i      Input first or only read (.fastq or .fastq.gz format) [REQUIRED]
   -n      Input reverse read (.fastq or fastq.gz [OPTIONAL]
   -r      MAG genome file (.fa or .fasta) [REQUIRED]
   -g      MAG prokka output (.gff)
   -o      Output files prefix (include path) [REQUIRED]
EOF
}

# variables
threads=
ref=
reads=
reads2=
outprefix=
gff=

while getopts "t:m:i:n:r:o:g:" OPTION
do
     case ${OPTION} in
         t)
             threads=${OPTARG}
             ;;
         i)
             reads=${OPTARG}
             ;;
         n)
             reads2=${OPTARG}
             ;;
         r)
             ref=${OPTARG}
             ;;
         o)
             outprefix=${OPTARG}
             ;;
         g)
             gff=${OPTARG}
             ;;
         ?)
             usage
             exit
             ;;
     esac
done

# check arguments
if [[ -z ${threads} ]] || [[ -z ${reads} ]] || [[ -z ${ref} ]] || [[ -z ${outprefix} ]] || [[ -z ${gff} ]]
then
     echo "ERROR : Please supply correct arguments"
     usage
     exit 1
fi


timestamp() {
  date +"%H:%M:%S"
}

readname=$(basename ${outprefix})
refix=$(basename ${ref%%.fa*})
if [ ${threads} -eq 1 ]
then
    threads_sam=1
else
    threads_sam=$((${threads}-1))
fi

# index reference file
if [[ ! -s ${outprefix}.1.bt2 ]]
then
    echo "$(timestamp) [ CMseq ] Indexing MAG FASTA file"
    bowtie2-build ${ref} ${outprefix}
else
    echo "$(timestamp) [ CMseq ] Index found, skipping bowtie2-build"
fi


# initial mapping and sorting
echo "$(timestamp) [ CMseq ] Mapping reads against MAG ..."
if [[ -z ${reads2} ]]
then
    bowtie2 --very-sensitive-local -x ${outprefix} -p ${threads} -U ${reads} | samtools view -@ ${threads_sam} -uS - | samtools sort -@ ${threads_sam} - -o ${outprefix}.bam
    samtools index -@ ${threads_sam} ${outprefix}.bam
else
    bowtie2 --very-sensitive-local -x ${outprefix} -p ${threads} -1 ${reads} -2 ${reads2} | samtools view -@ ${threads_sam} -uS - | samtools sort -@ ${threads_sam} - -o ${outprefix}.bam
    samtools index -@ ${threads_sam} ${outprefix}.bam
fi

# calculate polymorphic sites
echo "$(timestamp) [ CMseq ] Checking polymorphic sites ..."
scripts/polymut.py -c ${ref} ${outprefix}.bam --gff_file ${gff} --mincov 10 --minqual 30 --dominant_frq_thrsh 0.8 > ${outprefix}.cmseq.csv
rm -f ${outprefix}.bam ${outprefix}.bai

library(ggplot2)


args <- commandArgs(trailingOnly = TRUE)

checkm_sr=read.csv(args[1],header=TRUE,stringsAsFactor=FALSE)
checkm_coas=read.csv(args[2],header=TRUE,stringsAsFactor=FALSE)

checkm_sr$Status2="Single run"
checkm_coas$Status2="Co-assembly"

checkm=rbind.data.frame(checkm_sr,checkm_coas)
checkm$Status=factor(checkm$Status, levels=c("Single run", "Co-assembly"))

setwd("data/figures/")
ggplot(checkm, aes(x=Status, y=completeness, fill=Status)) +
  geom_boxplot(outlier.shape=NA)+theme_classic()+ylab("Completeness") + xlab("Assembly approach")+scale_fill_manual(breaks=c("Single run","Co-assembly"), values=c("#BBBBBB","#4477AA"))+guides(color=guide_legend(title="Approach"))+theme(legend.position="none")
ggsave("checkm_completeness.png",plot = last_plot())

ggplot(checkm, aes(x=Status, y=contamination, fill=Status))+geom_boxplot(outlier.shape=NA)+theme_classic()+ylab("Contamination") + xlab("Assembly approach")+scale_fill_manual(breaks=c("Single run","Co-assembly"), values=c("#BBBBBB","#4477AA"))+theme(legend.position="none")+ylim(0,10)
ggsave("checkm_contam.png",width=5,height=5)

library(ggplot2)

args <- commandArgs(trailingOnly = TRUE)

singlerun=read.table(args[1],header=FALSE,stringsAsFactor=FALSE)
coassembly=read.table(args[2],header=FALSE,fill=TRUE,stringsAsFactor=FALSE)
colnames(singlerun)=c("path","strain")
colnames(coassembly)=c("path","strain")

coassembly$status="Co-assembly"
singlerun$status="Single run"
coassembly$strain=as.numeric(as.vector(coassembly$strain))
singlerun$strain=as.numeric(as.vector(singlerun$strain))
coassembly=coassembly[complete.cases(coassembly$strain),]
singlerun=singlerun[complete.cases(singlerun$strain),]

all=rbind(coassembly,singlerun)
all$status=factor(all$status, levels=c("Single run", "Co-assembly"))
ggplot(all, aes(x=status, y=strain, fill=status)) +geom_boxplot(outlier.shape=NA)+theme_classic()+ylab("Strain Heterogeneity") + xlab("Approach")+scale_fill_manual(breaks=c("Single run","Co-assembly"), values=c("#BBBBBB","#4477AA"))+guides(color=guide_legend(title="Approach"))+ylim(0,12)+theme(legend.position="none")
ggsave("data/figures/cmseq_plot.png",width=5,height=5)

library(ggplot2)
library(gridExtra)

args <- commandArgs(trailingOnly = TRUE)

readcounts=read.table(args[1],header=TRUE,stringsAsFactor=FALSE)
df_flag=read.table(args[2],header=FALSE,stringsAsFactor=FALSE)
bwa.counts_sr=read.table(args[3],header=FALSE,stringsAsFactor=FALSE)
bwa.counts_coas=read.table(args[4],header=FALSE,stringsAsFactor=FALSE)
summary_out=args[5]

readcounts$Run=gsub("data/00_preprocessing/processed/singlerun/","",readcounts$Run)
readcounts=readcounts[!grepl("_2.fastq", readcounts$Run),]
readcounts$Run=gsub("_1.fastq","",readcounts$Run)
readcounts$Readcount=2*readcounts$Readcount

colnames(df_flag)=c("Run","Assembly")
df_comb=merge(df_flag,readcounts,by="Run")
df_comb$percassemb=df_comb$Assembly/df_comb$Readcount*100

colnames(bwa.counts_sr)=c("Run","Catalogue")
bwa.counts_sr$`Assembly Approach`="Single run"

colnames(bwa.counts_coas)=c("Run","Catalogue")
bwa.counts_coas$`Assembly Approach`="Coassembly"

bwa.counts=rbind.data.frame(bwa.counts_sr,bwa.counts_coas)
merged=merge(df_comb,bwa.counts, by="Run")
merged$percmags=merged$Catalogue/merged$Readcount*100
merged$percassemb=merged$Assembly/merged$Readcount*100

write.csv(merged,summary_out, quote=FALSE,row.names=FALSE)

ggplot(merged, aes(x=percmags, y=percassemb, color=`Assembly Approach`)) + geom_point() +xlab("Reads mapping to MAGs (%)") + ylab("Reads mapping to assembly (%)")+theme_classic()+scale_color_manual(breaks=c("Coassembly","Single run"), values=c("#a8ddb5","#c994c7"))+xlim(0,100)+ylim(0,100)
ggsave("data/figures/perassemb_perref.png",width=5,height=5)

import argparse
import sys

def ren_fasta(args):
	n = 0
	for line in open(args.fasta_file, "rU"):
		if line[0] == ">":
			name = line.strip("\n").replace(">","")
			n += 1
			print ">%s_%i\t%s" % (args.prefix, n, name)
		else:
			print line.strip("\n")


if __name__ == '__main__':
        parser = argparse.ArgumentParser(description='Rename multifasta file')
        parser.add_argument('-f', dest='fasta_file', help='Input FASTA file')
	parser.add_argument('-p', dest='prefix', help='Header prefix')
        if len(sys.argv) == 1:
                parser.print_help()
                sys.exit()
        else:
                args = parser.parse_args()
                ren_fasta(args)