import logging
import re

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

log = logging.getLogger(__name__)


class MultiqcModule(BaseMultiqcModule):
    """
    The module parses the report and count files generated by
    [Lima](https://github.com/PacificBiosciences/barcoding), a PacBio tool to
    demultiplex PacBio single-molecule sequencing data.

    By default, Lima will use `barcode1--barcode2` as the sample names. To prevent
    these barcodes showing up in the General Statistics table of MultiQC, the Lima
    results are added to their own section.

    If you want to include the Lima results in the General Statistics table, you
    can rename the `barcode1--barcode2` filenames to their apropriate samples using
    the [--replace-names](https://docs.seqera.io/multiqc/#sample-name-replacement)
    option. Each sample that is specified in this way will be moved from the Lima
    section to the General Statistics table.
    """

    def __init__(self):
        # Initialse the parent object
        super(MultiqcModule, self).__init__(
            name="Lima",
            anchor="lima",
            href="https://github.com/PacificBiosciences/barcoding",
            info="Demultiplex PacBio single-molecule sequencing reads.",
            extra="Lima, the PacBio barcode demultiplexer, is the standard tool to identify barcode "
            "sequences in PacBio single-molecule sequencing data. Starting in SMRT Link v5.1.0, it is the "
            "tool that powers the Demultiplex Barcodes GUI-based analysis application.",
            # No publication / DOI // doi=
        )

        # To store the summary data
        self.lima_summary = dict()
        self.lima_counts = dict()

        # Parse the output files
        self.parse_summary_files()
        self.parse_counts_files()

        # Remove filtered samples
        self.lima_summary = self.ignore_samples(self.lima_summary)
        self.lima_counts = self.ignore_samples(self.lima_counts)

        # Write the data files to disk
        if not self.lima_summary and not self.lima_counts:
            raise ModuleNoSamplesFound
        log.info(f"Found {len(self.lima_summary)} reports")
        log.info(f"Found {len(self.lima_counts)} samples")

        if self.lima_summary:
            self.write_data_file(self.lima_summary, "multiqc_lima_summary")
        if self.lima_counts:
            self.write_data_file(self.lima_counts, "multiqc_lima_counts")

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

        # Add a graph of all filtered ZMWs
        self.plot_filter_data()

        # Here, we make a difference between samples that have been renamed,
        # and samples that have their name derived from lima itself. Since lima
        # uses barcode1--barcode2 as sample names, these will never match the
        # other samples in the MultiQC report.
        #
        # Therefore, we do two different things here.
        # 1. All samples that have been renamed will be added to the general
        #       statistics table.
        # 2. All samples that have not been renamed will be added to their own
        #       table in the Lima section of the report.
        lima_renamed_count = dict()
        lima_original_count = dict()

        for sample in self.lima_counts:
            if sample in config.sample_names_replace.values():
                lima_renamed_count[sample] = self.lima_counts[sample]
            else:
                lima_original_count[sample] = self.lima_counts[sample]

        # Get the headers and tconfig for the table
        headers, tconfig = self.make_headers_config()

        # Add a graph for the data values in the counts file, for the samples
        # that haven't been renamed
        if lima_original_count:
            self.make_counts_table(lima_original_count, headers, tconfig)

        # Add renamed samples to the general statistics table, since we assume
        # they are named consistenly now
        if lima_renamed_count:
            self.add_general_stats(lima_renamed_count, headers)

    def parse_summary_files(self):
        for f in self.find_log_files("lima/summary", filehandles=True):
            data = parse_PacBio_log(f["f"])
            if len(data) > 0:
                if f["s_name"] in self.lima_summary:
                    log.debug(f"Duplicate summary sample name found! Overwriting: {f['s_name']}")
                self.lima_summary[f["s_name"]] = data
                self.add_data_source(f)

    def parse_counts_files(self):
        for f in self.find_log_files("lima/counts", filehandles=True):
            data = self.parse_lima_counts(f["f"], f)
            # Check for duplicate samples
            for sample in data:
                if sample in self.lima_counts:
                    log.debug(f"Duplicate counts sample name found! Overwriting: {sample}")
            # After warning the user, overwrite the samples
            self.lima_counts.update(data)
            self.add_data_source(f)

    def parse_lima_counts(self, file_content, f):
        """Parse lima counts file"""

        # The first line is the header
        header = next(file_content).strip().split()

        # A dictionary to store the results
        lima_counts = dict()
        for line in file_content:
            spline = line.strip().split()
            data = {field: value for field, value in zip(header, spline)}

            first_barcode = data["IdxFirstNamed"]
            second_barcode = data["IdxCombinedNamed"]
            # The format barcode1--barcode2 is also used in
            sample = self.clean_s_name(f"{first_barcode}--{second_barcode}", f)
            counts = data["Counts"]
            mean_score = data["MeanScore"]
            lima_counts[sample] = {"Counts": int(counts), "MeanScore": float(mean_score)}

        return lima_counts

    def plot_filter_data(self):
        # First, we gather all filter results for each lima summary
        all_filters = set()
        for data in self.lima_summary.values():
            for reason in self.filter_and_pass(data):
                all_filters.add(reason)
        all_filters = list(all_filters)
        # Move passing reads to the start of the list
        all_filters.insert(0, all_filters.pop(all_filters.index("ZMWs above all thresholds")))

        # Flatten the data into a structure that can be plotted
        plot_data = dict()
        for s_name, data in self.lima_summary.items():
            plot_data[s_name] = dict()
            for reason in all_filters:
                plot_data[s_name][reason] = dict()
                # This is where the filter reasons are stored
                zmw_marginals = data["ZMWs below any threshold"]["ZMW marginals"]
                for filter_reason in zmw_marginals:
                    if filter_reason == reason:
                        plot_data[s_name][reason] = zmw_marginals[filter_reason]["count"]
                        break
                # If we didn't find filter reason for this filename, set it to zero
                else:
                    plot_data[s_name][reason] = 0
                # Unless 'reason' is actually the special case of the number of unfiltered reads
                if reason == "ZMWs above all thresholds":
                    plot_data[s_name][reason] = data["ZMWs above all thresholds"]["count"]

        # Plot configuration
        config = {
            "id": "lima-filter-graph",
            "title": "Lima: ZMW filtering results",
            "ylab": "Number of ZMWs",
            "xlab": "Lima summary file",
        }

        self.add_section(
            name="ZMW filtering",
            anchor="lima-filter",
            description=("The number of ZMWs that failed or passed all Lima filters"),
            helptext="""
                The number of ZMWs that passed all filters is shown as **ZMWs above all thresholds**,
                all other categories that are shown in the graph represent the number of ZMWs that
                were dropped for the specified reason.
            """,
            plot=bargraph.plot(plot_data, all_filters, config),
        )

    def filter_and_pass(self, data):
        """Get the reasons why each ZMW was filtered"""
        reasons = dict()

        # Add why ZMWs were filtered
        filter_data = data["ZMWs below any threshold"]["ZMW marginals"]
        for reason in filter_data:
            reasons[reason] = filter_data[reason]["count"]
        # Add the ZMWs that passed
        reasons["ZMWs above all thresholds"] = data["ZMWs above all thresholds"]["count"]

        return reasons

    def make_headers_config(self):
        """Prepare the headers and config for the lima counts table"""
        tconfig = {
            "id": "multiqc_lima_counts",
            "namespace": "Lima",
            "title": "Lima: Number of Reads",
            "ylab": "# Reads",
        }

        headers = {
            "Counts": {
                "title": f"Read Count ({config.long_read_count_prefix})",
                "description": f"Number of reads for each sample or barcode pair ({config.long_read_count_desc})",
                "modify": lambda x: x * config.long_read_count_multiplier,
                "shared_key": "long_read_counts",
                "format": "{:,.2f}",
                "scale": "PuBuGn",
            },
            "MeanScore": {
                "title": "Quality Score",
                "description": "The mean quality score of the reads for each sample or barcode pair",
                "scale": "Spectral",
            },
        }
        return headers, tconfig

    def make_counts_table(self, counts, headers, tconfig):
        self.add_section(
            name="Per sample count data",
            anchor="multiqc_lima_count",
            description="""
                Per sample or per barcode statistics from Lima.
                For instructions on how to display sample names instead of `barcode--barcode` pairs,
                please see the [MultiQC Lima documentation](https://docs.seqera.io/multiqc/#lima).
            """,
            plot=table.plot(counts, headers, tconfig),
        )

    def add_general_stats(self, counts, headers):
        """Add (renamed) samples to the general statistics table"""
        self.general_stats_addcols(counts, headers)


def parse_PacBio_log(file_content):
    """Parse PacBio log file"""
    data = dict()
    # This is a local dictionary to store which annotations belong to which
    # result dictionary. This will be used to structure the output, but will
    # not be part of the output itself
    annotations = dict()
    current_annotation = None

    for line in file_content:
        # Get rid of trailing newlines
        line = line.strip("\n")
        # Did we enter a new section with annotations for an earlier result?
        # If so, we will only add an empty dictionary with the correct name
        # These field are of the format "something something for (A):"
        section_header_pattern = " for [(][A-Z][)][:]$"
        if re.search(section_header_pattern, line):
            linedata = parse_line(line)
            ann = linedata["annotation"]
            # Cut off the ' for (B):' part
            name = line[:-9]
            # We make a new heading with the current name under the data that
            # matches the current annotation
            current_annotation = dict()
            # We keep the dictonary accessible under 'current_annotation',
            # so we can keep adding new data to it without having to keep track
            # of where it belongs
            annotations[ann][name] = current_annotation
            continue

        linedata = parse_line(line)

        # If we got no data, we reached the end of the annotated section
        if not linedata:
            current_annotation = None
            continue

        # Lets get the name of the data
        name = linedata.pop("name")
        # If we are in an annotated section, we add the data to the current
        # annotation
        if current_annotation is not None:
            current_annotation[name] = linedata
        # Otherwise, we add the newfound annotation to the dictionary in case
        # we find a corresponding section later on.
        # The data from the current line we add directly to the output data
        else:
            annotation = linedata.pop("annotation")
            annotations[annotation] = linedata
            data[name] = linedata

    return data


def parse_line(line):
    """Parse a line from the Lima log file"""
    data = dict()

    # If we got an empty line to parse
    if not line.strip():
        return data

    # Split the line on the colon character
    key, values = line.strip().split(":")

    # The key can have multiple parts
    keys = key.strip().split()

    # Does the key have an annotation (A), (B) etc
    if re.fullmatch("[(][A-Z][)]", keys[-1]):
        # We store the annotation without the bracets
        data["annotation"] = keys[-1][1:-1]
        # And we add the rest of the key as name
        data["name"] = " ".join(keys[:-1])
    # Otherwise, we just store the name
    else:
        data["name"] = " ".join(keys)

    # Parsing the values
    values = values.strip().split()
    # Are there are no values we are done
    if not values:
        return data

    # If there is a single value
    if len(values) == 1:
        value = values.pop()
        # Is the value a percentage
        if value.endswith("%"):
            data["percentage"] = float(value[:-1])
        # Otherwise, it is an integer
        else:
            data["count"] = int(value)
    elif len(values) == 2:
        # If there are multiple values, they are in the format
        # count (percentage%)
        count = values[0]
        percentage = values[1]

        # The percentage should be in the format:
        # (12.34%) or (100%) or (0%) or (-nan%)
        # So we can remove the bracets first
        percentage = percentage[1:-1]
        # Then we make sure it is one of the supported patterns
        assert re.fullmatch(r"(\d+\.\d+%|\d+%|-nan%)", percentage)

        # If the percentage is this weird nan, set it to 0
        if percentage == "-nan%":
            data["percentage"] = 0.0
        else:  # Otherwise, we cut of the % sign and convert to float
            data["percentage"] = float(percentage[:-1])
        # Add the count as an integer
        data["count"] = int(count)

    return data
