import numpy as np
import brian2 as b2
from .base_interface import BaseInterface

class MEAInterface(BaseInterface):
    def __init__(self, temperature=1.0, max_current_na=2.0):
        self.T_sys = temperature
        self.max_current = max_current_na * b2.nA
        self.last_spike_count = None

    def get_biological_equations(self):
        return '''
        I_stim = I_op : amp
        I_op : amp
        '''

    def reset(self):
        self.last_spike_count = None

    def read_from_organoid(self, spike_monitor, neurons_per_neobit, num_neobits):
        if self.last_spike_count is None: self.last_spike_count = np.zeros(num_neobits * neurons_per_neobit)
        curr = np.array(spike_monitor.count)
        delta = curr - self.last_spike_count
        self.last_spike_count = curr
        x_inst = np.zeros(num_neobits)
        for i in range(num_neobits):
            if np.sum(delta[i*neurons_per_neobit : (i+1)*neurons_per_neobit]) > 1: x_inst[i] = 1.0
        return x_inst

    def write_to_organoid(self, neuron_group, force_array, neurons_per_neobit):
        prob = 1.0 / (1.0 + np.exp(force_array / self.T_sys))
        currents = (prob - 0.5) * 2.0 * self.max_current
        neuron_group.I_op = np.repeat(currents, neurons_per_neobit)
