from multiprocessing import cpu_count, Pool

import numpy as np

from wfdb.io.annotation import rdann
from wfdb.io.download import get_record_list
from wfdb.io.record import rdsamp


class Comparitor(object):
    """
    The class to implement and hold comparisons between two sets of
    annotations. See methods `compare`, `print_summary` and `plot`.

    Attributes
    ----------
    ref_sample : ndarray
        An array of the reference sample locations.
    test_sample : ndarray
        An array of the comparison sample locations.
    window_width : int
        The width of the window.
    signal : 1d numpy array, optional
        The signal array the annotation samples are labelling. Only used
        for plotting.

    Examples
    --------
    >>> import wfdb
    >>> from wfdb import processing

    >>> sig, fields = wfdb.rdsamp('sample-data/100', channels=[0])
    >>> ann_ref = wfdb.rdann('sample-data/100','atr')
    >>> xqrs = processing.XQRS(sig=sig[:,0], fs=fields['fs'])
    >>> xqrs.detect()

    >>> comparitor = processing.Comparitor(ann_ref.sample[1:],
                                           xqrs.qrs_inds,
                                           int(0.1 * fields['fs']),
                                           sig[:,0])
    >>> comparitor.compare()
    >>> comparitor.print_summary()
    >>> comparitor.plot()

    """

    def __init__(self, ref_sample, test_sample, window_width, signal=None):
        if len(ref_sample) > 1 and len(test_sample) > 1:
            if min(np.diff(ref_sample)) < 0 or min(np.diff(test_sample)) < 0:
                raise ValueError(
                    (
                        "The sample locations must be monotonically"
                        + " increasing"
                    )
                )

        self.ref_sample = ref_sample
        self.test_sample = test_sample
        self.n_ref = len(ref_sample)
        self.n_test = len(test_sample)
        self.window_width = window_width

        # The matching test sample number for each reference annotation.
        # -1 for indices with no match
        self.matching_sample_nums = np.full(self.n_ref, -1, dtype="int")

        self.signal = signal
        # TODO: rdann return annotations.where

    def _calc_stats(self):
        """
        Calculate performance statistics after the two sets of annotations
        are compared.

        Parameters
        ----------
        N/A

        Returns
        -------
        N/A

        Example:
        -------------------
         ref=500  test=480
        {  30 { 470 } 10  }
        -------------------

        tp = 470
        fp = 10
        fn = 30

        sensitivity = 470 / 500
        positive_predictivity = 470 / 480

        """
        # Reference annotation indices that were detected
        self.matched_ref_inds = np.where(self.matching_sample_nums != -1)[0]
        # Reference annotation indices that were missed
        self.unmatched_ref_inds = np.where(self.matching_sample_nums == -1)[0]
        # Test annotation indices that were matched to a reference annotation
        self.matched_test_inds = self.matching_sample_nums[
            self.matching_sample_nums != -1
        ]
        # Test annotation indices that were unmatched to a reference annotation
        self.unmatched_test_inds = np.setdiff1d(
            np.array(range(self.n_test)),
            self.matched_test_inds,
            assume_unique=True,
        )

        # Sample numbers that were matched and unmatched
        self.matched_ref_sample = self.ref_sample[self.matched_ref_inds]
        self.unmatched_ref_sample = self.ref_sample[self.unmatched_ref_inds]
        self.matched_test_sample = self.test_sample[self.matched_test_inds]
        self.unmatched_test_sample = self.test_sample[self.unmatched_test_inds]

        # True positives = matched reference samples
        self.tp = len(self.matched_ref_inds)
        # False positives = extra test samples not matched
        self.fp = self.n_test - self.tp
        # False negatives = undetected reference samples
        self.fn = self.n_ref - self.tp
        # No tn attribute

        self.sensitivity = float(self.tp) / float(self.tp + self.fn)
        self.positive_predictivity = float(self.tp) / self.n_test

    def compare(self):
        """
        Main comparison function. Note: Make sure to be able to handle
        these ref/test scenarios:

        Parameters
        -------
        N/A

        Returns
        -------
        N/A

        Example
        -------
        A:
        o----o---o---o
        x-------x----x

        B:
        o----o-----o---o
        x--------x--x--x

        C:
        o------o-----o---o
        x-x--------x--x--x

        D:
        o------o-----o---o
        x-x--------x-----x

        """
        test_samp_num = 0
        ref_samp_num = 0

        # Iterate through the reference sample numbers
        while ref_samp_num < self.n_ref and test_samp_num < self.n_test:
            # Get the closest testing sample number for this reference sample
            closest_samp_num, smallest_samp_diff = self._get_closest_samp_num(
                ref_samp_num, test_samp_num
            )
            # Get the closest testing sample number for the next reference
            # sample. This doesn't need to be called for the last index.
            if ref_samp_num < self.n_ref - 1:
                (
                    closest_samp_num_next,
                    smallest_samp_diff_next,
                ) = self._get_closest_samp_num(ref_samp_num + 1, test_samp_num)
            else:
                # Set non-matching value if there is no next reference sample
                # to compete for the test sample
                closest_samp_num_next = -1

            # Found a contested test sample number. Decide which
            # reference sample it belongs to. If the sample is closer to
            # the next reference sample, leave it to the next reference
            # sample and label this reference sample as unmatched.
            if (
                closest_samp_num == closest_samp_num_next
                and smallest_samp_diff_next < smallest_samp_diff
            ):
                # Get the next closest sample for this reference sample,
                # if not already assigned to a previous sample.
                # It will be the previous testing sample number in any
                # possible case (scenario D below), or nothing.
                if closest_samp_num and (
                    not ref_samp_num
                    or closest_samp_num - 1
                    != self.matching_sample_nums[ref_samp_num - 1]
                ):
                    # The previous test annotation is inspected
                    closest_samp_num = closest_samp_num - 1
                    smallest_samp_diff = abs(
                        self.ref_sample[ref_samp_num]
                        - self.test_sample[closest_samp_num]
                    )
                    # Assign the reference-test pair if close enough
                    if smallest_samp_diff < self.window_width:
                        self.matching_sample_nums[
                            ref_samp_num
                        ] = closest_samp_num
                    # Set the starting test sample number to inspect
                    # for the next reference sample.
                    test_samp_num = closest_samp_num + 1

                # Otherwise there is no matching test annotation

            # If there is no clash, or the contested test sample is
            # closer to the current reference, keep the test sample
            # for this reference sample.
            else:
                # Assign the reference-test pair if close enough
                if smallest_samp_diff < self.window_width:
                    self.matching_sample_nums[ref_samp_num] = closest_samp_num
                # Increment the starting test sample number to inspect
                # for the next reference sample.
                test_samp_num = closest_samp_num + 1

            ref_samp_num += 1

        self._calc_stats()

    def _get_closest_samp_num(self, ref_samp_num, start_test_samp_num):
        """
        Return the closest testing sample number for the given reference
        sample number. Limit the search from start_test_samp_num.

        Parameters
        ----------
        ref_samp_num : int
            The desired reference sample number to get the closest result.
        start_test_samp_num
            The desired testing reference sample number to get the
            closest result.

        Returns
        -------
        closest_samp_num : int
            The closest sample number to the reference sample number.
        smallest_samp_diff : int
            The smallest difference between the reference sample and
            the testing sample.

        """
        if start_test_samp_num >= self.n_test:
            raise ValueError("Invalid starting test sample number.")

        ref_samp = self.ref_sample[ref_samp_num]
        test_samp = self.test_sample[start_test_samp_num]
        samp_diff = ref_samp - test_samp

        # Initialize running parameters
        closest_samp_num = start_test_samp_num
        smallest_samp_diff = abs(samp_diff)

        # Iterate through the testing samples
        for test_samp_num in range(start_test_samp_num, self.n_test):
            test_samp = self.test_sample[test_samp_num]
            samp_diff = ref_samp - test_samp
            abs_samp_diff = abs(samp_diff)

            # Found a better match
            if abs_samp_diff < smallest_samp_diff:
                closest_samp_num = test_samp_num
                smallest_samp_diff = abs_samp_diff

            # Stop iterating when the ref sample is first passed or reached
            if samp_diff <= 0:
                break

        return closest_samp_num, smallest_samp_diff

    def print_summary(self):
        """
        Print summary metrics of the annotation comparisons.

        Parameters
        ----------
        N/A

        Returns
        -------
        N/A

        """
        if not hasattr(self, "sensitivity"):
            self._calc_stats()

        print(
            "%d reference annotations, %d test annotations\n"
            % (self.n_ref, self.n_test)
        )
        print("True Positives (matched samples): %d" % self.tp)
        print("False Positives (unmatched test samples): %d" % self.fp)
        print("False Negatives (unmatched reference samples): %d\n" % self.fn)

        print(
            "Sensitivity: %.4f (%d/%d)"
            % (self.sensitivity, self.tp, self.n_ref)
        )
        print(
            "Positive Predictivity: %.4f (%d/%d)"
            % (self.positive_predictivity, self.tp, self.n_test)
        )

    def plot(self, sig_style="", title=None, figsize=None, return_fig=False):
        """
        Plot the comparison of two sets of annotations, possibly
        overlaid on their original signal.

        Parameters
        ----------
        sig_style : str, optional
            The matplotlib style of the signal
        title : str, optional
            The title of the plot
        figsize: tuple, optional
            Tuple pair specifying the width, and height of the figure.
            It is the'figsize' argument passed into matplotlib.pyplot's
            `figure` function.
        return_fig : bool, optional
            Whether the figure is to be returned as an output argument.

        Returns
        -------
        fig : matplotlib figure object
            The figure information for the plot.
        ax : matplotlib axes object
            The axes information for the plot.

        """
        import matplotlib.pyplot as plt

        fig = plt.figure(figsize=figsize)
        ax = fig.add_subplot(1, 1, 1)

        legend = [
            "Signal",
            "Matched Reference Annotations (%d/%d)" % (self.tp, self.n_ref),
            "Unmatched Reference Annotations (%d/%d)" % (self.fn, self.n_ref),
            "Matched Test Annotations (%d/%d)" % (self.tp, self.n_test),
            "Unmatched Test Annotations (%d/%d)" % (self.fp, self.n_test),
        ]

        # Plot the signal if any
        if self.signal is not None:
            ax.plot(self.signal, sig_style)

            # Plot reference annotations
            ax.plot(
                self.matched_ref_sample,
                self.signal[self.matched_ref_sample],
                "ko",
            )
            ax.plot(
                self.unmatched_ref_sample,
                self.signal[self.unmatched_ref_sample],
                "ko",
                fillstyle="none",
            )
            # Plot test annotations
            ax.plot(
                self.matched_test_sample,
                self.signal[self.matched_test_sample],
                "g+",
            )
            ax.plot(
                self.unmatched_test_sample,
                self.signal[self.unmatched_test_sample],
                "rx",
            )

            ax.legend(legend)

        # Just plot annotations
        else:
            # Plot reference annotations
            ax.plot(self.matched_ref_sample, np.ones(self.tp), "ko")
            ax.plot(
                self.unmatched_ref_sample,
                np.ones(self.fn),
                "ko",
                fillstyle="none",
            )
            # Plot test annotations
            ax.plot(self.matched_test_sample, 0.5 * np.ones(self.tp), "g+")
            ax.plot(self.unmatched_test_sample, 0.5 * np.ones(self.fp), "rx")
            ax.legend(legend[1:])

        if title:
            ax.set_title(title)

        ax.set_xlabel("time/sample")

        fig.show()

        if return_fig:
            return fig, ax


def compare_annotations(ref_sample, test_sample, window_width, signal=None):
    """
    Compare a set of reference annotation locations against a set of
    test annotation locations. See the Comparitor class  docstring
    for more information.

    Parameters
    ----------
    ref_sample : 1d numpy array
        Array of reference sample locations.
    test_sample : 1d numpy array
        Array of test sample locations to compare.
    window_width : int
        The maximum absolute difference in sample numbers that is
        permitted for matching annotations.
    signal : 1d numpy array, optional
        The original signal of the two annotations. Only used for
        plotting.

    Returns
    -------
    comparitor : Comparitor object
        Object containing parameters about the two sets of annotations.

    Examples
    --------
    >>> import wfdb
    >>> from wfdb import processing

    >>> sig, fields = wfdb.rdsamp('sample-data/100', channels=[0])
    >>> ann_ref = wfdb.rdann('sample-data/100','atr')
    >>> xqrs = processing.XQRS(sig=sig[:,0], fs=fields['fs'])
    >>> xqrs.detect()

    >>> comparitor = processing.compare_annotations(ann_ref.sample[1:],
                                                    xqrs.qrs_inds,
                                                    int(0.1 * fields['fs']),
                                                    sig[:,0])
    >>> comparitor.print_summary()
    >>> comparitor.plot()

    """
    comparitor = Comparitor(
        ref_sample=ref_sample,
        test_sample=test_sample,
        window_width=window_width,
        signal=signal,
    )
    comparitor.compare()

    return comparitor


def benchmark_mitdb(detector, verbose=False, print_results=False):
    """
    Benchmark a QRS detector against mitdb's records.

    Parameters
    ----------
    detector : function
        The detector function.
    verbose : bool, optional
        The verbose option of the detector function.
    print_results : bool, optional
        Whether to print the overall performance, and the results for
        each record.

    Returns
    -------
    comparitors : dictionary
        Dictionary of Comparitor objects run on the records, keyed on
        the record names.
    sensitivity : float
        Aggregate sensitivity.
    positive_predictivity : float
        Aggregate positive_predictivity.

    Notes
    -----
    TODO:
    - remove non-qrs detections from reference annotations
    - allow kwargs

    Examples
    --------
    >>> import wfdb
    >> from wfdb.processing import benchmark_mitdb, xqrs_detect

    >>> comparitors, spec, pp = benchmark_mitdb(xqrs_detect)

    """
    record_list = get_record_list("mitdb")
    n_records = len(record_list)

    # Function arguments for starmap
    args = zip(record_list, n_records * [detector], n_records * [verbose])

    # Run detector and compare against reference annotations for all
    # records
    with Pool(cpu_count() - 1) as p:
        comparitors = p.starmap(benchmark_mitdb_record, args)

    # Calculate aggregate stats
    sensitivity = np.mean([c.sensitivity for c in comparitors])
    positive_predictivity = np.mean(
        [c.positive_predictivity for c in comparitors]
    )

    comparitors = dict(zip(record_list, comparitors))

    print("Benchmark complete")

    if print_results:
        print(
            "\nOverall MITDB Performance - Sensitivity: %.4f, Positive Predictivity: %.4f\n"
            % (sensitivity, positive_predictivity)
        )
        for record_name in record_list:
            print("Record %s:" % record_name)
            comparitors[record_name].print_summary()
            print("\n\n")

    return comparitors, sensitivity, positive_predictivity


def benchmark_mitdb_record(rec, detector, verbose):
    """
    Benchmark a single mitdb record.

    Parameters
    ----------
    rec : str
        The mitdb record to be benchmarked.
    detector : function
        The detector function.
    verbose : bool
        Whether to print the record names (True) or not (False).

    Returns
    -------
    comparitor : Comparitor object
        Object containing parameters about the two sets of annotations.

    """
    sig, fields = rdsamp(rec, pn_dir="mitdb", channels=[0])
    ann_ref = rdann(rec, pn_dir="mitdb", extension="atr")

    qrs_inds = detector(sig=sig[:, 0], fs=fields["fs"], verbose=verbose)

    comparitor = compare_annotations(
        ref_sample=ann_ref.sample[1:],
        test_sample=qrs_inds,
        window_width=int(0.1 * fields["fs"]),
    )
    if verbose:
        print("Finished record %s" % rec)
    return comparitor