#!/usr/bin/env python

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

# prevent boring warning (version 1.0)
import logging
import os
import re
from collections import OrderedDict

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

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():
                U = v["Unclassified"]
                v["Classified"] = 100 - U
                self.sequana_data[k] = v
            # self.sequana_data[k]["Count"] = int(d["Count"])
        for k, v in self.sequana_data.items():
            self.sequana_data[k]["Count"] = int(v["Count"])

        print(self.sequana_data)

        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 = {
            "id": "classification_vs_unclassified",
            "title": "classification",
            "cpswitch": False,
            "xmin": 0,
            "xmax": 100,
            "xlab_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 barplot 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() if x != "Count"])

        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 = {
            "id": "kraken",
            "title": "Taxonomy by kingdom",
            # "percentages": True,
            "cpswitch": False,
            "xmin": 0,
            "xmax": 100,
            "xlab_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 any(["Count" in self.sequana_data[s] for s in self.sequana_data]):
            headers["count"] = {
                "title": "Total reads",
                "description": "total reads",
                "min": 0,
                # "max": 100,
                "scale": "RdYlGn",
                "format": "{0:d}",
                # "shared_key": "count",
            }
        print(headers)

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