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

class OptoInterface(BaseInterface):
    def __init__(self, temperature=1.0, max_rate_hz=200.0):
        self.T_sys = temperature
        self.max_rate_hz = max_rate_hz
        self.last_spike_count = None

    def get_biological_equations(self):
        return '''
        I_stim = I_blue + I_yellow : amp
        I_blue = 2.0 * nA * O_blue * (0*mV - v) / (60*mV) : amp
        I_yellow = 2.0 * nA * O_yellow * (-90*mV - v) / (60*mV) : amp
        dO_blue/dt = (phi_blue * (1 - O_blue) - 0.1/ms * O_blue) : 1
        dO_yellow/dt = (phi_yellow * (1 - O_yellow) - 0.05/ms * O_yellow) : 1
        phi_blue : Hz
        phi_yellow : Hz
        '''

    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):
        num_neobits = len(force_array)
        prob = 1.0 / (1.0 + np.exp(force_array / self.T_sys))
        blue = np.zeros(num_neobits); yellow = np.zeros(num_neobits)
        for i in range(num_neobits):
            if prob[i] > 0.5: blue[i] = (prob[i] - 0.5) * 2.0
            else: yellow[i] = (0.5 - prob[i]) * 2.0
        neuron_group.phi_blue = np.repeat(blue, neurons_per_neobit) * self.max_rate_hz * b2.Hz
        neuron_group.phi_yellow = np.repeat(yellow, neurons_per_neobit) * self.max_rate_hz * b2.Hz
