import numpy as np
import strax

export, __all__ = strax.exporter()


@export
class DistinctChannels(strax.LoopPlugin):
    """Compute the number of contributing PMTs that contribute to the alt_s1 but not to the main
    S1."""

    __version__ = "0.1.1"
    depends_on = ("event_basics", "peaks")
    loop_over = "events"
    dtype = [
        (
            "alt_s1_distinct_channels",
            np.int32,
            "Number of PMTs contributing to the secondary S1 that do not contribute to the main S1",
        ),
    ] + strax.time_fields

    def compute_loop(self, event, peaks):
        if event["alt_s1_index"] == -1:
            n_distinct = 0
        else:
            s1_a = peaks[event["s1_index"]]
            s1_b = peaks[event["alt_s1_index"]]
            s1_a_peaks = np.nonzero((s1_a["area_per_channel"] > 0) * 1)
            s1_b_peaks = np.nonzero((s1_b["area_per_channel"] > 0) * 1)
            n_distinct = 0
            for channel in range(len(s1_b_peaks[0])):
                if s1_b_peaks[0][channel] not in s1_a_peaks[0]:
                    n_distinct += 1

        return dict(
            alt_s1_distinct_channels=n_distinct, time=event["time"], endtime=event["endtime"]
        )
