from Bio import SeqIO
import os

min_size = snakemake.config["classify_contigs"]["min_size"]
max_size = snakemake.config["classify_contigs"]["max_size"]
assemblies = snakemake.input
output = snakemake.output["class_file"]
plasmids_files_path = snakemake.output["plasmids_list"]

# If only one sample, must converted into list
if len(assemblies) < 2:
    assemblies = [assemblies]

# Create folders
if not os.path.exists("Sequences/plasmids"):
    os.mkdir("Sequences/plasmids")
if not os.path.exists("Sequences/chromosomes"):
    os.mkdir("Sequences/chromosomes")

with open(output, "w") as report, open(plasmids_files_path, "w") as plasmids_list:
    # Write report first line
    report.write("Sample\tContig ID\tType\tLength\n")

    for file in assemblies:
        sample = file.split("/")[1] # sample ID (second term in path)
        assembly = SeqIO.parse(file, "fasta")
        for seq in assembly:
            # If plasmid
            if min_size <= len(seq.seq) <= max_size:
                with open(f"Sequences/plasmids/{sample}_plasmid_{seq.id}.fasta", "w") as out:
                    SeqIO.write([seq], out, "fasta")
                report.write(f"{sample}\t{seq.id}\tPlasmid\t{len(seq.seq)}\n")
                plasmids_list.write(f"Sequences/plasmids/{sample}_plasmid_{seq.id}.fasta\n")

            # If chromosome
            elif max_size < len(seq.seq):
                with open(f"Sequences/chromosomes/{sample}_chromosome_{seq.id}.fasta", "w") as out:
                    SeqIO.write([seq], out, "fasta")
                report.write(f"{sample}\t{seq.id}\tChromosome\t{len(seq.seq)}\n")

import pandas as pd

""" Gather results from Resfinder and mobsuite """

def add_resfinder(master_df, res_files):
    " Add resfinder results to the classification file "

    # Creation of dataframe for resfinder results
    res_df = pd.DataFrame({"Sample":[], "Contig ID":[], "ResFinder Resistance genes":[], "ResFinder phenotype":[]})

    # Parsing of all Resfinder files
    for f in res_files:
        sample = f.split("/")[1] # extracted from the path
        results = {} # stores the results per contig
        text = open(f, "r").readlines()[1:] # open and skip header
        for line in text:
            gene = line.split("\t")[0]
            contig = line.split("\t")[5]
            phenotype = line.split("\t")[7].split(", ") # list of phenotypes
            if contig not in results:
                results[contig] = [gene, phenotype] # list of genes in pos 0 and list of phenotype in pos 1
            else:
                results[contig][0] += f", {gene}" # add the new gene to the string in position 0 (list of genes)
                results[contig][1].extend(x for x in phenotype if x not in results[contig][1]) # add phenotype if not already present

        # Looping over results dict to complete dataframe. The label used is the length of the dataframe, creating a simple index.
        for elt in results:
            res_df.loc[len(res_df)] = [sample, elt, results[elt][0], ", ".join(results[elt][1])]

    # Merging the classification dataframe with the resfinder dataframe
    new_df = master_df.merge(res_df, how="outer", on=["Sample", "Contig ID"])

    return new_df

def add_mobsuite(master_df, mob_files):
    " Add mobsuite results to the master file "

    # Creation of dataframe for mobsuite results
    res_df = pd.DataFrame({"Sample":[], "Contig ID":[], "Mob-suite replicon types":[], "Mob-suite relaxase types":[]})

    # Parsing of all mobsuite files
    for f in mob_files:
        sample = f.split("/")[1].split("_")[0] # extracted from the path
        results = {} # stores the results per contig
        text = open(f, "r").readlines()[1:] # open and skip header
        for line in text:
            contig = line.split("\t")[0]
            replicon = line.split("\t")[5]
            relaxase = line.split("\t")[7]
            results[contig] = [replicon, relaxase]

        # Looping over results dict to complete dataframe. The label used is the length of the dataframe, creating a simple index.
        for elt in results:
            res_df.loc[len(res_df)] = [sample, elt, results[elt][0], results[elt][1]]

    # Merging the classification dataframe with the mobsuite dataframe
    new_df = master_df.merge(res_df, how="outer", on=["Sample", "Contig ID"])

    return new_df


if __name__ == "__main__":

    # Load dataframe and list of files from the Snakefile
    class_file = pd.read_csv(snakemake.input["classification"], sep='\t', header=0)
    class_file = class_file.astype({"Sample": str}) # Convert first column to string in case samples ID are just numbers
    res_files = snakemake.input["resfinder"]
    mob_files = snakemake.input["mobsuite"]

    # If only one sample, must converted into list
    if len(res_files) > 1:
        res_list = res_files
    else:
        res_list = [res_files[0]]

    if len(mob_files) > 1:
        mob_list = mob_files
    else:
        mob_list = [mob_files[0]]

    # Add information for each analysis
    infos_resfinder = add_resfinder(class_file, res_list)
    infos_mobsuite = add_mobsuite(infos_resfinder, mob_list)

    # Export the final dataframe
    final_df = infos_mobsuite.sort_values(by=["Sample", "Type", 'Length'])
    final_df.to_csv(snakemake.output[0], sep="\t")

shell:
    "NanoFilt -q {params.quality} -l {params.length} {input} --logfile {log}> {output}"

shell:
    "medaka_consensus -i {input.reads} -d {input.assembly} -t {threads} -m {params.model} -o 04-Medaka/{wildcards.sample} > {log} 2>&1"

shell:
    "flye --nano-hq {input} -t {threads} -o 03-assembly/{wildcards.sample}_long {params.others} > {log} 2> /dev/null"

shell:
    "fastp -i {input.R1} -I {input.R2} -o {output.R1} -O {output.R2} -q {params.qual} "
    "-u {params.unqual} -e {params.ave_qual} -l {params.length} -h {output.html} -j {output.json} 2> {log}"

shell:
    "mkdir QC/{wildcards.sample};"
    "fastqc {input} -o QC > /dev/null 2>&1"

shell:
    "unicycler -1 {input.short_1} -2 {input.short_2} -l {input.long} -o 03-assembly/{wildcards.sample}_hybrid -t {threads} --mode {params.mode} "
    "--verbosity 0 --keep 0"

shell:
    "bwa index {input.draft} > /dev/null 2>&1;"
    "bwa mem -t {threads} -a {input.draft} {input.short_1} > {output.alignments_1} 2> {log};"
    "bwa mem -t {threads} -a {input.draft} {input.short_1} > {output.alignments_2} 2> {log}"