import numpy as np
import numba
import strax
import straxen


export, __all__ = strax.exporter()


@export
class PeakProximity(strax.OverlapWindowPlugin):
    """Look for peaks around a peak to determine how many peaks are in proximity (in time) of a
    peak."""

    __version__ = "0.5.0"

    depends_on = "peak_basics"
    dtype = [
        ("n_competing", np.int32, "Number of nearby larger or slightly smaller peaks"),
        (
            "n_competing_left",
            np.int32,
            "Number of larger or slightly smaller peaks left of the main peak",
        ),
        (
            "t_to_prev_peak",
            np.int64,
            "Time between end of previous peak and start of this peak [ns]",
        ),
        ("t_to_next_peak", np.int64, "Time between end of this peak and start of next peak [ns]"),
        ("t_to_nearest_peak", np.int64, "Smaller of t_to_prev_peak and t_to_next_peak [ns]"),
    ] + strax.time_fields

    min_area_fraction = straxen.URLConfig(
        default=0.5,
        infer_type=False,
        help=(
            "The area of competing peaks must be at least "
            "this fraction of that of the considered peak"
        ),
    )

    nearby_window = straxen.URLConfig(
        default=int(1e7),
        infer_type=False,
        help="Peaks starting within this time window (on either side) in ns count as nearby.",
    )

    peak_max_proximity_time = straxen.URLConfig(
        default=int(1e8),
        infer_type=False,
        help="Maximum value for proximity values such as t_to_next_peak [ns]",
    )

    def get_window_size(self):
        return self.peak_max_proximity_time

    def compute(self, peaks):
        # only consider S1 and S2 peaks
        mask = np.isin(peaks["type"], [1, 2])
        windows = strax.touching_windows(peaks[mask], peaks, window=self.nearby_window)
        n_left, n_tot = self.find_n_competing(
            peaks[mask], peaks, windows, fraction=self.min_area_fraction
        )

        t_to_prev_peak = np.ones(len(peaks), dtype=np.int64) * self.peak_max_proximity_time
        t_to_prev_peak[1:] = peaks["time"][1:] - peaks["endtime"][:-1]

        t_to_next_peak = t_to_prev_peak.copy()
        t_to_next_peak[:-1] = peaks["time"][1:] - peaks["endtime"][:-1]

        return dict(
            time=peaks["time"],
            endtime=strax.endtime(peaks),
            n_competing=n_tot,
            n_competing_left=n_left,
            t_to_prev_peak=t_to_prev_peak,
            t_to_next_peak=t_to_next_peak,
            t_to_nearest_peak=np.minimum(t_to_prev_peak, t_to_next_peak),
        )

    @staticmethod
    @numba.jit(nopython=True, nogil=True, cache=True)
    def find_n_competing(peaks_in_roi, peaks, windows, fraction):
        n_left = np.zeros(len(peaks), dtype=np.int32)
        n_tot = n_left.copy()
        areas_in_roi = peaks_in_roi["area"]
        areas = peaks["area"]

        dig = np.searchsorted(peaks_in_roi["center_time"], peaks["center_time"])
        for i, peak in enumerate(peaks):
            left_i, right_i = windows[i]
            threshold = areas[i] * fraction
            n_left[i] = np.sum(areas_in_roi[left_i : dig[i]] > threshold)
            n_tot[i] = n_left[i] + np.sum(areas_in_roi[dig[i] : right_i] > threshold)

        return n_left, n_tot
