import numpy as np
from .baseanalysis import BaseAnalysis


class ConnectivityEstimator(BaseAnalysis):
    """
    Estimate the connectivity of a network.

    Output Keys: p_ee, p_ii, p_ie, p_ei
    """
    def analyse(self):
        if self.network.nr_exc > 1:
            p_ee = np.count_nonzero(self.network["E", "E"]) / self.network.nr_exc / (self.network.nr_exc - 1)
        else:
            p_ee = 0
        if self.network.nr_inh > 1:
            p_ii = np.count_nonzero(self.network["I", "I"]) / self.network.nr_inh / (self.network.nr_inh - 1)
        else:
            p_ii = 0
        if self.network.nr_exc > 0 and self.network.nr_inh > 0:
            p_ei = np.count_nonzero(self.network["E", "I"]) / self.network.nr_exc / self.network.nr_inh
            p_ie = np.count_nonzero(self.network["I", "E"]) / self.network.nr_exc / self.network.nr_inh
        else:
            p_ei = 0
            p_ie = 0
        return {"p_ee": p_ee, "p_ii": p_ii, "p_ei": p_ei, "p_ie": p_ie}
