from collections import defaultdict

import logging
from typing import Dict

from multiqc.base_module import BaseMultiqcModule, ModuleNoSamplesFound
from multiqc.plots import bargraph, table

log = logging.getLogger(__name__)


class MultiqcModule(BaseMultiqcModule):
    """
    The module parses results generated by the `xengsort classify` command.

    **Note**: MultiQC parses the standard output from xengsort, hence one has to redirect
    command line output to a file in order to use it with the MultiQC module. Also note that 
    the tool does not register any sample name information in the output, so MultiQC
    attempts to fetch the sample name from the file name by default. Example command that
     would help MultiQC recognize data for a sample named "SAMPLE":

    ```sh
    xengsort classify --index myindex \
      --fastq paired.1.fq.gz --pairs paired.2.fq.gz \
      --prefix myresults \
      --classification count \
      > SAMPLE.txt
    ```
    """

    def __init__(self):
        super(MultiqcModule, self).__init__(
            name="Xengsort",
            anchor="xengsort",
            href="https://gitlab.com/genomeinformatics/xengsort",
            info="Fast xenograft read sorter based on space-efficient k-mer hashing",
            doi="doi.org/10.4230/LIPIcs.WABI.2020.4",
        )

        # Find and load any Xenome reports
        self.percents = dict()
        self.counts = dict()
        for f in self.find_log_files("xengsort"):
            self._parse_log(f)

        # Filter to strip out ignored sample names
        self.percents = self.ignore_samples(self.percents)
        self.counts = self.ignore_samples(self.counts)
        if len(self.percents) == 0:
            raise ModuleNoSamplesFound
        log.info(f"Found {len(self.counts)} reports")

        # Superfluous function call to confirm that it is used in this module
        # Replace None with actual version if it is available
        self.add_software_version(None)

        # Write parsed report data to a file
        self.write_data_file(self.percents, "multiqc_xengsort_percents")
        self.write_data_file(self.counts, "multiqc_xengsort_counts")

        self._build_table()
        self._build_plot()

    def _parse_log(self, f):
        lines = iter(f["f"].splitlines())
        for line in lines:
            if "\t" not in line:
                continue
            fields = line.strip().split("\t")
            if set(fields) == {"prefix", "host", "graft", "ambiguous", "both", "neither"}:
                values = next(lines).strip().split("\t")
                data = dict(zip(fields, values))
                s_name = data.pop("prefix")
                f["s_name"] = s_name
                data = {k: int(v) for k, v in data.items()}
                percents = {k: v / sum(data.values()) * 100 for k, v in data.items()}

                if s_name in self.counts:
                    log.debug(f"Duplicate sample name found! Overwriting: {s_name}")
                self.add_data_source(f, s_name)
                self.counts[s_name] = data
                self.percents[s_name] = percents
                break

    def _build_table(self):
        """
        Prepare headers and data for a table. Add a section with a table,
        and add a few columns into the general stats.
        """
        headers: Dict[str, Dict] = {}
        table_data = defaultdict(dict)

        scale_by_cls = {
            "graft": "Blues",
            "host": "Reds",
            "both": "Purples",
            "ambiguous": "Greys",
            "neither": "Greys",
        }
        for sn, data in self.percents.items():
            for cls in ["graft", "host", "ambiguous", "both", "neither"]:
                table_data[sn][f"{cls}_reads_pct"] = data.get(cls)
                headers[f"{cls}_reads_pct"] = {
                    "rid": f"xengsort_{cls}_reads_pct",  # to make the ID unique from xenome
                    "title": f"{cls.capitalize()} reads",
                    "description": f"share of {cls} reads in the sample",
                    "min": 0,
                    "suffix": "%",
                    "scale": scale_by_cls[cls],
                    "format": "{:,.2f}",
                    "hidden": cls in ["both", "neither", "ambiguous"],
                }
        self.general_stats_addcols(table_data, headers)
        detail_headers = headers.copy()
        for metric in headers:
            detail_headers[metric]["hidden"] = False

        for sn, data in self.counts.items():
            for cls in ["graft", "host", "ambiguous", "both", "neither"]:
                table_data[sn][f"{cls}_reads_cnt"] = data.get(cls)
                detail_headers[f"{cls}_reads_cnt"] = {
                    "rid": f"xengsort_{cls}_reads_cnt",  # to make the ID unique from xenome
                    "title": f"{cls.capitalize()} reads",
                    "description": f"number of {cls} reads in the sample",
                    "min": 0,
                    "scale": scale_by_cls.get(cls),
                    "format": "{:,d}",
                    "hidden": True,
                }

        self.add_section(
            name="Summary table",
            anchor="xengsort-summary-table-section",
            plot=table.plot(
                table_data,
                detail_headers,
                {
                    "id": "xengsort-summary-table",
                    "namespace": "Xengsort",
                    "title": "Xengsort: Summary",
                },
            ),
        )

    def _build_plot(self):
        """
        Create two bar plots: based on summary and detail data.
        """
        cats = {
            "graft": {"name": "Graft", "color": "#377eb8"},  # blue
            "host": {"name": "Host", "color": "#e41a1c"},  # red
            "both": {"name": "Both", "color": "#984ea3"},  # purple
            "ambiguous": {"name": "Ambiguous", "color": "#616161"},  # grey
            "neither": {"name": "Neither", "color": "#b3b3b3"},  # light grey
        }
        self.add_section(
            description=f"This plot shows the number of reads classified by {self.name}",
            helptext="""
            There are 5 possible categories:  
            * **Graft**: reads found in graft species, e.g. human
            * **Host**: reads found in host species, e.g. mouse
            * **Both**: reads found in either of the species
            * **Neither**: reads was found in neither of the species
            * **Ambiguous**: reads origin could not be adequately determined.  
            """,
            name="Summary classification",
            anchor="xengsort_summary_bar_plot_section",
            plot=bargraph.plot(
                self.counts,
                cats,
                {
                    "id": "xengsort_summary_bar_plot",
                    "title": f"{self.name}: summary classification",
                    "ylab": "# Reads",
                    "cpswitch_counts_label": "Number of reads",
                    "cpswitch_c_active": False,
                    "logswitch": True,
                },
            ),
        )
