import numpy as np
import pandas as pd

from agora.abc import ParametersABC, ProcessABC
from agora.utils.lineage import mb_array_to_indices


class LineageProcessParameters(ParametersABC):
    """
    Parameters
    """

    def __init__(
        self,
    ):
        super().__init__()

    @classmethod
    def default(cls):
        return cls.from_dict({})


class LineageProcess(ProcessABC):
    """
    Lineage process that must be passed a (N,3) lineage matrix (where the coliumns are trap, mother, daughter respectively)
    """

    def __init__(self, parameters: LineageProcessParameters):
        super().__init__(parameters)

    def run(
        self,
    ):
        pass

    def filter_signal_cells(self, signal: pd.DataFrame):
        # av_indices = np.signal.index.intersection(mb_array_to_indices(self.lineage))]
        # av_indices = np.array(list(signal.index))

        sig_ind = np.array(list(signal.index)).T[:, None, :]
        mo_av = (self.lineage[:, :2].T[:, :, None] == sig_ind).all(axis=0).any(axis=1)
        da_av = (
            (self.lineage[:, [0, 2]].T[:, :, None] == sig_ind).all(axis=0).any(axis=1)
        )

        return self.lineage[mo_av & da_av]
