import brian2 as b2
from .base_telemetry import BaseTelemetry

class MeaTelemetry(BaseTelemetry):
    def __init__(self, record_indices):
        self.record_indices = record_indices
        self.state_monitor = None

    def attach(self, organoid):
        from pyopu.organoids.base_organoid import BaseOrganoid
        if not isinstance(organoid, BaseOrganoid):
            raise TypeError("MeaTelemetry must be attached to an Organoid.")
        if 'I_op' not in organoid.neuron_group.equations.names:
            raise ValueError("MeaTelemetry attached to non-MEA organoid.")
        self.organoid = organoid
        self.state_monitor = b2.StateMonitor(organoid.neuron_group, ['I_op'], record=self.record_indices)
        organoid.register_hardware(self.state_monitor)

    def reset(self):
        if hasattr(self, 'organoid') and self.state_monitor is not None:
            if self.state_monitor in self.organoid.network:
                self.organoid.network.remove(self.state_monitor)
            self.state_monitor = b2.StateMonitor(self.organoid.neuron_group, ['I_op'], record=self.record_indices)
            self.organoid.register_hardware(self.state_monitor)

    def get_data(self):
        return {
            "times_ms": self.state_monitor.t / b2.ms,
            "injected_current_nA": self.state_monitor.I_op / b2.nA
        }

class OptoTelemetry(BaseTelemetry):
    def __init__(self, record_indices):
        self.record_indices = record_indices
        self.state_monitor = None

    def attach(self, organoid):
        from pyopu.organoids.base_organoid import BaseOrganoid
        if not isinstance(organoid, BaseOrganoid):
            raise TypeError("OptoTelemetry must be attached to an Organoid.")
        if 'phi_blue' not in organoid.neuron_group.equations.names:
            raise ValueError("OptoTelemetry attached to non-Opto organoid.")
        self.organoid = organoid
        self.state_monitor = b2.StateMonitor(organoid.neuron_group, ['phi_blue', 'phi_yellow'], record=self.record_indices)
        organoid.register_hardware(self.state_monitor)

    def reset(self):
        if hasattr(self, 'organoid') and self.state_monitor is not None:
            if self.state_monitor in self.organoid.network:
                self.organoid.network.remove(self.state_monitor)
            self.state_monitor = b2.StateMonitor(self.organoid.neuron_group, ['phi_blue', 'phi_yellow'], record=self.record_indices)
            self.organoid.register_hardware(self.state_monitor)

    def get_data(self):
        return {
            "times_ms": self.state_monitor.t / b2.ms,
            "light_blue_hz": self.state_monitor.phi_blue / b2.Hz,
            "light_yellow_hz": self.state_monitor.phi_yellow / b2.Hz
        }
