# actors/cortex.py
import time
from actor import Actor


class Cortex(Actor):
    def __init__(self, cid, exoself_pid, sensor_pids, neuron_pids, actuator_pids):
        super().__init__(f"Cortex-{cid}")
        self.cid = cid
        self.sensors = sensor_pids
        self.neurons = neuron_pids
        self.actuators = actuator_pids
        self.exoself_pid = exoself_pid

        self.awaiting_sync = set()  # AIDs, von denen wir im aktuellen Schritt noch Sync erwarten
        self.fitness_acc = 0.0  # akkumulierte Fitness über die Episode
        self.ef_acc = 0  # akkumulierte EndFlags im aktuellen Schritt/Episode
        self.cycle_acc = 0  # Anzahl abgeschlossener Sense-Think-Act Zyklen
        self.active = False  # aktiv/inaktiv (wartet auf reactivate)
        self._t0 = None  # Startzeit der Episode

    async def _kick_sensors(self):
        for s in self.sensors:
            await s.send(("sync",))

    def _reset_for_new_cycle(self):
        self.awaiting_sync = set(a.aid for a in self.actuators)

    def _reset_for_new_episode(self):
        self.fitness_acc = 0.0
        self.ef_acc = 0
        self.cycle_acc = 0
        self._reset_for_new_cycle()
        self._t0 = time.perf_counter()
        self.active = True

    async def run(self):
        # Initialisierung: falls Actuatoren schon gesetzt sind, Episode starten
        if self.actuators:
            self._reset_for_new_episode()
            await self._kick_sensors()

        while True:
            msg = await self.inbox.get()
            tag = msg[0]

            # Optional: Exoself kann AIDs registrieren, bevor wir starten
            if tag == "register_actuators":
                _, aids = msg
                # Map aids auf echte Actuatoren falls nötig; hier übernehmen wir sie direkt
                self.awaiting_sync = set(aids)
                if not self.active:
                    self._reset_for_new_episode()
                    await self._kick_sensors()
                continue

            if tag == "sync" and self.active:
                print("CORTEX: got sync message: ", msg)
                # Aktuator meldet Schrittabschluss
                _t, aid, fitness, halt_flag = msg

                print("----------------")
                print("_t:", _t)
                print("aid:", aid)
                print("fitness:", fitness)
                print("halt_flag:", halt_flag)
                print("----------------")

                # akkumulieren
                self.fitness_acc += float(fitness)
                self.ef_acc += int(halt_flag)

                # diesen Aktuator als "eingetroffen" markieren
                if aid in self.awaiting_sync:
                    self.awaiting_sync.remove(aid)

                print("CORTEX: awaiting sync: ", self.awaiting_sync)

                # Wenn alle Aktuatoren gemeldet haben, ist ein Zyklus fertig
                if not self.awaiting_sync:
                    print("CORTEX: cycle completed.")
                    self.cycle_acc += 1

                    # Episodenende, wenn irgendein Aktuator EndFlag setzt
                    if self.ef_acc > 0:
                        elapsed = time.perf_counter() - self._t0
                        # An Exoself berichten
                        await self.exoself_pid.send((
                            "evaluation_completed",
                            self.fitness_acc,
                            self.cycle_acc,
                            elapsed
                        ))
                        # inaktiv werden – warten auf reactivate
                        self.active = False
                        # Flags für nächste Episode werden erst bei reactivate gesetzt
                    else:
                        # Nächsten Zyklus starten
                        self.ef_acc = 0
                        self._reset_for_new_cycle()
                        await self._kick_sensors()

                continue

            if tag == "reactivate":
                # neue Episode starten
                self._reset_for_new_episode()
                await self._kick_sensors()
                continue

            if tag == "terminate":
                # geordnet alle Kinder beenden
                for a in (self.sensors + self.actuators + self.neurons):
                    await a.send(("terminate",))
                return

            elif tag == "backup_from_neuron":
                # vom Neuron an Cortex -> an Exoself weiterreichen
                await self.exoself_pid.send(msg)