""" MultiQC module to parse output from UMI-tools """


import logging
import os
from collections import OrderedDict

from multiqc import config
from multiqc.modules.base_module import BaseMultiqcModule
from multiqc.plots import bargraph, beeswarm

# Initialise the logger
log = logging.getLogger(__name__)


class MultiqcModule(BaseMultiqcModule):
    """
    umitools module class, parses dedup logs
    """

    def __init__(self):

        # Initialise the parent object
        super(MultiqcModule, self).__init__(
            name="UMI-tools",
            anchor="umitools",
            href="https://github.com/CGATOxford/UMI-tools",
            info="contains tools for dealing with Unique Molecular Identifiers (UMIs)/(RMTs) and scRNA-Seq barcodes.",
            doi="10.1101/gr.209601.116",
        )

        # Find and load any umitools log files
        self.umitools_data = dict()
        for f in self.find_log_files("umitools"):
            # Parse the log file for sample name and statistics
            input_fname, data = self.parse_logs(f["f"])
            if len(data) > 1:
                # Clean the sample name
                f["s_name"] = self.clean_s_name(input_fname, f)
                # Log a warning if the log file matches an existing sample name
                if f["s_name"] in self.umitools_data:
                    log.debug("Duplicate sample name found! Overwriting: {}".format(f["s_name"]))
                # Store the data in the overall data dictionary
                self.umitools_data[f["s_name"]] = data
                # Add the log file information to the multiqc_sources.txt
                self.add_data_source(f)

        # Check the log files against the user supplied list of samples to ignore
        self.umitools_data = self.ignore_samples(self.umitools_data)

        # If no files are found, raise an exception
        if len(self.umitools_data) == 0:
            raise UserWarning

        # Log the number of reports found
        log.info("Found {} reports".format(len(self.umitools_data)))

        # Write parsed report data to a file
        self.write_data_file(self.umitools_data, "multiqc_umitools")

        # write data to the general statistics table
        self.umitools_general_stats_table()

        # add a section with a deduplicated reads plot to the report
        self.add_section(
            name="Deduplicated Reads",
            anchor="umitools-dedup-plot",
            plot=self.umitools_deduplication_plot(),
        )

        # add a section with a beeswarm plot of UMI stats to the report
        self.add_section(
            name="UMI Stats",
            anchor="umitools-umi-stats",
            description="Statistics from running `umi_tools dedeup`",
            helptext="""
            - **Positions Dedup**: Total number of positions deduplicated
            - **Total UMIs**: Total UMIs found in sample
            - **Unique UMIs**: Unique UMIs found in sample
            - **Mean #UMI**: Mean number of unique UMIs per position
            - **Max #UMI**: Max number of unique UMIs per position
            """,
            plot=self.umitools_umi_stats_swarm(),
        )

    def parse_logs(self, f):
        # Initialise a dictionary to hold the data from this log file
        logdata = {}
        # Initialise an empty string for the sample name
        parsed_fname = None
        # Initialise a boolean to determine if the log file is from the dedup command
        dedup_log = False

        # Make a dictionary to hold the data lookup table and data types
        lookup_dict = {
            "total_umis": ("INFO total_umis ", int),
            "unique_umis": ("INFO #umis ", int),
            "input_reads": ("INFO Reads: Input Reads: ", int),
            "output_reads": ("INFO Number of reads out: ", int),
            "positions_deduplicated": ("INFO Total number of positions deduplicated: ", int),
            "mean_umi_per_pos": ("INFO Mean number of unique UMIs per position: ", float),
            "max_umi_per_pos": ("INFO Max. number of unique UMIs per position: ", int),
        }

        # iterate through each line of the log file
        for line in f.splitlines():

            # search for the dedup output string and set the boolean
            if line.startswith("# output generated by dedup"):
                dedup_log = True

            # search for the sample name
            if line.startswith("# stdin"):
                # parse the line and write to the sample name
                parsed_fname = os.path.basename(line.partition("name=")[2].partition(" mode=")[0].strip("'"))

            if dedup_log == True and parsed_fname is not None:
                # iterate through each item in the lookup table
                for key, value in lookup_dict.items():
                    # search for the lookup value
                    if value[0] in line:
                        # parse the line and write to the data dictionary
                        logdata[key] = value[1](line.partition(value[0])[2])

        # calculate a few simple supplementary stats
        try:
            logdata["percent_passing_dedup"] = round(((logdata["output_reads"] / logdata["input_reads"]) * 100.0), 2)
        except (KeyError, ZeroDivisionError):
            pass
        try:
            logdata["removed_reads"] = logdata["input_reads"] - logdata["output_reads"]
        except (KeyError):
            pass

        return parsed_fname, logdata

    def umitools_general_stats_table(self):
        """Take the parsed stats from the umitools report and add it to the
        basic stats table at the top of the report"""

        headers = OrderedDict()
        headers["output_reads"] = {
            "title": "{} Unique Reads".format(config.read_count_prefix),
            "description": "Reads remaining after deduplication ({})".format(config.read_count_desc),
            "min": 0,
            "modify": lambda x: x * config.read_count_multiplier,
            "shared_key": "read_count",
            "scale": "PuRd",
        }
        headers["percent_passing_dedup"] = {
            "title": "% Pass Dedup",
            "description": "% processed reads that passed deduplication",
            "max": 100,
            "min": 0,
            "suffix": "%",
            "scale": "RdYlGn",
        }
        self.general_stats_addcols(self.umitools_data, headers)

    def umitools_deduplication_plot(self):
        """Generate a plot the read deduplication rates for the main report"""

        # Specify the order of the different possible categories
        keys = OrderedDict()
        keys["output_reads"] = {"color": "#7fc97f", "name": "Reads remaining"}
        keys["removed_reads"] = {"color": "#fdc086", "name": "Reads removed"}

        # Config for the plot
        config = {
            "id": "umitools_deduplication",
            "title": "UMI-tools: Deduplication Counts",
            "ylab": "# Reads",
            "cpswitch_counts_label": "Number of Reads",
        }

        return bargraph.plot(self.umitools_data, keys, config)

    def umitools_umi_stats_swarm(self):
        """Generate a swarmplot of umi stats for the main report"""

        headers = OrderedDict()
        headers["positions_deduplicated"] = {
            "title": "Positions Dedup",
            "description": "Total number of positions deduplicated",
            "min": 0,
            "format": "{:,.0f}",
            "scale": "Greens",
        }
        headers["total_umis"] = {
            "title": "Total UMIs",
            "description": "Total UMIs found in sample",
            "min": 0,
            "format": "{:,.0f}",
            "scale": "Blues",
        }
        headers["unique_umis"] = {
            "title": "Unique UMIs",
            "description": "Unique UMIs found in sample",
            "min": 0,
            "format": "{:,.0f}",
            "scale": "Purples",
        }
        headers["mean_umi_per_pos"] = {
            "title": "Mean #UMI",
            "description": "Mean number of unique UMIs per position",
            "min": 0,
            "format": "{:,.2f}",
            "scale": "Reds",
        }
        headers["max_umi_per_pos"] = {
            "title": "Max #UMI",
            "description": "Max number of unique UMIs per position",
            "min": 0,
            "format": "{:,.0f}",
            "scale": "Oranges",
        }

        # Config for the table
        config = {
            "id": "umitools_stats_swarmplot",
            "table_title": "UMI-tools: UMI stats",
        }

        return beeswarm.plot(self.umitools_data, headers, config)
