# code from pairtools software (2025) but not distributed within multiqc.
# added here to avoid another dependency in sequana
# changes made:
# - integrated within sequana as a plugin
# - fixed all pconfig arguments due to deprecated in multiqc.
import logging
import os
import re
from bisect import bisect
from itertools import zip_longest
from math import sqrt
from operator import add, itemgetter
from typing import Dict

import yaml
from multiqc import BaseMultiqcModule, ModuleNoSamplesFound
from multiqc.plots import bargraph, linegraph

log = logging.getLogger(__name__)


def humanize_genomic_dist(dist, units=1):
    """
    turn genomic distance (in basepairs) into
    human readable string (supports Mb, Kb and bp)
    """
    _bp = int(dist) * units
    if _bp // 1_000_000:
        return f"{_bp // 1_000_000}Mb"
    elif _bp // 1_000:
        return f"{_bp // 1_000}Kb"
    else:
        return f"{_bp}bp"


def humanize_genomic_interval(start, end, units=1):
    """
    turn genomic distance interval (basepairs) into
    human readable string
    assuming end > start
    """
    # translate into bp, given provided units
    start_bp = start if start is None else int(start) * units
    end_bp = end if end is None else int(end) * units
    # special case 1st and last intervals:
    if end_bp is None:  # last open-ended
        return f">{humanize_genomic_dist(start_bp)}"
    elif start_bp <= 0:  # first one
        return f"<{humanize_genomic_dist(end_bp)}"
    elif start_bp // 1_000_000:  # both are Mb
        return f"{start_bp // 1_000_000}-{end_bp // 1_000_000}Mb"
    elif end_bp // 1_000_000:  # only end is Mb, start -> float
        return f"{start_bp / 1_000_000:.1f}-{end_bp // 1_000_000}Mb"
    elif start_bp // 1_000:  # both are Kb
        return f"{start_bp // 1_000}-{end_bp // 1_000}Kb"
    elif end_bp // 1_000:  # only end is Kb, start -> float
        return f"{start_bp / 1_000:.1f}-{end_bp // 1_000}Kb"
    else:
        return f"{start_bp}-{end_bp}bp"


def edges_to_intervals(edges):
    """
    turn an array of internal edges into an array of
    intervals, i.e.:
    [1,2,3] -> [(0,1),(1,2),(2,3),(3,None)]
    [0,1,2,3] -> [(0,1),(1,2),(2,3),(3,None)]
    """
    _edges = [*edges, None] if 0 in edges else [0, *edges, None]
    return list(zip(_edges[:-1], _edges[1:]))


def cumsums_to_rangesums(cumsums, total):
    """
    transform cumulative counts to counts in the intervals
    using the total counts information, e.g.:
    [a1,... aN] ->
    [total, a1,... aN, 0] ->
    [total-a1, a1-a2, ... aN-0]
    """
    _cumcsums = [total, *cumsums, 0]
    return [(_h - _l) for _h, _l in zip(_cumcsums[:-1], _cumcsums[1:])]


def parse_dist_range(dist_range):
    """
    turn dist_range_str (as string) into a tuple (start, end).
    Where, end is None for the last open interval.
    """
    dist_range_str = str(dist_range)
    if dist_range_str.endswith("+"):  # last open-ended interval "10000+"
        _start = dist_range_str.strip("+")
        return (int(_start), None)
    elif "-" in dist_range_str:  # other intervals, like "100-2000" and "0-10"
        _start, _end = dist_range_str.split("-")
        return (int(_start), int(_end))
    else:  # unparseable range
        raise ValueError(f"cannot parse distance range {dist_range_str}")


def geometric_mean(start, end):
    """
    return geometric mean for (start, end) interval
    + special casing open ended interval (start, None)
    """
    if start == 0:
        return 0
    elif end is None:
        return start
    else:
        return sqrt(start * end)


def read_stats_from_file(
    file_handle,
    field_separator="\t",
    key_separator="/",
    flat_keys_required=[
        "total",
        "total_unmapped",
        "total_single_sided_mapped",
        "total_mapped",
        "total_dups",
        "total_nodups",
        "cis",
        "trans",
    ],
    pair_orientations_required=["++", "+-", "-+", "--"],
    dist_edges=[0, 100, 500, 1000, 2000, 5000, 10000, 15000, 20000],
):
    """
    parse .stats file generated by pairtools into a dictionary.
    Pairtools-generated .stats file are yaml/json-like, but
    historically they have adopted a custom syntax using "/" to
    separate hierarchical keys of nested dictionaries.

    This function will attempt to parse .stats file, filling out
    output dictionary stat_from_file.

    It'll return None is one of the flat_keys_required is missing
    And it'll has None as a value for some of the internal saections
    if there are some parsing issues or inconsistencies.

    Parameters
    ----------
    file_handle: .stats file handle

    Returns
    -------
    stat_from_file : dict | None
        dictionary with stats valued parsed from .stats
    """

    # instantiate an empty dict for stats
    stat_from_file = {}

    # cis_1kb+, cis_2kb+, cis_4kb+, cis_10kb+, cis_20kb+, cis_40kb+
    cis_dist_pattern = r"cis_(\d+)kb\+"
    cis_dist_list = []

    stat_from_file["pair_types"] = {}

    # store all 4 possible orientations of dist_freq
    dist_freq_orient = {}
    pairs_by_dist_orient = {}
    for _po in pair_orientations_required:
        dist_freq_orient[_po] = []
        pairs_by_dist_orient[_po] = [0 for i in dist_edges]

    # line by line parsing
    for _line in file_handle:
        fields = _line.strip().split(field_separator)
        # skip anything but key[TAB]value-like lines:
        if len(fields) != 2:
            continue
        else:
            _nested_key, _value = fields
            # pursue only numerical values
            try:
                _value = int(_value)
            except ValueError:
                continue
            # parse nested_key
            _key_fields = _nested_key.split(key_separator)
            # (A) key is a simple flat-key (not nested):
            if len(_key_fields) == 1:
                (_key,) = _key_fields
                # (a) one of the required flat key, store as is:
                if _key in flat_keys_required:
                    stat_from_file[_key] = _value
                # (b) one of the cis-dist keys:
                elif _key_match := re.fullmatch(cis_dist_pattern, _key):
                    dist = int(_key_match.group(1))  # extract "dist" from _key
                    cis_dist_list.append((dist, _value))  # [(dist,count)]
                # (c) unknown or unused key - skipping
                else:
                    continue
            # (B) nested _key - 'pair_types', 'dist_freq' will be used
            elif len(_key_fields) > 1:
                # get the first _key and keep the remaining keys in '_subkeys'
                _key, *_subkeys = _key_fields
                # (a) counts by pairtype, e.g. pair_types/UU
                if (_key == "pair_types") and (len(_subkeys) == 1):
                    (_sub_key,) = _subkeys
                    stat_from_file[_key][_sub_key] = _value
                # (b) counts by distance and by pair orientation, e.g. dist_freq/100-3000/-+
                elif (_key == "dist_freq") and (len(_subkeys) == 2):
                    distance_range_str, pair_orientation = _subkeys
                    # skip entries that cannot be parsed
                    try:
                        start, end = parse_dist_range(distance_range_str)
                    except ValueError:
                        continue
                    if pair_orientation in pair_orientations_required:
                        dist = geometric_mean(start, end)
                        _index = bisect(dist_edges, dist) - 1
                        pairs_by_dist_orient[pair_orientation][_index] += _value
                        if dist > 0:
                            dist_freq_orient[pair_orientation].append((dist, _value))
                    else:
                        continue
                # (c) skip all other keys
                else:
                    continue
            # (C) empty or unparseable _nested_key
            else:
                continue

    # subsequent parsing and sanitation

    # we have to have ALL of the required flat fields
    for k in flat_keys_required:
        if k not in stat_from_file:
            return None  # complete failure - entire sample is broken

    # calculation several ratios/fractions based on required fields for general stats table:
    stat_from_file["frac_unmapped"] = stat_from_file["total_unmapped"] / stat_from_file["total"] * 100.0
    stat_from_file["frac_single_sided_mapped"] = (
        stat_from_file["total_single_sided_mapped"] / stat_from_file["total"] * 100.0
    )
    # total_mapped = total_dups + total_nodups
    stat_from_file["frac_mapped"] = stat_from_file["total_mapped"] / stat_from_file["total"] * 100.0
    stat_from_file["frac_dups"] = stat_from_file["total_dups"] / stat_from_file["total"] * 100.0
    # the rest of stats are based on nodups:
    stat_from_file["frac_cis"] = stat_from_file["cis"] / stat_from_file["total_nodups"] * 100.0

    # 'pair_types':
    if not stat_from_file["pair_types"]:
        stat_from_file["pair_types"] = None

    # 'cis_dist'
    if len(cis_dist_list) > 0:
        sorted_dists, cumcounts = zip(*sorted(cis_dist_list, key=itemgetter(0)))
        # cumulative counts -> per-interval counts before
        counts = cumsums_to_rangesums(cumcounts, stat_from_file["cis"])
        # sorted distance -> human readable intervals:
        interval_names = []
        for start_kb, end_kb in edges_to_intervals(sorted_dists):
            _int_name = humanize_genomic_interval(start_kb, end_kb, units=1_000)
            interval_names.append(f"cis: {_int_name}")
        # add trans counts and dist range at the end:
        counts.append(stat_from_file["trans"])
        interval_names.append("trans")
        # store as interval_name -> count dict, ready for plotting:
        stat_from_file["cis_dist"] = dict(zip(interval_names, counts))
    else:
        stat_from_file["cis_dist"] = None

    # 'dist_freq'
    # make sure we have identical number of dist/counts per orientation and it is > 0
    _npo = len(dist_freq_orient[pair_orientations_required[0]])
    if (_npo > 0) and all(len(dist_freq_orient[_po]) == _npo for _po in pair_orientations_required):
        stat_from_file["dist_freq"] = {}
        counts_all_orientations = [0 for i in range(_npo)]
        for _po in pair_orientations_required:
            # parse distance ranges and orientations to store
            dists, counts = zip(*sorted(dist_freq_orient[_po], key=itemgetter(0)))
            stat_from_file["dist_freq"][_po] = dict(zip(dists, counts))
            # accumulate counts for different orientations
            counts_all_orientations = map(add, counts_all_orientations, counts)
        # sum all "pair_orientations_required"
        stat_from_file["dist_freq"]["all"] = dict(zip(dists, counts_all_orientations))
    else:
        stat_from_file["dist_freq"] = None

    # "pairs_by_strand" - assuming dist_freq worked out only !
    if stat_from_file["dist_freq"] is not None:
        stat_from_file["pairs_by_strand"] = {}
        for _po in pair_orientations_required:
            # interval names for pairs by distance by orientation
            for _index, (start, end) in enumerate(edges_to_intervals(dist_edges)):
                _int_name = humanize_genomic_interval(start, end)
                # re-organize the dict to make it ready for plotting ...
                stat_from_file["pairs_by_strand"][_int_name] = {
                    _po: pairs_by_dist_orient[_po][_index] for _po in pair_orientations_required
                }
    else:
        stat_from_file["pairs_by_strand"] = None

    return stat_from_file


class MultiqcModule(BaseMultiqcModule):
    """
    The module parses summary statistics generated by pairtools's `dedup` and `stats` tools.
    """

    def __init__(self):
        super(MultiqcModule, self).__init__(
            name="pairtools",
            anchor="pairtools",
            href="https://github.com/mirnylab/pairtools",
            info="Toolkit for Chromatin Conformation Capture experiments. Handles short-reads paired reference "
            "alignments, extracts 3C-specific information, and perform common tasks such as sorting, filtering, "
            "and deduplication.",
            doi=["10.5281/zenodo.1490830", "10.1101/2023.02.13.528389"],
        )

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

        # Find and load any pairtools stats summary files:
        self.pairtools_stats = {}
        for f in self.find_log_files("pairtools", filehandles=True):
            s_name = f["s_name"]
            if (_sample_report := self.parse_pairtools_stats(f)) is None:
                log.warning(f"{s_name} is missing important metrics will not be reported !")
                continue
            else:
                self.pairtools_stats[s_name] = _sample_report
                # Add file to multiqc_sources.txt
                self.add_data_source(f, s_name=s_name)

        # Filter to strip out ignored sample names
        self.pairtools_stats = self.ignore_samples(self.pairtools_stats)
        if len(self.pairtools_stats) == 0:
            raise ModuleNoSamplesFound

        log.info(f"Found {len(self.pairtools_stats)} reports")

        # Write parsed report data to a file
        self.write_data_file(self.pairtools_stats, "multiqc_pairtools")

        # load various parameters stored in a separate yml (e.g. color schemes)
        with open(os.path.join(os.path.dirname(__file__), "params.yml"), "r") as fp:
            self.params = yaml.safe_load(fp)

        self.pairtools_general_stats()

        # Report sections
        self.add_section(
            name="Pairs by alignment status",
            anchor="pair-types",
            description="Number of pairs classified according to their alignment status,"
            " including uniquely mapped (UU), unmapped (NN), duplicated (DD), and others.",
            helptext="""For further details check
            <a href=\"https://pairtools.readthedocs.io/en/latest/formats.html#pair-types\" > pairtools</a>
            documentation.""",
            plot=self.pair_types_chart(),
        )

        self.add_section(
            name="Pre-filtered pairs by genomic location (overview)",
            anchor="cis-ranges-trans",
            description="Distribution of pre-filtered pairs (mapping uniquely and rescued) by genomic"
            " separation for <it>cis-</it>pairs and <it>trans-</it>pairs.",
            helptext="""
            Samples can have different distributions of pairs by genomic locations for various biological
            and technical reasons. Biological examples: cell-cycle stages, differentitation stages difference,
            mutations affecting genome organization, etc.
            Technical differences arise due to the fact that pre-filtered read pairs still include artifacts:
            Short-range cis-pairs (<1kb) are typically enriched in ligation artifacts (self-circles, dangling-ends, etc).
            Elevated number of trans interactions typically suggests increased noise levels - e.g. random ligations etc.""",
            content=self.pairs_by_cisrange_trans(),
        )

        self.add_section(
            name="Pre-filtered pairs as a function of genomic separation (in detail)",
            anchor="scalings-plots",
            description="Number of interactions (pre-filtered pairs) as a function"
            " of genomic separation, log-binned."
            " Click on an individual curve to reveal information for different"
            " read pair orientations.",
            helptext="""Short-range cis-pairs are typically enriched in ligation artifacts.
            Frequency of interactions for read pairs of different orientations
            ++,+-,-+ and -- (FF, FR, RF, RR) provide insight into these artifacts.
            For example, dangling-ends manifest themselves as FR-pairs, while self-circles - RF.
            Thus, enrichment of FR/RF pairs at a given genomic separation can hint at the level
            of contamination.""",
            plot=self.pairs_with_genomic_separation(),
        )

        self.add_section(
            name="Fraction of read pairs by strand orientation",
            anchor="read-orientation",
            description="Number of pre-filtered pairs reported for every type"
            " of read pair orientation. Numbers are reported for different"
            " ranges of genomic separation and combined.",
            helptext="""Short-range cis-pairs are typically enriched in technical artifacts.
            Frequency of interactions for read pairs of different orientations
            ++,+-,-+ and -- (FF, FR, RF, RR) provide insight into these technical artifacts.
            For example, dangling-ends manifest themselves as FR-pairs, while self-circles - RF.
            Thus, enrichment of FR/RF pairs at a given genomic separation can hint at the level
            of contamination.""",
            plot=self.pairs_by_strand_orientation(),
        )

    @staticmethod
    def parse_pairtools_stats(f):
        """
        Parse a pairtools summary stats
        """
        f_handle = f["f"]
        log.info(f"parsing .stats file: {f_handle.name}")
        return read_stats_from_file(f_handle)

    def pair_types_chart(self):
        """
        Generate the pair types report: a stacked bargraph
        with the number of reads per pair type displayed
        in a pre-defined order and coloring.
        """

        # Construct a data structure for the plot: { sample: {ptype : count} }
        # and keep track of the observed pair types - for consistent visuals:
        data_dict = {}
        _observed_ptypes = set()
        for _sample_name, _sample_data in self.pairtools_stats.items():
            if (_sample_data_dict := _sample_data["pair_types"]) is None:
                # if any of samples is None -> None
                desc = "Skipping the plot as some samples are missing pair types data"
                return f'<div class="alert alert-info">{desc}</div>'
            else:
                data_dict[_sample_name] = _sample_data_dict
                _observed_ptypes |= set(_sample_data_dict)

        # display common pairtypes first with pre-defined colors and order :
        ptypes_annotated: Dict = {}
        for _ptype, ptype_color in self.params["pairtypes_colors"].items():
            if any(t.lower() == _ptype.lower() for t in _observed_ptypes):
                ptypes_annotated[_ptype] = {"color": ptype_color, "name": _ptype}
            else:
                # display remaining pairtypes second with automatic colors :
                ptypes_annotated[_ptype] = {"name": _ptype}

        # multiqc interactive plotting call :
        return bargraph.plot(
            data_dict,
            ptypes_annotated,
            pconfig={
                "id": "pair_types",
                "title": "pairtools: pair types report",
                "ylab": "# Reads",
                "cpswitch_counts_label": "Number of Reads",
            },
        )

    def pairs_by_cisrange_trans(self):
        """
        Cis-pairs split into several ranges of genomic separation, and trans-category.

        Here we extract relevant sub-dict "cis_dist" for each sample, make sure there are
        no None-s and dist_names are the same across all samples.
        """

        data_dict = {}
        _dist_category_set = set()
        for _sample_name, _sample_data in self.pairtools_stats.items():
            if (_sample_data_dict := _sample_data["cis_dist"]) is None:
                # if any of samples is None -> None
                desc = "Skipping the plot as some samples are missing cis distance data"
                return f'<div class="alert alert-info">{desc}</div>'
            else:
                data_dict[_sample_name] = _sample_data_dict
                # make sure all dist_categories are identical
                _dist_category_set |= _sample_data_dict.keys()
                if len(_dist_category_set) > 1:
                    desc = "Skipping the plot as different samples have different distance categories"
                    return f'<div class="alert alert-info">{desc}</div>'

        # extract uniq distance categories
        uniq_dist_categories = _dist_category_set

        # assign colors to distance ranges (as many as possible):
        key_dict: Dict = {}
        for color, _cat_names in zip_longest(self.params["cis_range_colors"], uniq_dist_categories):
            for _cat_name in _cat_names:
                key_dict[_cat_name] = {"name": _cat_name}
                if color:
                    key_dict[_cat_name]["color"] = color
                if _cat_name == "trans":
                    key_dict[_cat_name]["color"] = "#1035AC"  # dark-ish blue

        return bargraph.plot(
            data_dict,
            key_dict,
            pconfig={
                "id": "pair_cis_ranges",
                "title": "pairtools: cis pairs broken into ranges",
                "ylab": "# Reads",
                "cpswitch_counts_label": "Number of reads",
            },
        )

    def pairs_by_strand_orientation(self):
        """
        number of cis-pairs with genomic separation
        """

        # traverse sample_data to extract dist_categories
        _dist_categories = None
        for _sample_name, _sample_data in self.pairtools_stats.items():
            if (_sample_data_dict := _sample_data["pairs_by_strand"]) is None:
                # if any of samples is None -> None
                desc = "Skipping the plot as some samples are missing distance frequencies data"
                return f'<div class="alert alert-info">{desc}</div>'
            else:
                # extract distance categories that should be identical across sample
                _dist_categories = list(_sample_data_dict.keys())

        # initialize data_dict - with distance categories as the outter layer
        if _dist_categories is None:
            desc = "No pairs_by_strand data to plot"
            return f'<div class="alert alert-info">{desc}</div>'

        data_dict: Dict[str, Dict] = {_cat: {} for _cat in _dist_categories}

        # traverse sample_data again to reorder data for plotting
        for _sample_name, _sample_data in self.pairtools_stats.items():
            _sample_data_dict = _sample_data["pairs_by_strand"]
            for _cat, _innermost_dict in _sample_data_dict.items():
                data_dict[_cat][_sample_name] = _innermost_dict

        # annotate read orientations with nice colors:
        keys_annotated = {}
        for key in ["++", "-+", "+-", "--"]:
            name = self.params["pairs_orientation_names"][key]
            color = self.params["pairs_orientation_colors"][key]
            keys_annotated[key] = {"color": color, "name": name}

        return bargraph.plot(
            [data_dict[_dist] for _dist in _dist_categories],
            [keys_annotated for _dist in _dist_categories],
            pconfig={
                "id": "pair_by_orient_cis_ranges",
                "title": "Cis pairs broken into ranges and read orientations",
                "ylab": "# Reads",
                "cpswitch_counts_label": "Number of reads",
                "data_labels": list(data_dict.keys()),
            },
        )

    def pairs_with_genomic_separation(self):
        """
        number of cis-pairs with genomic separation
        # dist_freq/56234133-100000000/-+
        """

        _strand_categories = ["all", "++", "+-", "--", "-+"]
        # initialize data_dict - with distance categories as the outter layer
        data_dict: Dict[str, Dict] = {_cat: {} for _cat in _strand_categories}

        # traverse sample_data to extract dist_categories
        for _sample_name, _sample_data in self.pairtools_stats.items():
            if (_sample_data_dict := _sample_data["dist_freq"]) is None:
                # if any of samples is None -> None
                desc = "Skipping the plot as some samples are missing distance frequencies data"
                return f'<div class="alert alert-info">{desc}</div>'

            # reorder data in _sample_data_dict to extract strand_category on top:
            for _cat in _strand_categories:
                data_dict[_cat][_sample_name] = _sample_data_dict[_cat]

        return linegraph.plot(
            [data_dict[_cat] for _cat in _strand_categories],
            pconfig={
                "id": "broom_plot",
                "title": "pairtools: Pairs by distance and orientation",
                "xlab": "Genomic separation (bp)",
                "ylab": "number of pairs",
                "xlog": True,
                "ylog": True,
                "data_labels": [
                    {"name": "P(s)", "ylab": "number of pairs"},
                    {"name": "FF", "ylab": "number of pairs"},
                    {"name": "FR", "ylab": "number of pairs"},
                    {"name": "RF", "ylab": "number of pairs"},
                    {"name": "RR", "ylab": "number of pairs"},
                ],
            },
        )

    def pairtools_general_stats(self):
        """Add columns to General Statistics table"""
        headers = {
            "total": {
                "title": "Read pairs",
                "description": "Total read pairs before mapping",
                "scale": "Greys",
                "shared_key": "read_count",
            },
            "frac_unmapped": {
                "title": "Unmapped",
                "description": "% of pairs (w.r.t. total) with both sides unmapped",
                "suffix": "%",
                "scale": "PuRd",
            },
            "frac_single_sided_mapped": {
                "title": "One-sided",
                "description": "% of pairs (w.r.t. total) with one side mapped",
                "suffix": "%",
                "scale": "Purples",
            },
            "frac_mapped": {
                "title": "Two-sided",
                "description": "% of pairs (w.r.t. total) with both sides mapped",
                "suffix": "%",
                "scale": "RdYlGn",
            },
            "frac_dups": {
                "title": "Duplicated",
                "description": "% of duplicated pairs (w.r.t. mapped)",
                "suffix": "%",
                "scale": "YlOrRd",
            },
            "total_nodups": {
                "title": "Unique pairs",
                "description": "Mapped pairs after deduplication",
                "scale": "Greys",
                "shared_key": "read_count",
            },
            "frac_cis": {
                "title": "Cis",
                "description": "% of cis-pairs (w.r.t mapped)",
                "suffix": "%",
                "scale": "PuOr-rev",
            },
        }

        self.general_stats_addcols(self.pairtools_stats, headers)
