#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import sys
from pathlib import Path
from snakemake.logging import logger
import logging
from pprint import pprint as pp
# load own functions
from culebrONT import CulebrONT, logo
from culebrONT.global_variables import (VARIANT_SUFFIX, EXTERNAL, SELECTION_FILE_NAME, SELECTION_COLUMNS,
                                        TAXDUMP_FILES, GENOMESCOPE_READS)

# GLOBAL VARIABLES
culebront = CulebrONT(workflow=workflow, config=config)
tools_config = culebront.tools_config


#print(culebront.string_to_dag)
# pp(culebront.export_use_yaml)

# print for debug:
# logger.debug(pp(culebront))
# exit()

# Getting paths on usefully variables
output_dir = config['DATA']['OUTPUT']
fastq_dir = config['DATA']['FASTQ']
reference_file = config['DATA']['REF']

# tools versions are collected by the rules write_versions_script and tools_versions (see below)


################ WILDCARDS  ################
FASTQ, = glob_wildcards(f"{config['DATA']['FASTQ']}{{fastq}}{culebront.fastq_files_ext}") if config['DATA']['FASTQ'] else ([],)
nb_racon_rounds = culebront.nb_racon_rounds
nb_pilon_rounds = culebront.nb_pilon_rounds

# load some variables from culebront class
add_circular_name = culebront.add_circular_name

#output of Minipolish, -> (input de Tag Circular Minipolish)
TMP = culebront.TMP
#output of Tag Circular Minipolish , -> (input de fixstart)
TCM = culebront.TCM
#output of the assemblers, -> (input de tag_circular)
TAG = culebront.TAG

#######################################
# Change workdir to output path
workdir:config["DATA"]["OUTPUT"]
#######################################

# {assemblers} is a variant name (FLYE, FLYE_1) or an external assembly (EXTERNAL_<name> or <name>);
# each assembler rule narrows it
wildcard_constraints:
    assemblers = "[A-Za-z0-9_]+",
    fastq = "[^/]+",

tool_of = culebront.variant_tool

def get_fastq(wildcards):
    return f"{fastq_dir}{wildcards.fastq}{culebront.fastq_files_ext}"


def draft_to_racon(wildcards):
    n = int(wildcards.nb)
    if n == 1:
        if config['CIRCULAR']:
            return f"{output_dir}{{fastq}}/ASSEMBLERS/{{assemblers}}/POLISHING/ROTATE/rotate_{n}/assembly.racon{n}.fasta"
        else:
            return f"{output_dir}{{fastq}}/ASSEMBLERS/{{assemblers}}/ASSEMBLER/assembly{add_circular_name}.fasta"
    elif n > 1:
        if config['CIRCULAR']:
            return f"{output_dir}{{fastq}}/ASSEMBLERS/{{assemblers}}/POLISHING/ROTATE/rotate_{n}/assembly.racon{n}.fasta"
        else:
            return f"{output_dir}{{fastq}}/ASSEMBLERS/{{assemblers}}/POLISHING/RACON/racon_{n-1}/assembly.racon{n-1}.fasta"
    else:
        raise ValueError(f"loop numbers must be 1 or greater: received {n}")

def draft_to_rotate(wildcards):
    n = int(wildcards.nb)
    if n == 1:
        return  f"{output_dir}{{fastq}}/ASSEMBLERS/{{assemblers}}/ASSEMBLER/assemblyCIRCULARISED_circTag.fasta"
    elif n > 1:
        return f"{output_dir}{{fastq}}/ASSEMBLERS/{{assemblers}}/POLISHING/RACON/racon_{n-1}/assembly.racon{n-1}.fasta"
    else:
        raise ValueError(f"loop numbers must be 1 or greater: received {n}")

def draft_to_correction(wildcards):
    if 'MINIASM' == wildcards.assemblers:
        return f"{output_dir}{{fastq}}/ASSEMBLERS/{{assemblers}}/POLISHING/RACON/racon_{nb_racon_rounds}/assembly.racon{nb_racon_rounds}{TCM}.fasta"
    elif config['POLISHING']['RACON']:
        return f"{output_dir}{{fastq}}/ASSEMBLERS/{{assemblers}}/POLISHING/RACON/racon_{nb_racon_rounds}/assembly.racon{nb_racon_rounds}.fasta"
    else:
        return f"{output_dir}{{fastq}}/ASSEMBLERS/{{assemblers}}/ASSEMBLER/assembly{add_circular_name}.fasta"


def draft_to_pilon(wildcards):
    n = int(wildcards.nbp)
    if n == 1:
        return f"{output_dir}{{fastq}}/ASSEMBLERS/{{assemblers}}/CORRECTION/PILON/pilon_1/assembly.pilon1.fasta"
    elif n > 1:
        return f"{output_dir}{{fastq}}/ASSEMBLERS/{{assemblers}}/CORRECTION/PILON/pilon_{n-1}/assembly.pilon{n-1}.fasta"
    else:
        raise ValueError(f"loop numbers for pilon must be 1 or greater: received {n}")

def get_illumina(wildcards):
    r1 = ""
    r2 = ""
    for element1 in culebront.R1 :
        if wildcards.fastq in element1:
            r1=element1
    for element2 in culebront.R2 :
        if wildcards.fastq in element2:
            r2=element2
    return f"{r1} {r2}"

################################ FINAL ####################################
rule final:
    input:
        # f"{output_dir}FINAL_REPORT/CulebrONT_report.html",
        f"{output_dir}REPORT/Book/index.html"
        # f"{output_dir}REPORT/QMD/about_workflow.qmd"


################################ SELECTION ####################################
def selected_files(pattern, fastq, with_reads=False):
    """files of a rule for the assemblies of the Sélection d'assemblages (evaluate=yes) of a sample"""
    return [expand(pattern, fastq=row["fastq"], assemblers=row["assemblers"], quality_step=row["quality_step"])[0]
            for row in culebront.selected_rows(fastq, with_reads)]


################################ ASSEMBLY ####################################
include: f"{culebront.install_path}/snakefiles/assemblers.snake"

############################### CIRCULARISATION ##############################
include: f"{culebront.install_path}/snakefiles/circular.snake"

################################ POLISHING ####################################
include: f"{culebront.install_path}/snakefiles/polishers.snake"

################################ CORRECTION ####################################
include: f"{culebront.install_path}/snakefiles/correction.snake"

################################ QUALITY ####################################
include: f"{culebront.install_path}/snakefiles/quality.snake"

################################ RAPPORT ####################################
include: f"{culebront.install_path}/snakefiles/reported.snake"


rule copy_final_assemblies:
    """ build stats on assembly"""
    threads: 1
    input:
        fasta_list = rules.preparing_fasta_to_quality.output.renamed
    output:
        fasta_final = f"{output_dir}ASSEMBLIES/{{fastq}}-{{assemblers}}-{{quality_step}}{add_circular_name}.fasta"
    params:
        sample = '{fastq}'
    log:
        output = f"{output_dir}ASSEMBLIES/LOGS/{{fastq}}_{{assemblers}}_{{quality_step}}_CP-FINAL-ASSEMBLIES.o",
        error = f"{output_dir}ASSEMBLIES/LOGS/{{fastq}}_{{assemblers}}_{{quality_step}}_CP-FINAL-ASSEMBLIES.e"
    message:
        """
        Launching {rule}
        threads : {threads}
        input:
             fasta_list : {input.fasta_list}
        output:
             fasta_final : {output.fasta_final}
        log:
            output : {log.output}
            error: {log.error}
        """
    shell:
        """
        cp {input.fasta_list} {output.fasta_final}
        """


def input_report_by_sample(wildcards):
    """results of the activated tools for the assemblies of the Sélection d'assemblages of a sample (or of the
    external group), one file per assembly, keys read by report_helpers.py"""
    fastq = wildcards.fastq
    reads = fastq != EXTERNAL
    dico_final = {"fasta": selected_files(rules.preparing_fasta_to_quality.output.renamed, fastq)}
    if reads:
        dico_final["time"] = expand(rules.run_benchmark_time.output.stat_time, fastq=fastq)
    if config['QUALITY']['QUAST']:
        dico_final["quast"] = expand(rules.run_quast.output.report, fastq=fastq)
    if config['QUALITY']['BUSCO']:
        dico_final["busco"] = selected_files(rules.run_busco.output.summary, fastq)
    if config['QUALITY']['BLOBTOOLS']:
        dico_final["blob_table"] = selected_files(rules.run_blobtools.output.table, fastq)
        dico_final["blob_snail"] = selected_files(rules.run_blobtools.output.snail, fastq)
        dico_final["blob_cumulative"] = selected_files(rules.run_blobtools.output.cumulative, fastq)
        # blob plot only with the ONT coverage (assemblies attached to a sample)
        dico_final["blob_plot"] = selected_files(rules.run_blobtools_blob_plot.output.blob, fastq, with_reads=True)
    if config['QUALITY']['GENOMESCOPE'] and reads:
        dico_final["genomescope_summary"] = expand(rules.run_genomescope.output.summary, fastq=fastq)
        dico_final["genomescope_linear"] = expand(rules.run_genomescope.output.linear_plot, fastq=fastq)
        dico_final["genomescope_log"] = expand(rules.run_genomescope.output.log_plot, fastq=fastq)
    if config['QUALITY']['GCI'] and reads:
        dico_final["gci"] = selected_files(rules.run_gci.output.gci, fastq, with_reads=True)
    if config['QUALITY']['LTR_RETRIEVER']:
        dico_final["lai"] = selected_files(rules.run_ltr_retriever.output.lai, fastq)
    if config['QUALITY']['ASSEMBLYTICS']:
        dico_final["assemblytics_summary"] = selected_files(rules.run_assemblytics.output.summary, fastq)
        dico_final["assemblytics_dotplot"] = selected_files(rules.run_assemblytics.output.png_dotplot, fastq)
        dico_final["assemblytics_nchart"] = selected_files(rules.run_assemblytics.output.png_Nchart, fastq)
    if config['QUALITY']['MERQURY'] and reads:
        dico_final["merqury_qv"] = selected_files(rules.run_merqury.output.qv_global, fastq, with_reads=True)
        dico_final["merqury_completeness"] = selected_files(rules.run_merqury.output.comp_stats, fastq, with_reads=True)
    if config['QUALITY']['FLAGSTATS'] and reads:
        dico_final["flagstat"] = selected_files(rules.run_flagstat.output.txt, fastq, with_reads=True)
    return dico_final


def report_assemblies(wildcards):
    """(assembler, quality step, fasta) of the assemblies of the Sélection d'assemblages of a sample"""
    return [(row["assemblers"], row["quality_step"],
             selected_files(rules.preparing_fasta_to_quality.output.renamed, wildcards.fastq)[i])
            for i, row in enumerate(culebront.selected_rows(wildcards.fastq))]


rule report_by_sample :
    """ create a quarto report by sample/fastq"""
    threads: 1
    input:
        unpack(input_report_by_sample)
    output:
        qmd_by_sample = f"{output_dir}REPORT/QMD/{{fastq}}/{{fastq}}.ipynb"
    params:
        sample_name = '{fastq}',
        assemblies = report_assemblies,
        skipped = culebront.quality_skipped,
    log:
        output=f"{output_dir}REPORT/LOGS/{{fastq}}_QMD.o",
        error=f"{output_dir}REPORT/LOGS/{{fastq}}_QMD.e"
    message:
        """
        Launching {rule}
        threads : {threads}
        input:
             input_report_by_sample : {input}
        output:
             qmd_by_sample : {output.qmd_by_sample}
        log:
            output : {log.output}
            error: {log.error}
        """
    script:
        f"{culebront.snakemake_scripts}/generate_samples.py"


localrules: write_versions_script, download_busco_lineage, write_assembly_selection, pre_evaluation

rule write_versions_script:
    """write the bash script that prints the version of each activated tool"""
    output:
        script = f"{output_dir}LOGS/tools_versions.sh"
    params:
        content = culebront.versions_script()
    run:
        Path(output.script).write_text(params.content)


rule medaka_version:
    """version of medaka, which has its own image"""
    threads: 1
    output:
        csv = f"{output_dir}LOGS/versions_medaka.csv"
    container:
        tools_config['APPTAINER'].get('MEDAKA', '')
    envmodules:
        tools_config["ENV-MODULES"]["MEDAKA"]
    shell:
        """
        printf 'MEDAKA,None,"%s",@medaka\\n' "$(medaka --version 2>/dev/null | tr '\\n' ' ' | sed 's/ *$//')" > {output.csv}
        """


rule ltr_retriever_version:
    """version of LTR_retriever, which has its own image"""
    threads: 1
    output:
        csv = f"{output_dir}LOGS/versions_ltr_retriever.csv"
    container:
        tools_config['APPTAINER'].get('LTR_RETRIEVER', '')
    envmodules:
        tools_config["ENV-MODULES"].get("LTR_RETRIEVER", "")
    shell:
        """
        printf 'LTR_RETRIEVER,None,"%s",@ltr_retriever\\n' "$(LTR_retriever -h 2>&1 | grep -o 'v[0-9][0-9.]*' | head -1)" > {output.csv}
        """


rule tools_versions:
    """collect the version of each activated tool, in the container or with the modules of the active Profile"""
    threads: 1
    input:
        script = rules.write_versions_script.output.script,
        medaka = rules.medaka_version.output.csv if "MEDAKA" in culebront.correction_tools_activated else [],
        ltr_retriever = rules.ltr_retriever_version.output.csv if "LTR_RETRIEVER" in culebront.quality_tools_activated else [],
    output:
        csv = f"{output_dir}versions.csv"
    log:
        error = f"{output_dir}LOGS/tools_versions.e"
    container:
        tools_config['APPTAINER']['TOOLS']
    envmodules:
        *culebront.versions_modules()
    shell:
        """
        (bash {input.script}; cat {input.medaka} {input.ltr_retriever} /dev/null) > {output.csv} 2>{log.error}
        """


rule report_about_workflow:
    """build qmd for tools version dag and config"""
    threads: 1
    input:
        dag = rules.rule_graph.output.dag,
        report_snakemake = rules.run_report_snakemake.output.report_snakemake,
        versions = rules.tools_versions.output.csv
    output:
        about_ipynb = f"{output_dir}REPORT/QMD/about_workflow.ipynb"
    params:
        config_yaml = culebront.export_use_yaml,
        versions = f"{output_dir}versions.csv",
    log:
        output=f"{output_dir}REPORT/LOGS/report_about_workflow_QMD.o",
        error=f"{output_dir}REPORT/LOGS/report_about_workflow_QMD.e"
    message:
        """
        Launching {rule}
        threads : {threads}
        input:
             dag : {input.dag}
             versions : {params.versions}
        output:
             about_ipynb : {output.about_ipynb}
        log:
            output : {log.output}
            error: {log.error}
        """
    script:
        f"{culebront.snakemake_scripts}/generate_tools.py"

rule ipynb_convert_samples_qmd:
    """build qmd for tools version dag and config"""
    threads: 1
    input:
        ipynb = rules.report_by_sample.output.qmd_by_sample,
    output:
        sample_qmd = f"{output_dir}REPORT/QMD/{{fastq}}/{{fastq}}.qmd"
    log:
        output=f"{output_dir}REPORT/LOGS/ipynb_convert_{{fastq}}_qmd.o",
        error=f"{output_dir}REPORT/LOGS/ipynb_convert_{{fastq}}_qmd.e"
    message:
        """
        Launching {rule}
        threads : {threads}
        input:
             ipynb : {input.ipynb}
        output:
             sample_qmd : {output.sample_qmd}
        log:
            output : {log.output}
            error: {log.error}
        """
    container:
        tools_config['APPTAINER']['TOOLS']
    envmodules:
        tools_config["ENV-MODULES"]["QUARTO"]
    shell:
        "quarto convert {input.ipynb} 1>{log.output} 2>{log.error}"


rule ipynb_convert_qmd:
    """build qmd for tools version dag and config"""
    threads: 1
    input:
        ipynb = rules.report_about_workflow.output.about_ipynb,
    output:
        about_qmd = f"{output_dir}REPORT/QMD/about_workflow.qmd"
    log:
        output=f"{output_dir}REPORT/LOGS/ipynb_convert_qmd.o",
        error=f"{output_dir}REPORT/LOGS/ipynb_convert_qmd.e"
    message:
        """
        Launching {rule}
        threads : {threads}
        input:
             ipynb : {input.ipynb}
        output:
             about_qmd : {output.about_qmd}
        log:
            output : {log.output}
            error: {log.error}
        """
    container:
        tools_config['APPTAINER']['TOOLS']
    envmodules:
        tools_config["ENV-MODULES"]["QUARTO"]
    shell:
        "quarto convert {input.ipynb} 1>{log.output} 2>{log.error}"

rule edit_quarto:
    """edit quarto config file to add page of samples"""
    threads: 1
    input:
        list_sample_qmd = lambda wildcards: expand(rules.ipynb_convert_samples_qmd.output.sample_qmd, fastq=culebront.report_samples),
        quarto = f"{culebront.snakemake_scripts}/report_template/_quarto.yml",
    output:
        quarto_conf = f"{output_dir}REPORT/QMD/_quarto.yml"
    params:
        out_dir = directory(f"{output_dir}REPORT/Book/")
    log:
        output=f"{output_dir}REPORT/LOGS/edit_quarto.o",
        error=f"{output_dir}REPORT/LOGS/edit_quarto.e"
    message:
        """
        Launching {rule}
        threads : {threads}
        input:
             list_sample_qmd : {input.list_sample_qmd}
             quarto_conf : {input.quarto}
        output:
             quarto_conf : {output.quarto_conf}
        log:
            output : {log.output}
            error: {log.error}
        """
    script:
        f"{culebront.snakemake_scripts}/edit_quarto.py"


rule build_book:
    """copy template files to build report"""
    threads: 1
    input:
        qmd = rules.ipynb_convert_qmd.output.about_qmd,
        quarto_conf = rules.edit_quarto.output.quarto_conf,
        index_qmd = f"{culebront.snakemake_scripts}/report_template/index.qmd",
        list_sample_qmd = lambda wildcards: expand(rules.ipynb_convert_samples_qmd.output.sample_qmd, fastq=culebront.report_samples),
        ref_bib = f"{culebront.snakemake_scripts}/report_template/references.bib",
        ref_qmd = f"{culebront.snakemake_scripts}/report_template/references.qmd",
        csl = f"{culebront.snakemake_scripts}/report_template/cellalina.csl",
        scss = f"{culebront.snakemake_scripts}/report_template/offline.scss",
        helpers = f"{culebront.snakemake_scripts}/report_template/report_helpers.py",
        extensions = f"{culebront.snakemake_scripts}/report_template/_extensions",
        logo = f"{culebront.install_path}/culebront_logo.png",
    output:
        html = f"{output_dir}REPORT/Book/index.html"
    params:
        book_dir_qmd = directory(f"{output_dir}REPORT/QMD/"),
        input_dir = f"{culebront.snakemake_scripts}/",
        link_html = f"{output_dir}report.html"
    log:
        output=f"{output_dir}REPORT/LOGS/build_book.o",
        error=f"{output_dir}REPORT/LOGS/build_book.e"
    message:
        """
        Launching {rule}
        threads : {threads}
        input:
             qmd : {input.qmd}
        output:
             about_qmd : {output.html}
        log:
            output : {log.output}
            error: {log.error}
        """
    container:
        tools_config['APPTAINER']['TOOLS']
    envmodules:
        tools_config["ENV-MODULES"]["QUARTO"]
    shell:
        """
            (mkdir -p {params.book_dir_qmd};
            cp {input.index_qmd} {params.book_dir_qmd};
            cp {input.ref_bib} {params.book_dir_qmd};
            cp {input.csl} {params.book_dir_qmd};
            cp {input.logo} {params.book_dir_qmd};
            cp {input.ref_qmd} {input.scss} {input.helpers} {params.book_dir_qmd};
            cp -r {input.extensions} {params.book_dir_qmd};
            # plotly.js of the report environment, loaded once by the sample pages (no CDN, see generate_samples.py)
            cp $(python3 -c "import os, plotly; print(os.path.join(os.path.dirname(plotly.__file__), 'package_data', 'plotly.min.js'))") {params.book_dir_qmd};
            cd {params.book_dir_qmd};
            quarto render --to html
            # redirection, not a link: through a link, the browser resolves the paths of the book from OUTPUT/
            rm -f {params.link_html}
            printf '<!doctype html><meta charset="utf-8"><meta http-equiv="refresh" content="0; url=REPORT/Book/index.html"><a href="REPORT/Book/index.html">CulebrONT report</a>\\n' > {params.link_html}
            ) 1>{log.output} 2>{log.error}
        """


################################ PRE-EVALUATION ####################################
# after copy_final_assemblies, used by the Sélection d'assemblages
include: f"{culebront.install_path}/snakefiles/pre_evaluation.snake"
