from .baseanalysis import BaseAnalysis
import numpy as np


class ReciprocityEstimator(BaseAnalysis):
    """
    Estimate the network's reciprocities; here "reciprocity_ei" means:
        Given a connection from E -> I, what is the probability
        of the reciprocated connection to also exist?

    Output keys: reciprocity_ee, reciprocity_ii, reciprocity_ei, reciprocity_ie
    """

    def analyse(self):
        reciprocity = {"reciprocity_ee": self._reciprocity("E", "E"),
                       "reciprocity_ii": self._reciprocity("I", "I"),
                       "reciprocity_ei": self._reciprocity("E", "I"),
                       "reciprocity_ie": self._reciprocity("I", "E")}
        return reciprocity

    def _reciprocity(self, start, end):
        nr_connections_from_start_to_end = np.count_nonzero(self.network[start, end])
        if nr_connections_from_start_to_end == 0:
            return 0
        nr_reciprocated_connections = np.count_nonzero(self.network[start, end] * self.network[end, start].T)
        return nr_reciprocated_connections / nr_connections_from_start_to_end
