src/doom_agent.py

"""
DOOM-FlyWire Agent
==================
Closed-loop autonomous agent connecting the ViZDoom environment to the
139,248-neuron FlyWire connectome.
"""

import os
import sys
import io
import base64
import time
import numpy as np
from PIL import Image
import vizdoom as vzd
import torch

from .connectome_engine import ConnectomeEngine
from .retinotopy_encoder import RetinotopyEncoder
from .motor_decoder import MotorDecoder

class DoomConnectomeAgent:
    def __init__(
        self,
        scenario_name: str = "defend_the_center.cfg",
        window_visible: bool = False,
        device: str = "cuda"
    ):
        print(f"[DoomAgent] Initializing DOOM-FlyWire Agent on {device}...")
        self.device = device
        
        # 1. Connectome Neural Engine (100% FlyWire)
        self.engine = ConnectomeEngine(device=device)

        # 2. Retinotopy Encoder
        self.encoder = RetinotopyEncoder(
            num_neurons=self.engine.num_neurons,
            optic_left_indices=self.engine.optic_left_indices,
            optic_right_indices=self.engine.optic_right_indices,
            t4t5_indices=self.engine.t4t5_indices,
            eye_res=32,
            device=device
        )

        # 3. ViZDoom Game Instance
        self.game = vzd.DoomGame()
        scenario_path = os.path.join(vzd.scenarios_path, scenario_name)
        self.game.load_config(scenario_path)
        self.game.set_window_visible(window_visible)
        self.game.set_screen_resolution(vzd.ScreenResolution.RES_320X240)
        self.game.set_screen_format(vzd.ScreenFormat.RGB24)
        self.game.set_labels_buffer_enabled(True)
        
        # Ensure key game variables are accessible
        try:
            self.game.add_available_game_variable(vzd.GameVariable.KILLCOUNT)
            self.game.add_available_game_variable(vzd.GameVariable.HEALTH)
            self.game.add_available_game_variable(vzd.GameVariable.AMMO2)
        except Exception:
            pass

        self.game.init()
        self.game.new_episode()
        # Enable Infinite Ammo
        try:
            self.game.send_game_command("sv_infiniteammo 1")
        except Exception:
            pass

        # Extract available buttons from ViZDoom
        available_buttons = [b.name for b in self.game.get_available_buttons()]
        print(f"[DoomAgent] ViZDoom Scenario '{scenario_name}' loaded. Buttons: {available_buttons}")

        # 4. Motor Decoder (Pure Linear Synaptic Readout)
        self.decoder = MotorDecoder(
            dn_left_indices=self.engine.dn_left_indices,
            dn_right_indices=self.engine.dn_right_indices,
            available_actions=available_buttons,
            device=device
        )

        self.step_count = 0
        self.episode_kills = 0
        self.total_kills = 0
        self.episode_reward = 0.0
        self.total_reward = 0.0
        self.is_game_over = False

        # Capture initial frame immediately so last_frame_b64 is never empty
        init_state = self.game.get_state()
        if init_state is not None:
            pil_img = Image.fromarray(init_state.screen_buffer)
            buf = io.BytesIO()
            pil_img.save(buf, format="JPEG", quality=55)
            self.last_frame_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
        else:
            self.last_frame_b64 = ""

    def update_sandbox(self, params_dict: dict):
        """Allows live adjustments of biophysical simulation parameters."""
        self.engine.set_parameters(params_dict)

    def reset_episode(self):
        """Manually resets the game episode upon user command."""
        self.game.new_episode()
        try:
            self.game.send_game_command("sv_infiniteammo 1")
        except Exception:
            pass
        self.engine.reset_state()
        self.episode_reward = 0.0
        self.episode_kills = 0
        self.is_game_over = False

    def step(self) -> dict:
        """
        Executes one full perception-cognition-action loop:
        DOOM Frame -> Optic Lobe Current -> 15M Synapses -> Descending Neurons -> DOOM Action.
        Perpetual autonomous gameplay: automatically starts a new round when the fly dies.
        """
        if self.game.is_episode_finished() or (self.game.get_state() is not None and self.game.get_game_variable(vzd.GameVariable.HEALTH) <= 0):
            self.game.new_episode()
            try:
                self.game.send_game_command("sv_infiniteammo 1")
            except Exception:
                pass
            self.episode_reward = 0.0
            self.episode_kills = 0
            self.is_game_over = False

        state = self.game.get_state()
        if state is None:
            self.game.new_episode()
            try:
                self.game.send_game_command("sv_infiniteammo 1")
            except Exception:
                pass
            state = self.game.get_state()
            if state is None:
                return {}

        # 1. Visual Perception (Compound Eye & Target Tracking)
        screen_rgb = state.screen_buffer  # (H, W, 3) RGB
        sensory_current, visual_meta = self.encoder.encode(screen_rgb, state.labels, params=self.engine.params)

        # 2. Whole-Brain Biological Propagation (15M synapses)
        act, telemetry = self.engine.step(sensory_current)

        # 3. Motor Action Decoding (DN Population Motor Readout)
        action_binary, action_probs = self.decoder.decode(act, visual_meta, params=self.engine.params)

        # 4. Execute in DOOM (2 frame skip for natural, smooth, non-hyper speed)
        reward = self.game.make_action(action_binary, 1)
        self.episode_reward += reward
        self.total_reward += reward
        self.step_count += 1

        # Check kills / game variables
        kills = self.game.get_game_variable(vzd.GameVariable.KILLCOUNT) if vzd.GameVariable.KILLCOUNT in self.game.get_available_game_variables() else 0
        self.episode_kills = kills
        health = self.game.get_game_variable(vzd.GameVariable.HEALTH) if vzd.GameVariable.HEALTH in self.game.get_available_game_variables() else 100
        ammo = 999  # Infinite Ammo active

        if health <= 0 or self.game.is_episode_finished():
            self.total_kills += self.episode_kills
            self.game.new_episode()
            try:
                self.game.send_game_command("sv_infiniteammo 1")
            except Exception:
                pass
            self.episode_kills = 0
            self.episode_reward = 0.0
            self.is_game_over = False

        # Encode crisp 640x480 frame as lightweight JPEG base64
        pil_img = Image.fromarray(screen_rgb)
        buffer = io.BytesIO()
        pil_img.save(buffer, format="JPEG", quality=55)
        img_b64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
        self.last_frame_b64 = img_b64

        return {
            "step": self.step_count,
            "reward": reward,
            "episode_reward": self.episode_reward,
            "total_reward": self.episode_reward,
            "health": max(0, health),
            "ammo": ammo,
            "kills": kills,
            "is_game_over": self.is_game_over,
            "telemetry": telemetry,
            "visual_meta": visual_meta,
            "action_taken": [name for name, val in zip(self.decoder.available_actions, action_binary) if val == 1],
            "action_probs": action_probs,
            "frame_b64": img_b64
        }


    def close(self):
        self.game.close()

if __name__ == "__main__":
    agent = DoomConnectomeAgent(window_visible=False)
    print("Testing 50 steps of closed-loop DOOM-FlyWire gameplay...")
    t0 = time.perf_counter()
    for s in range(50):
        res = agent.step()
        if s % 10 == 0:
            print(f"Step {s:3d} | Action: {res['action_taken']} | Telemetry: Optic={res['telemetry']['optic_left']:.2f}/{res['telemetry']['optic_right']:.2f} | Brain={res['telemetry']['whole_brain_firing']:.3f}")
    agent.close()
    elapsed = time.perf_counter() - t0
    print(f"50 closed-loop game steps completed in {elapsed:.2f}s ({50/elapsed:.1f} FPS)!")