

rule kraken_to_krona:
    input: __kraken__output_wkdir + "kraken.out.summary"
    output: __kraken__output_html
    run:
        shell("ktImportText %s -o %s" % (input, output))


rule kraken_translate:
    input: __kraken__output
    output:
        summary = __kraken__output_wkdir + "kraken.out.summary",
        png = __kraken__output_wkdir + "kraken.png",
    run:
        from sequana.kraken import KrakenResults
        k = KrakenResults(input[0])
        k.kraken_to_krona(output_filename=output.summary)
        df = k.plot(kind='pie')
        if df is not None and len(df):
            from pylab import savefig
            savefig(output.png)
        else:
            from sequana import sequana_data
            import shutil
            filename = sequana_data("no_data.jpg", "images")
            shutil.copy(filename, output.png)


rule kraken:
    """Kraken taxonomic sequence classification system

    Required input:
        - __kraken__input

    Required output:
        - __kraken__output_wkdir: working directory
        - __kraken__output: the kraken final output
        - __kraken__output_csv: summary in csv format
        - __kraken__output_json: summary in json format

    Config:
        .. code-block:: yaml

            kraken:
                database_directory:  # a valid path to a Kraken database

    See `KrakenBuilder <http://sequana.readthedocs.io/en/master/references.html#sequana.kraken_builder.KrakenBuilder>`_ to build your own database or visit 
    https://github.com/sequana/data for a database toy example.

    References:
        - https://github.com/DerrickWood/kraken

   """
    input: __kraken__input
    output:
        unclassified = __kraken__output_wkdir + "kraken_unclassified.fastq.gz",
        kraken = __kraken__output,
        kraken_csv = __kraken__output_csv,
        kraken_json = __kraken__output_json
    params:
        database = config['kraken']['database_directory']
    threads: config['kraken']['threads']
    run:

        # Note the unclassified[0:-3] to ignore the extension .gz
        outname = output.unclassified[0:-3]
        if len(input) == 1:
            shell("kraken -db {params.database} {input[0]} "
                " --threads {threads}"
                " --out {output.kraken} "
                " --unclassified-out %s " % outname
                )
        elif len(input) == 2:
            shell("kraken -db {params.database}  {input[0]} {input[1]} "
                " --threads {threads} "
                " --paired  --out {output.kraken}"
                " --unclassified-out %s " % outname
                )
        # Now we compress the unclassified file
        shell("pigz -p 4 " + outname)

        from sequana.kraken import KrakenResults
        k = KrakenResults(output.kraken)
        k.kraken_to_csv(output.kraken_csv, params.database)
        k.kraken_to_json(output.kraken_json, params.database)


