import brian2 as b2
from .base_telemetry import BaseTelemetry

class SpikeTelemetry(BaseTelemetry):
    def __init__(self):
        self.spike_monitor = None

    def attach(self, organoid):
        from pyopu.organoids.base_organoid import BaseOrganoid
        if not isinstance(organoid, BaseOrganoid):
            raise TypeError("SpikeTelemetry must be attached to an Organoid.")
        self.organoid = organoid
        self.spike_monitor = b2.SpikeMonitor(organoid.neuron_group)
        organoid.register_hardware(self.spike_monitor)

    def reset(self):
        if hasattr(self, 'organoid') and self.spike_monitor is not None:
            if self.spike_monitor in self.organoid.network:
                self.organoid.network.remove(self.spike_monitor)
            self.spike_monitor = b2.SpikeMonitor(self.organoid.neuron_group)
            self.organoid.register_hardware(self.spike_monitor)

    def get_data(self):
        return {
            "times_ms": self.spike_monitor.t / b2.ms,
            "neuron_ids": self.spike_monitor.i
        }

class OpsinTelemetry(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("OpsinTelemetry must be attached to an Organoid.")
        if 'O_blue' in organoid.neuron_group.equations.names:
            self.organoid = organoid
            self.state_monitor = b2.StateMonitor(
                organoid.neuron_group, ['O_blue', 'O_yellow'], record=self.record_indices
            )
            organoid.register_hardware(self.state_monitor)
        else:
            raise ValueError("OpsinTelemetry attached to an organoid without Opsins.")

    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, ['O_blue', 'O_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,
            "O_blue": self.state_monitor.O_blue,
            "O_yellow": self.state_monitor.O_yellow
        }
