import numpy as np
import brian2 as b2
from .base_organoid import BaseOrganoid
from pyopu.telemetry.base_telemetry import Observable

class LIFOrganoid(BaseOrganoid, Observable):
    def __init__(self, num_neobits, neurons_per_neobit, interface, noise_sigma=5.0):
        Observable.__init__(self)
        self.num_neobits = num_neobits
        self.neurons_per_neobit = neurons_per_neobit
        self.x_state = np.zeros(num_neobits)
        
        b2.start_scope()
        b2.defaultclock.dt = 0.1 * b2.ms
        
        base_eqs = '''
        dv/dt = (-70*mV - v + 100*Mohm * I_stim + noise_sigma * (20*ms)**0.5 * xi) / (20*ms) : volt (unless refractory)
        '''
        full_eqs = base_eqs + interface.get_biological_equations()
        
        self.neuron_group = b2.NeuronGroup(
            num_neobits * neurons_per_neobit, full_eqs, 
            threshold='v > -50*mV', reset='v = -75*mV', refractory=2*b2.ms, method='heun',
            namespace={'noise_sigma': noise_sigma * b2.mV}
        )
        self.neuron_group.v = -70*b2.mV + np.random.rand(self.neuron_group.N) * 20*b2.mV
        self._control_spike_mon = b2.SpikeMonitor(self.neuron_group)
        self.network = b2.Network(self.neuron_group, self._control_spike_mon)

    def register_hardware(self, brian_object):
        self.network.add(brian_object)

    def run_step(self, dt_ms):
        self.network.run(dt_ms * b2.ms)
