import yaml 
import csv
import pandas as pd
import os
from Bio import SeqIO

cwd = os.getcwd()
shell.executable('bash')

##### Configuration #####
configfile: workflow.current_basedir + "/config.yaml"

run_name = ""
if config.get("run_name"):
    run_name = config["run_name"]
else:
    run_name = "nCoV2019"


config["basecalled_path"] = config["basecalled_path"].rstrip("/")
config["run_path"] = config["basecalled_path"].replace("fastq_pass","").rstrip("/")
if config.get("fast5_path"):
    config["fast5_path"] = config["fast5_path"].rstrip("/")
else:
    print("No fast5 path given, assuming fast5_pass in same directory as fastq_pass")
    fast5_path = config["basecalled_path"].replace("fastq_pass","").rstrip('/') + "/fast5_pass"
    config["fast5_path"] = fast5_path

if config.get("annotated_path"):
    config["annotated_path"] = config["annotated_path"].rstrip("/")
else:
    config["annotated_path"] = "annotations"


barcode_string = ''
barcodes = []
barcode_to_sample = {}

if config.get("samples"):
    handle = config["samples"]
    samples = yaml.safe_load(handle)

    for s in samples:
        if s.lower()=="negative":
            pass
        else:
            for barcode in samples[s]:
                barcode_string += barcode + " "
                barcodes.append(barcode)
                barcode_to_sample[barcode]=s
    barcode_string = barcode_string.rstrip(" ")
    config["barcodes"]= barcodes
    print("Found samples in configuration.\n")
    print("barcode\tsample")
    for barcode in barcode_to_sample:
        print(barcode,'\t',barcode_to_sample[barcode])

elif config.get("barcodes"):
    barcode_config = config["barcodes"]
    barcodes = barcode_config.split(",")
    config["barcodes"] = barcodes

    for barcode in barcodes:
        barcode_to_sample[barcode] = barcode
        barcode_string+=barcode + " "

    print("Found barcodes in configuration.\n")
    print("barcode\tsample")
    for barcode in barcode_to_sample:
        print(barcode,'\t',barcode_to_sample[barcode])
else:
    barcodes = ["none"]
    barcode_string = "none"
    barcode_to_sample["none"] = "nCoV2019"

    print("Found no barcode or sample information in configuration.\n")
    print("barcode\tsample")
    for barcode in barcode_to_sample:
        print(barcode,'\t',barcode_to_sample[barcode])

##### Workflow #####

rule all:
    input:
        expand("binned_{barcode}.fastq",barcode=barcodes),
        f"{run_name}_sequencing_summary.txt",
        expand("{barcode}.consensus.fasta",barcode=barcodes),
        f"consensus_sequences/{run_name}.fasta"

rule barcoding:
    input:
    params:
        path_to_reads = config["basecalled_path"],
        barcode_str = "{barcode}",
        prefix = "./binned"
    output:
        "binned_{barcode}.fastq"
    shell:
        "artic guppyplex "
        "--skip-quality-check "
        "--min-length 400 "
        "--max-length 700 "
        "--directory {params.path_to_reads:q}/{params.barcode_str} "
        "--prefix {params.prefix:q}"


rule get_sequencing_summary_file:
    input:

    params:
        path_to_run = config["run_path"]
    output:
        f"{run_name}_sequencing_summary.txt"
    run:
        summary_files = []
        for r,d,f in os.walk(str(params.path_to_run)):
            for fn in f:
                if fn.startswith("sequencing_summary"):
                    summary_files.append(r + '/' + fn)
        print(f"Number of sequencing summary files found: {len(summary_files)}")

        summaryfh = open(str(output[0]), "w")
        dfs = []
        for summaryfn in summary_files:
            df = pd.read_csv(summaryfn, sep="\t")
            # support for local basecalling
            if 'filename_fast5' in df.columns:
                df['filename'] = df['filename_fast5']
            dfs.append(df)

        pd.concat(dfs).to_csv(summaryfh, sep="\t", index=False)
        summaryfh.close()

rule consensus_get:
    input:
        summary = rules.get_sequencing_summary_file.output,
        reads = "binned_{barcode}.fastq"

    params:
        primer_scheme_dir = workflow.current_basedir + "/../../../primer_schemes",
        primer_version = "nCoV-2019/V3",
        path_to_fast5 = config["fast5_path"],
        sample_name = "{barcode}"


    output:
        "{barcode}.trimmed.rg.sorted.bam",
        "{barcode}.merged.vcf",
        "{barcode}.consensus.fasta"

    shell:
        "set +o pipefail;"
        "artic minion "
        "--normalise 200 "
        "--threads 1 "
        "--scheme-directory {params.primer_scheme_dir} "
        "--read-file {input.reads:q} "
        "--fast5-directory {params.path_to_fast5:q}/ "
        "--sequencing-summary {input.summary} "
        "{params.primer_version} "
        "{params.sample_name}"

rule tidy_up:
    input:
        expand("{barcode}.consensus.fasta", barcode = config["barcodes"]),
        expand("{barcode}.trimmed.rg.sorted.bam", barcode = config["barcodes"])
    params:
        barcode_to_sample = barcode_to_sample
    output:
        f"consensus_sequences/{run_name}.fasta"
    run:
        if not os.path.exists("./CLIMB"):
            shell("mkdir ./CLIMB")
        for barcode in barcodes:
            sample = params.barcode_to_sample[barcode]
            if not os.path.exists("./CLIMB/"+sample):
                shell("mkdir ./CLIMB/" + sample)

        with open(str(output[0]),"w") as fw:
            for r,d,f in os.walk(cwd):
                for fn in f:
                    if fn.endswith(".consensus.fasta"):
                        barcode = fn.rstrip(".consensus.fasta")

                        if barcode in params.barcode_to_sample.keys():

                            for record in SeqIO.parse(r + '/' + fn,"fasta"):
                                barcode = fn.split(".")[0]
                                sample = params.barcode_to_sample[barcode]
                                fw2 = open(cwd + "/CLIMB/"+sample+ "/" + sample + ".fasta","w")

                                shell("cp " + r + '/' + barcode + ".trimmed.rg.sorted.bam" + " ./CLIMB/"+sample+ "/" + sample + ".bam")
                                fw.write(f">{sample} barcode={barcode} length={len(record.seq)}\n{record.seq}\n")
                                fw2.write(f">{sample} barcode={barcode} length={len(record.seq)}\n{record.seq}\n")
                                fw2.close()
                            else:
                                print("Barcode not found", barcode)
