#!/usr/bin/env python

""" MultiQC module to parse output from sequana"""
import os
import re
import json

from collections import OrderedDict

# prevent boring warning (version 1.0)
import logging

logging.captureWarnings(True)
from multiqc import config
from multiqc.modules.base_module import BaseMultiqcModule
from multiqc.plots import linegraph, table, heatmap, bargraph

logging.captureWarnings(False)

import colorlog

log = colorlog.getLogger(__name__)


class MultiqcModule(BaseMultiqcModule):
    """

    Here keys are the sample name

    {
     "25": {
      "Bacteria": 0.4,
      "Metazoa": 2.0,
      "Unclassified": 0.0,
      "Viruses": 97.6
     },
     "26": {
      "Bacteria": 8.800000000000002,
      "Metazoa": 0.4,
      "Unclassified": 18.4,
      "Viruses": 72.39999999999998
     },
    }

    .. warning:: this is not really a multiqc here because we read a file
        that summarizes several samples together.

    """

    def __init__(self):
        # Initialise the parent object
        super(MultiqcModule, self).__init__(
            name="Sequana/Kraken",  # name that appears at the top
            anchor="sequana_kraken",  #
            target="sequana",  # Name show that link to the following href
            href="http://github.com/sequana/sequana/",
            info="sequana_kraken pipeline multi summary",
        )

        self.sequana_data = {}

        # In theory only one file
        for myfile in self.find_log_files("sequana_kraken"):
            name = myfile["s_name"]
            d = json.loads(myfile["f"])

            for k, v in d.items():
                S = sum(list(v.values()))
                U = v["Unclassified"]
                v["Classified"] = S - U
                self.sequana_data[k] = v

        if len(self.sequana_data) == 0:
            log.debug("No samples found: sequana_kraken")
            raise UserWarning

        """info = "<ul>"
        for this in sorted(self.sequana_data.keys()):
            info += '<li><a href="{}/summary.html">{}</a></li>'.format(this,this,this)
        info += "</ul>"
        href="http://sequana.readthedocs.io/en/main/"
        target = "Sequana"
        mname = '<a href="{}" target="_blank">{}</a> individual report pages:'.format(href, target)
        self.intro = '<p>{} {}</p>'.format( mname, info)
        """

        log.info("Found {} reports".format(len(self.sequana_data)))

        self.populate_columns()
        self.add_kraken()
        self.add_classified_vs_unclassified()

    def add_classified_vs_unclassified(self):
        data = {}
        for sample_name in self.sequana_data.keys():
            C = self.sequana_data[sample_name]["Classified"]
            U = self.sequana_data[sample_name]["Unclassified"]
            data[sample_name] = {"unclassified": U, "classified": C}

        pconfig = {
            "title": "classification",
            "cpswitch": False,
            "min": 0,
            "max": 100,
            "format": "{0:.2f}",
            "logswitch": False,
        }

        keys = OrderedDict()
        keys["classified"] = {"color": "green", "name": "Classified"}
        keys["unclassified"] = {"color": "red", "name": "Unclassified"}

        self.add_section(
            name="Classification ratio",
            anchor="classification_ratio",
            description="The following barplots shows proportion of classified and unclassified reads in percentage",
            helptext="",
            plot=bargraph.plot(data, keys, pconfig),
        )

    def _set_nan_to_zero(self, x):
        try:
            x + 0
            return x
        except:
            return 0

    def add_kraken(self):
        data = {}

        # First, we figure out all possible names
        kingdoms = set([x for k in self.sequana_data.keys() for x in self.sequana_data[k].keys()])

        colors = ["Archaea", "Bacteria", "Eukaryota", "Viruses", "Metazoa", "Fungi", "Unclassified", "Classified"]

        for sample_name in self.sequana_data.keys():
            for kingdom in sorted(kingdoms):
                if kingdom not in self.sequana_data[sample_name]:
                    self.sequana_data[sample_name][kingdom] = 0

            data[sample_name] = {"others": 0}
            for kingdom in sorted(kingdoms):
                if kingdom not in colors:
                    # here we add together non-superkingdom + other artifical
                    # sequences
                    data[sample_name]["others"] += self._set_nan_to_zero(self.sequana_data[sample_name][kingdom])
                else:
                    data[sample_name][kingdom.lower()] = self._set_nan_to_zero(self.sequana_data[sample_name][kingdom])
            data[sample_name]["unclassified"] = self._set_nan_to_zero(self.sequana_data[sample_name]["Unclassified"])

        pconfig = {
            "title": "Taxonomy by kingdom",
            # "percentages": True,
            "cpswitch": False,
            "min": 0,
            "max": 100,
            "format": "{0:.2f}",
            "logswitch": False,
        }

        keys = OrderedDict()
        # superkingdom:
        keys["archea"] = {"color": "orange", "name": "Archea"}
        keys["bacteria"] = {"color": "#b1084c", "name": "Bacteria"}
        keys["eukaryota"] = {"color": "green", "name": "Eukaryota"}
        keys["viruses"] = {"color": "#437bb1", "name": "Viruses"}
        # kingdom:
        keys["metazoa"] = {"color": "green", "name": "Metazoa"}
        keys["fungi"] = {"color": "purple", "name": "Fungi"}
        # others
        keys["unclassified"] = {"color": "grey", "name": "Unclassified"}
        keys["others"] = {"color": "blue", "name": "Others"}
        # subkingdom
        # keys['viridiplantae'] = {'color': 'yellow', 'name': 'Viridiplantae'}
        # keys['dikarya'] = {'color': 'brown', 'name': 'dikarya'}

        self.add_section(
            name="Taxonomy by kingdom",
            anchor="taxonomy",
            description="The following barplots summarizes the kraken analysis for each sample. ",
            helptext="",
            plot=bargraph.plot(data, keys, pconfig),
        )

    def populate_columns(self):
        headers = {}

        if any(["Classified" in self.sequana_data[s] for s in self.sequana_data]):
            headers["Classified"] = {
                "title": "Classified reads (%)",
                "description": "classified reads (%)",
                "min": 0,
                "max": 100,
                "scale": "RdYlGn",
                "format": "{0:.2f}",
                "shared_key": "count",
            }

        for name in ["Viruses", "Bacteria", "Eukaryota", "Archea", "Fungi"]:
            if any([name in self.sequana_data[s] for s in self.sequana_data]):
                headers[name] = {
                    "title": "Reads classified as {} (%)".format(name.lower()),
                    "description": "Reads classified as {} (%)".format(name.lower()),
                    "min": 0,
                    "max": 100,
                    "scale": "RdYlGn",
                    "format": "{0:.2f}",
                    "shared_key": "count",
                }

        if len(headers.keys()):
            self.general_stats_addcols(self.sequana_data, headers)
