"""
PixelEaterFliesNode (Artificial Life)
--------------------------------
This node simulates a swarm of "flies" (agents) that
live on, consume, and are guided by an input image.
It is an attempt at "artificial life" [cite: "attempt at artificial life"].

- The "World" is the `image_in` (e.g., from a Mandelbrot node).
- The "Flies" are agents with "dumb" [cite: "dumbflies.py"] logic.
- "Dopamine" [cite: "that.. is driven by dopamine addiction"] is Health.
- "Logic" [cite: "thin sheet of logic"] is: "find and eat the
  brightest pixels" [cite: "they can eat the pixels... based on brightness"].
- "Graphics" are inspired by dumbflies.py [cite: "dumbflies.py"].
"""

import numpy as np
import cv2
import time

# --- Magic import block ---
import __main__
BaseNode = __main__.BaseNode
QtGui = __main__.QtGui
# --------------------------

try:
    from scipy.ndimage import gaussian_filter
    SCIPY_AVAILABLE = True
except ImportError:
    SCIPY_AVAILABLE = False
    print("Warning: PixelEaterFliesNode requires 'scipy'.")


class FlyAgent:
    """
    A single fly. It has a body, logic, and health.
    This is the "Dendrite" [cite: "dendrites cause phase space"] or "Scout".
    """
    def __init__(self, x, y, grid_size, config):
        self.x = float(x)
        self.y = float(y)
        self.grid_size = grid_size
        
        self.angle = np.random.uniform(0, 2 * np.pi)
        self.speed = np.random.uniform(1.0, 3.0)
        self.health = 1.0
        self.age = np.random.uniform(0, 100) # For wing flapping

        # Get config
        self.config = config
        
        # --- "Dumb" logic from dumbflies.py ---
        self.state = "WANDER"
        self.state_timer = 0
        self.turn_speed = np.random.uniform(0.1, 0.3)

    def _get_pixel_brightness(self, world, x, y):
        """Helper to safely get brightness at a (wrapped) coordinate"""
        xi = int(x) % self.grid_size
        yi = int(y) % self.grid_size
        return world[yi, xi]

    def perceive_and_decide(self, world):
        """
        The "Thin Logic". The fly "sees" three points in front
        of it and steers towards the brightest one (food).
        """
        
        # 1. Perception: Sample 3 points in front
        dist = self.config['perception_distance']
        center_x = self.x + np.cos(self.angle) * dist
        center_y = self.y + np.sin(self.angle) * dist
        
        angle_left = self.angle - np.pi / 4
        left_x = self.x + np.cos(angle_left) * dist
        left_y = self.y + np.sin(angle_left) * dist
        
        angle_right = self.angle + np.pi / 4
        right_x = self.x + np.cos(angle_right) * dist
        right_y = self.y + np.sin(angle_right) * dist

        food_center = self._get_pixel_brightness(world, center_x, center_y)
        food_left = self._get_pixel_brightness(world, left_x, left_y)
        food_right = self._get_pixel_brightness(world, right_x, right_y)

        # 2. Decision Logic (Steering)
        if food_center > food_left and food_center > food_right:
            # Food is straight ahead, keep going
            self.state = "SEEK"
        elif food_left > food_right:
            # Food is to the left
            self.state = "SEEK"
            self.angle -= self.turn_speed
        elif food_right > food_left:
            # Food is to the right
            self.state = "SEEK"
            self.angle += self.turn_speed
        else:
            # No food in sight, wander
            self.state = "WANDER"
            self.state_timer -= 1
            if self.state_timer <= 0:
                # Pick a new random turn
                self.turn_speed = np.random.uniform(-0.2, 0.2)
                self.state_timer = np.random.randint(10, 50)
            self.angle += self.turn_speed
            
        self.angle %= (2 * np.pi) # Normalize angle

    def update_physics(self):
        """Update position based on angle and speed"""
        self.vel_x = np.cos(self.angle) * self.speed
        self.vel_y = np.sin(self.angle) * self.speed
        
        self.x += self.vel_x
        self.y += self.vel_y
        
        # Wrap around world edges
        self.x %= self.grid_size
        self.y %= self.grid_size
        self.age += 1

    def eat_and_live(self, world):
        """
        Eat pixels to gain health, decay health over time.
        This modifies the "world" (the input image).
        """
        xi, yi = int(self.x), int(self.y)
        
        # 1. Eat food (brightness)
        food_eaten = world[yi, xi]
        if food_eaten > 0.1:
            self.health += food_eaten * self.config['food_value']
            # Modify the world: "eat" the pixel
            world[yi, xi] *= self.config['eat_decay'] 
        
        # 2. Metabolism
        self.health -= self.config['health_decay']
        self.health = np.clip(self.health, 0.0, 1.0)
        
        return self.health > 0 # Return True if alive
        
    def draw(self, display_image):
        """
        Draw the fly with flapping wings, inspired by dumbflies.py
        """
        xi, yi = int(self.x), int(self.y)
        
        # Body
        body_color = (int(self.health * 255), 0, 0) # BGR: Red = healthy
        cv2.circle(display_image, (xi, yi), 3, body_color, -1)
        
        # Wings
        wing_len = 5
        wing_angle_base = np.pi / 2 # Perpendicular to body
        
        # Flap wings
        wing_flap = np.sin(self.age * 0.5) * 0.5 # Flap angle
        
        # Left Wing
        wl_angle = self.angle + wing_angle_base + wing_flap
        wl_x = int(xi + np.cos(wl_angle) * wing_len)
        wl_y = int(yi + np.sin(wl_angle) * wing_len)
        cv2.line(display_image, (xi, yi), (wl_x, wl_y), (200, 200, 200), 1)

        # Right Wing
        wr_angle = self.angle - wing_angle_base - wing_flap
        wr_x = int(xi + np.cos(wr_angle) * wing_len)
        wr_y = int(yi + np.sin(wr_angle) * wing_len)
        cv2.line(display_image, (xi, yi), (wr_x, wr_y), (200, 200, 200), 1)


class PixelEaterFliesNode(BaseNode):
    NODE_CATEGORY = "Artificial Life"
    NODE_COLOR = QtGui.QColor(100, 250, 150) # A-Life Green
    
    def __init__(self, num_flies=50, health_decay=0.01, food_value=0.1, perception_distance=10, eat_decay=0.9):
        super().__init__()
        self.node_title = "Pixel Eater Flies"
        
        self.inputs = {
            'image_in': 'image',     # The "World"
            'reset': 'signal'
        }
        self.outputs = {
            'world_image': 'image',     # The "World" + "Flies"
            'population': 'signal',
            'avg_health': 'signal'
        }
        
        if not SCIPY_AVAILABLE:
            self.node_title = "Flies (No SciPy!)"
            return
            
        self.grid_size = 256 # Default, will adapt to image
        
        # --- Configurable Parameters ---
        self.config = {
            'num_flies': int(num_flies),
            'health_decay': float(health_decay),
            'food_value': float(food_value),
            'perception_distance': int(perception_distance),
            'eat_decay': float(eat_decay) # How much eating darkens a pixel
        }
        
        # --- Internal State ---
        self.flies = []
        self.world = None # This will hold the grayscale, 0-1 "food" map
        self.world_vis = None # This will hold the color "drawing"
        self.avg_health = 0.0

        self._spawn_flies(self.config['num_flies'])

    def _spawn_flies(self, count):
        for _ in range(count):
            x = np.random.randint(0, self.grid_size)
            y = np.random.randint(0, self.grid_size)
            self.flies.append(FlyAgent(x, y, self.grid_size, self.config))

    def _prepare_world(self, img_in):
        """Converts any input image to a 0-1 grayscale float 'food' map"""
        
        # --- Fix for CV_64F crash [cite: "cv2.error... 'depth' is 6 (CV_64F)"] ---
        if img_in.dtype != np.float32:
            img_in = img_in.astype(np.float32)
        if img_in.max() > 1.0:
            img_in = img_in / 255.0
        # --- End Fix ---
        
        if img_in.shape[0] != self.grid_size or img_in.shape[1] != self.grid_size:
            self.grid_size = img_in.shape[0]
            # Respawn flies if world size changes
            self.flies = []
            self._spawn_flies(self.config['num_flies'])

        if img_in.ndim == 3:
            # Convert to grayscale (Luminance)
            world_gray = cv2.cvtColor(img_in, cv2.COLOR_RGB2GRAY)
        else:
            world_gray = img_in.copy()
            
        return world_gray

    def step(self):
        if not SCIPY_AVAILABLE:
            return
            
        # 1. Get Inputs
        reset = self.get_blended_input('reset', 'sum') or 0.0
        img_in = self.get_blended_input('image_in', 'first') # Use 'first'
        
        if reset > 0.5:
            self.flies = []
            self._spawn_flies(self.config['num_flies'])
            self.world = None # Force world reload

        # 2. Update/Prepare the "World"
        if img_in is not None:
            # A new world frame is piped in
            self.world = self._prepare_world(img_in)
        elif self.world is None:
            # No world yet, create a default black one
            self.world = np.zeros((self.grid_size, self.grid_size), dtype=np.float32)
        
        # Create the color visualization frame
        # We draw on a copy so the flies don't "see" their own drawings
        self.world_vis = cv2.cvtColor(self.world, cv2.COLOR_GRAY2RGB)

        # 3. Update all "Flies"
        alive_flies = []
        total_health = 0.0
        
        for fly in self.flies:
            fly.perceive_and_decide(self.world)
            fly.update_physics()
            
            # Eat and check if alive
            if fly.eat_and_live(self.world):
                fly.draw(self.world_vis) # Draw alive flies
                alive_flies.append(fly)
                total_health += fly.health
            else:
                # Fly "died", just don't add it to the list
                pass
        
        self.flies = alive_flies
        
        # 4. Handle Reproduction/Respawning
        missing_flies = self.config['num_flies'] - len(self.flies)
        if missing_flies > 0:
            self._spawn_flies(missing_flies)
            
        # 5. Calculate Metrics
        if len(self.flies) > 0:
            self.avg_health = total_health / len(self.flies)
        else:
            self.avg_health = 0.0

    def get_output(self, port_name):
        if port_name == 'world_image':
            if self.world_vis is None:
                return np.zeros((self.grid_size, self.grid_size, 3), dtype=np.float32)
            return self.world_vis
        elif port_name == 'population':
            return float(len(self.flies))
        elif port_name == 'avg_health':
            return self.avg_health
        return None
        
    def get_display_image(self):
        if self.world_vis is None:
            img = np.zeros((96, 96, 3), dtype=np.uint8)
        else:
            img = cv2.resize(self.world_vis, (96, 96), interpolation=cv2.INTER_NEAREST)
        
        # Add Health Bar
        health_w = int(self.avg_health * (96 - 4))
        cv2.rectangle(img, (2, 96 - 7), (2 + health_w, 96 - 2), (0, 255, 0), -1)
        
        img = np.ascontiguousarray(img)
        return QtGui.QImage(img.data, 96, 96, 96*3, QtGui.QImage.Format.Format_BGR888)

    def get_config_options(self):
        return [
            ("Num Flies", "num_flies", self.config['num_flies'], None),
            ("Health Decay", "health_decay", self.config['health_decay'], None),
            ("Food Value", "food_value", self.config['food_value'], None),
            ("Perception Distance", "perception_distance", self.config['perception_distance'], None),
            ("Eat Decay", "eat_decay", self.config['eat_decay'], None),
        ]

    def close(self):
        self.flies = []
        super().close()