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


import logging
import re

from multiqc import config
from multiqc.modules.base_module import BaseMultiqcModule, ModuleNoSamplesFound
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)
            if data and 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(f"Duplicate sample name found! Overwriting: {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)
                # Add version info
                self.add_software_version(data["version"], f["s_name"])

        # 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 ModuleNoSamplesFound

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

        # 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):
        # Check this is a dedup log
        if "# output generated by dedup" not in f["f"]:
            log.debug(f"Skipping as not a dedup log: {f['fn']}")
            return None, None

        # Get the s_name from the input file if possible
        # stdin : <_io.TextIOWrapper name='M18-39155_T1.Aligned.sortedByCoord.out.bam' mode='r' encoding='UTF-8'>
        s_name_re = r"stdin\s+:\s+<_io\.TextIOWrapper name='([^\']+)'"
        s_name_match = re.search(s_name_re, f["f"])
        if s_name_match:
            s_name = s_name_match.group(1)
        else:
            s_name = f["s_name"]

        # Initialise a dictionary to hold the data from this log file
        data = {}

        # Search for values using regular expressions
        regexes = {
            "total_umis": r"INFO total_umis (\d+)",
            "unique_umis": r"INFO #umis (\d+)",
            "input_reads": r"INFO Reads: Input Reads: (\d+)",
            "output_reads": r"INFO Number of reads out: (\d+)",
            "positions_deduplicated": r"INFO Total number of positions deduplicated: (\d+)",
            "mean_umi_per_pos": r"INFO Mean number of unique UMIs per position: ([\d\.]+)",
            "max_umi_per_pos": r"INFO Max. number of unique UMIs per position: (\d+)",
            "version": r"# UMI-tools version: ([\d\.]+)",
        }
        for key, regex in regexes.items():
            re_matches = re.search(regex, f["f"])
            if re_matches:
                if key == "version":
                    data[key] = re_matches.group(1)
                else:
                    data[key] = float(re_matches.group(1))

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

        return s_name, data

    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 = {
            "output_reads": {
                "title": f"{config.read_count_prefix} Unique Reads",
                "description": f"Reads remaining after deduplication ({config.read_count_desc})",
                "min": 0,
                "modify": lambda x: x * config.read_count_multiplier,
                "shared_key": "read_count",
                "scale": "PuRd",
            },
            "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 = {
            "output_reads": {"color": "#7fc97f", "name": "Reads remaining"},
            "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 = {
            "positions_deduplicated": {
                "title": "Positions Dedup",
                "description": "Total number of positions deduplicated",
                "min": 0,
                "format": "{:,.0f}",
                "scale": "Greens",
            },
            "total_umis": {
                "title": "Total UMIs",
                "description": "Total UMIs found in sample",
                "min": 0,
                "format": "{:,.0f}",
                "scale": "Blues",
            },
            "unique_umis": {
                "title": "Unique UMIs",
                "description": "Unique UMIs found in sample",
                "min": 0,
                "format": "{:,.0f}",
                "scale": "Purples",
            },
            "mean_umi_per_pos": {
                "title": "Mean #UMI",
                "description": "Mean number of unique UMIs per position",
                "min": 0,
                "format": "{:,.2f}",
                "scale": "Reds",
            },
            "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)
