"""Population manager — the primordial soup of random tapes."""

from dataclasses import dataclass, field

import numpy as np

from .interpreter import execute


@dataclass(slots=True)
class InteractionResult:
    """Result of a single soup interaction."""
    ops_count: int
    idx_a: int
    idx_b: int
    mutated: bool = False
    depth: int = 0


class Soup:
    """Manages a population of byte tapes that interact pairwise.

    Each interaction:
      1. Pick two tapes at random.
      2. Concatenate them (128 bytes).
      3. Execute as embodied BF.
      4. Split back into two tapes and write them home.
      5. Optionally apply a random bit-flip mutation.
      6. Track symbiogenesis depth.
    """

    def __init__(self, n_tapes: int = 1024, tape_len: int = 64,
                 max_steps: int = 10_000, seed: int | None = None,
                 mutation_rate: float = 0.0, max_tree_depth: int = 0):
        self.n_tapes = n_tapes
        self.tape_len = tape_len
        self.max_steps = max_steps
        self.mutation_rate = mutation_rate
        self.max_tree_depth = max_tree_depth
        self.rng = np.random.default_rng(seed)
        self.tapes = self.rng.integers(
            0, 256, size=(n_tapes, tape_len), dtype=np.uint8
        )
        self.depths = np.zeros(n_tapes, dtype=np.uint32)

    def interact(self) -> InteractionResult:
        """Run one interaction cycle."""
        idx_a, idx_b = self.rng.choice(self.n_tapes, size=2, replace=False)

        # Concatenate two tapes into a single 128-byte working tape
        combined = bytearray(
            bytes(self.tapes[idx_a]) + bytes(self.tapes[idx_b])
        )

        combined, ops = execute(combined, max_steps=self.max_steps)

        # Split back and write home
        half = self.tape_len
        self.tapes[idx_a] = np.frombuffer(combined[:half], dtype=np.uint8)
        self.tapes[idx_b] = np.frombuffer(combined[half:], dtype=np.uint8)

        # Depth tracking: offspring depth = max(parent_depths) + 1 if ops > 0
        parent_depth = max(int(self.depths[idx_a]), int(self.depths[idx_b]))
        new_depth = parent_depth + 1 if ops > 0 else parent_depth
        if self.max_tree_depth > 0:
            new_depth = min(new_depth, self.max_tree_depth)
        self.depths[idx_a] = new_depth
        self.depths[idx_b] = new_depth

        # Mutation: per-interaction probability of random bit-flip
        mutated = False
        if self.mutation_rate > 0 and self.rng.random() < self.mutation_rate:
            target = self.rng.choice([idx_a, idx_b])
            byte_idx = self.rng.integers(0, self.tape_len)
            bit_idx = self.rng.integers(0, 8)
            self.tapes[target, byte_idx] ^= np.uint8(1 << bit_idx)
            mutated = True

        return InteractionResult(
            ops_count=ops, idx_a=idx_a, idx_b=idx_b,
            mutated=mutated, depth=new_depth,
        )

    def get_state(self) -> np.ndarray:
        """Return a copy of the full soup state."""
        return self.tapes.copy()

    def get_depths(self) -> np.ndarray:
        """Return a copy of the depth array."""
        return self.depths.copy()

    def set_state(self, state: np.ndarray, rng_state: dict | None = None,
                  depths: np.ndarray | None = None):
        """Restore soup from a saved state."""
        self.tapes = state.copy()
        if rng_state is not None:
            self.rng.bit_generator.state = rng_state
        if depths is not None:
            self.depths = depths.astype(np.uint32, copy=True)
        else:
            self.depths = np.zeros(state.shape[0], dtype=np.uint32)

    def get_rng_state(self) -> dict:
        """Return serializable RNG state for checkpointing."""
        return self.rng.bit_generator.state
