import scipy as sp

from .statistics import pair_correlation
from .util import add_default_output
from ..modeltest import Task


class SORNRecorder:
    def __init__(self, nr_neurons, nr_time_steps):
        self.history = sp.zeros((nr_time_steps, nr_neurons))

    def record(self, time_step, name, value):
        if name != "current_active_neurons":
            return
        self.history[time_step] = value


class UnsynchronizedActivityTask(Task):
    """
    Accepts the SORN model.

    Checks if the activity of the network is synchronized or not.
    """

    @add_default_output
    def __call__(self, model) -> dict:
        model.recorder = SORNRecorder(model.nr_neurons, model.NR_TIME_STEPS)
        model()  # need to call this to produce the history
        show_last_n_time_steps = 100
        correlation = pair_correlation(model.recorder.history[-show_last_n_time_steps:])
        return {"correlation": correlation, "history": model.recorder.history.copy()}
