"""Tests for the swarm layer — agent identity, fusion, provenance, diversity."""

import tempfile
from pathlib import Path

import numpy as np
import pytest

from abiogenesis.swarm import Agent, KTermFusion, DiversityProbe, Provenance
from abiogenesis.soup import Soup
from abiogenesis.interpreter import execute
from abiogenesis.swarm_runner import SwarmExperiment


# ---------------------------------------------------------------------------
# Agent creation
# ---------------------------------------------------------------------------

class TestAgent:
    def test_agent_has_uuid(self):
        tape = np.zeros(64, dtype=np.uint8)
        agent = Agent(tape)
        assert len(agent.id) == 12
        assert agent.id.isalnum()

    def test_agent_has_name(self):
        tape = np.zeros(64, dtype=np.uint8)
        agent = Agent(tape)
        assert agent.name.startswith("agent-")

    def test_agent_custom_name(self):
        tape = np.zeros(64, dtype=np.uint8)
        agent = Agent(tape, name="alice")
        assert agent.name == "alice"

    def test_agent_copies_tape(self):
        """Agent should own a copy, not a reference."""
        tape = np.ones(64, dtype=np.uint8)
        agent = Agent(tape)
        tape[0] = 99
        assert agent.tape[0] == 1  # agent's copy unchanged

    def test_agent_provenance_initialized(self):
        tape = np.zeros(64, dtype=np.uint8)
        agent = Agent(tape)
        assert agent.provenance.agent_id == agent.id
        assert agent.provenance.parent_ids == []
        assert agent.provenance.depth == 0
        assert agent.provenance.merge_interaction == 0
        assert agent.provenance.created_at == 0

    def test_agent_unique_ids(self):
        tape = np.zeros(64, dtype=np.uint8)
        agents = [Agent(tape) for _ in range(100)]
        ids = [a.id for a in agents]
        assert len(set(ids)) == 100  # all unique


# ---------------------------------------------------------------------------
# KTermFusion
# ---------------------------------------------------------------------------

class TestKTermFusion:
    def test_fuse_modifies_tapes(self):
        """Fusion should modify agent tapes (same as soup interact)."""
        rng = np.random.default_rng(42)
        tape_a = rng.integers(0, 256, size=64, dtype=np.uint8)
        tape_b = rng.integers(0, 256, size=64, dtype=np.uint8)

        agent_a = Agent(tape_a)
        agent_b = Agent(tape_b)
        original_a = agent_a.tape.copy()
        original_b = agent_b.tape.copy()

        fusion = KTermFusion(max_steps=10_000)
        ops, _ = fusion.fuse(agent_a, agent_b, interaction=0)

        # At least one tape should be different (ops > 0 for most random tapes)
        # or both unchanged if ops == 0
        if ops > 0:
            changed = (
                not np.array_equal(agent_a.tape, original_a) or
                not np.array_equal(agent_b.tape, original_b)
            )
            assert changed

    def test_fuse_matches_soup_interact(self):
        """KTermFusion should produce identical results to Soup.interact()
        for the same tape content."""
        rng = np.random.default_rng(123)
        tape_a = rng.integers(0, 256, size=64, dtype=np.uint8)
        tape_b = rng.integers(0, 256, size=64, dtype=np.uint8)

        # Method 1: Direct execute (what Soup does)
        combined = bytearray(bytes(tape_a) + bytes(tape_b))
        combined, expected_ops = execute(combined, max_steps=10_000)
        expected_a = np.frombuffer(combined[:64], dtype=np.uint8).copy()
        expected_b = np.frombuffer(combined[64:], dtype=np.uint8).copy()

        # Method 2: KTermFusion
        agent_a = Agent(tape_a)
        agent_b = Agent(tape_b)
        fusion = KTermFusion(max_steps=10_000)
        ops, _ = fusion.fuse(agent_a, agent_b, interaction=0)

        assert ops == expected_ops
        np.testing.assert_array_equal(agent_a.tape, expected_a)
        np.testing.assert_array_equal(agent_b.tape, expected_b)

    def test_fuse_returns_ops_count(self):
        tape = np.zeros(64, dtype=np.uint8)
        agent_a = Agent(tape)
        agent_b = Agent(tape)
        fusion = KTermFusion()
        ops, _ = fusion.fuse(agent_a, agent_b, interaction=0)
        assert isinstance(ops, int)
        assert ops >= 0

    def test_fuse_increments_interaction_count(self):
        tape = np.zeros(64, dtype=np.uint8)
        agent_a = Agent(tape)
        agent_b = Agent(tape)
        fusion = KTermFusion()
        fusion.fuse(agent_a, agent_b, interaction=0)
        assert agent_a.interaction_count == 1
        assert agent_b.interaction_count == 1

    def test_fuse_max_steps_respected(self):
        """With max_steps=0, no ops should execute."""
        rng = np.random.default_rng(42)
        tape_a = rng.integers(0, 256, size=64, dtype=np.uint8)
        tape_b = rng.integers(0, 256, size=64, dtype=np.uint8)
        agent_a = Agent(tape_a)
        agent_b = Agent(tape_b)
        fusion = KTermFusion(max_steps=0)
        ops, _ = fusion.fuse(agent_a, agent_b, interaction=0)
        assert ops == 0


# ---------------------------------------------------------------------------
# Provenance tracking
# ---------------------------------------------------------------------------

class TestProvenance:
    def test_provenance_updates_after_fusion(self):
        rng = np.random.default_rng(42)
        tape_a = rng.integers(0, 256, size=64, dtype=np.uint8)
        tape_b = rng.integers(0, 256, size=64, dtype=np.uint8)

        agent_a = Agent(tape_a)
        agent_b = Agent(tape_b)
        id_a, id_b = agent_a.id, agent_b.id

        fusion = KTermFusion()
        fusion.fuse(agent_a, agent_b, interaction=5)

        # Both agents should reference each other as parents
        assert id_a in agent_a.provenance.parent_ids
        assert id_b in agent_a.provenance.parent_ids
        assert id_a in agent_b.provenance.parent_ids
        assert id_b in agent_b.provenance.parent_ids

        # Merge interaction should be recorded
        assert agent_a.provenance.merge_interaction == 5
        assert agent_b.provenance.merge_interaction == 5

    def test_depth_increments_on_active_interaction(self):
        """Depth should increase when ops > 0."""
        rng = np.random.default_rng(42)
        tape_a = rng.integers(0, 256, size=64, dtype=np.uint8)
        tape_b = rng.integers(0, 256, size=64, dtype=np.uint8)

        agent_a = Agent(tape_a)
        agent_b = Agent(tape_b)
        assert agent_a.provenance.depth == 0

        fusion = KTermFusion()
        ops, _ = fusion.fuse(agent_a, agent_b, interaction=0)

        if ops > 0:
            assert agent_a.provenance.depth == 1
            assert agent_b.provenance.depth == 1

    def test_depth_no_increment_on_zero_ops(self):
        """Depth should not increase when ops == 0."""
        # All zeros tape produces no ops (no BF instructions)
        tape = np.zeros(64, dtype=np.uint8)
        agent_a = Agent(tape)
        agent_b = Agent(tape.copy())

        fusion = KTermFusion()
        ops, _ = fusion.fuse(agent_a, agent_b, interaction=0)
        assert ops == 0
        assert agent_a.provenance.depth == 0
        assert agent_b.provenance.depth == 0

    def test_depth_accumulates_over_multiple_fusions(self):
        """Depth should accumulate across multiple active fusions."""
        rng = np.random.default_rng(99)
        fusion = KTermFusion()

        # Create agents that will produce ops
        agents = []
        for _ in range(4):
            tape = rng.integers(0, 256, size=64, dtype=np.uint8)
            agents.append(Agent(tape))

        # Fuse multiple times and track depth increase
        for i in range(10):
            a_idx = rng.integers(0, len(agents))
            b_idx = rng.integers(0, len(agents))
            while b_idx == a_idx:
                b_idx = rng.integers(0, len(agents))
            fusion.fuse(agents[a_idx], agents[b_idx], interaction=i)

        # At least some agents should have depth > 1
        max_depth = max(a.provenance.depth for a in agents)
        assert max_depth >= 1

    def test_provenance_preserves_created_at(self):
        tape = np.zeros(64, dtype=np.uint8)
        agent = Agent(tape)
        agent.provenance.created_at = 42

        other = Agent(np.ones(64, dtype=np.uint8))
        fusion = KTermFusion()
        fusion.fuse(agent, other, interaction=100)

        # created_at should be preserved through fusion
        assert agent.provenance.created_at == 42


# ---------------------------------------------------------------------------
# DiversityProbe
# ---------------------------------------------------------------------------

class TestDiversityProbe:
    def test_probe_returns_reading(self):
        rng = np.random.default_rng(42)
        agents = [Agent(rng.integers(0, 256, size=64, dtype=np.uint8))
                  for _ in range(32)]
        probe = DiversityProbe()
        reading = probe.measure(agents, interaction=0)

        assert reading.interaction == 0
        assert reading.shannon_entropy >= 0
        assert reading.compression_ratio > 0
        assert reading.alert_level in ("none", "warning", "critical")

    def test_probe_delegates_to_crystallization_probe(self):
        """DiversityProbe should produce same results as CrystallizationProbe
        when given the same data."""
        from abiogenesis.probe import CrystallizationProbe

        rng = np.random.default_rng(42)
        tapes = rng.integers(0, 256, size=(32, 64), dtype=np.uint8)
        agents = [Agent(tapes[i]) for i in range(32)]

        # Direct probe
        direct = CrystallizationProbe(window_size=3, critical_threshold=-0.001)
        direct_reading = direct.measure(tapes, interaction=0)

        # Swarm probe
        swarm_probe = DiversityProbe(window_size=3, critical_threshold=-0.001)
        swarm_reading = swarm_probe.measure(agents, interaction=0)

        assert direct_reading.shannon_entropy == swarm_reading.shannon_entropy
        assert direct_reading.compression_ratio == swarm_reading.compression_ratio
        assert direct_reading.unique_tapes == swarm_reading.unique_tapes

    def test_probe_history(self):
        rng = np.random.default_rng(42)
        agents = [Agent(rng.integers(0, 256, size=64, dtype=np.uint8))
                  for _ in range(16)]
        probe = DiversityProbe()
        probe.measure(agents, interaction=0)
        probe.measure(agents, interaction=100)
        probe.measure(agents, interaction=200)

        history = probe.get_history()
        assert len(history) == 3
        assert history[0].interaction == 0
        assert history[2].interaction == 200


# ---------------------------------------------------------------------------
# Mutation
# ---------------------------------------------------------------------------

class TestSwarmMutation:
    def test_mutation_applies_at_expected_rate(self):
        """With mutation_rate=1.0, every interaction should mutate."""
        config = {
            "n_agents": 16,
            "tape_len": 64,
            "mutation_rate": 1.0,
            "seed": 42,
            "snapshot_interval": 1_000_000,
            "checkpoint_interval": 1_000_000,
        }
        exp = SwarmExperiment(config=config)

        # Record initial state
        initial_tapes = [a.tape.copy() for a in exp.agents]

        # Run a few interactions
        # We can't call run() because it prints, so replicate the core loop
        n = 50
        mutations = 0
        for _ in range(n):
            idx_a, idx_b = exp.rng.choice(len(exp.agents), size=2, replace=False)
            agent_a = exp.agents[idx_a]
            agent_b = exp.agents[idx_b]
            exp.fusion.fuse(agent_a, agent_b, exp.interaction_count)
            if exp.config["mutation_rate"] > 0 and exp.rng.random() < exp.config["mutation_rate"]:
                target = exp.agents[exp.rng.choice([idx_a, idx_b])]
                byte_idx = exp.rng.integers(0, exp.config["tape_len"])
                bit_idx = exp.rng.integers(0, 8)
                target.tape[byte_idx] ^= np.uint8(1 << bit_idx)
                mutations += 1
            exp.interaction_count += 1

        # With rate=1.0, all should mutate
        assert mutations == n


# ---------------------------------------------------------------------------
# SwarmExperiment
# ---------------------------------------------------------------------------

class TestSwarmExperiment:
    def test_init_creates_agents(self):
        config = {"n_agents": 32, "tape_len": 64, "seed": 42,
                  "snapshot_interval": 1_000_000,
                  "checkpoint_interval": 1_000_000}
        exp = SwarmExperiment(config=config)
        assert len(exp.agents) == 32
        for agent in exp.agents:
            assert len(agent.tape) == 64

    def test_run_1000_interactions(self):
        """SwarmExperiment should run 1000 interactions without error."""
        config = {
            "n_agents": 32,
            "tape_len": 64,
            "seed": 42,
            "mutation_rate": 1e-3,
            "snapshot_interval": 500,
            "checkpoint_interval": 1_000_000,
            "output_dir": tempfile.mkdtemp(),
        }
        exp = SwarmExperiment(config=config)
        exp.run(n_interactions=1000)
        assert exp.interaction_count == 1000
        assert len(exp.snapshot_interactions) >= 2  # at 500 and 1000

    def test_run_records_fusion_log(self):
        config = {
            "n_agents": 16,
            "tape_len": 64,
            "seed": 42,
            "mutation_rate": 0.0,
            "snapshot_interval": 1_000_000,
            "checkpoint_interval": 1_000_000,
            "output_dir": tempfile.mkdtemp(),
        }
        exp = SwarmExperiment(config=config)
        exp.run(n_interactions=100)
        graph = exp.get_fusion_graph()
        assert len(graph) == 100
        # Each entry should have required keys
        entry = graph[0]
        assert "interaction" in entry
        assert "agent_a" in entry
        assert "agent_b" in entry
        assert "ops" in entry

    def test_checkpoint_save_load_roundtrip(self):
        """Checkpoint save and load should preserve state."""
        tmpdir = tempfile.mkdtemp()
        config = {
            "n_agents": 16,
            "tape_len": 64,
            "seed": 42,
            "mutation_rate": 1e-3,
            "snapshot_interval": 100,
            "checkpoint_interval": 1_000_000,
            "output_dir": tmpdir,
        }

        # Run and save
        exp1 = SwarmExperiment(config=config)
        exp1.run(n_interactions=500)

        ckpt_path = str(Path(tmpdir) / "test_ckpt.npz")
        exp1.save_checkpoint(ckpt_path)

        # Load
        exp2 = SwarmExperiment.load_checkpoint(ckpt_path)

        # Verify
        assert exp2.interaction_count == exp1.interaction_count
        assert len(exp2.agents) == len(exp1.agents)
        for a1, a2 in zip(exp1.agents, exp2.agents):
            np.testing.assert_array_equal(a1.tape, a2.tape)
            assert a1.id == a2.id
            assert a1.name == a2.name
            assert a1.provenance.depth == a2.provenance.depth
        assert exp2.snapshot_compression == exp1.snapshot_compression

    def test_deterministic_with_seed(self):
        """Same seed should produce same results."""
        config = {
            "n_agents": 16,
            "tape_len": 64,
            "seed": 42,
            "mutation_rate": 1e-3,
            "snapshot_interval": 1_000_000,
            "checkpoint_interval": 1_000_000,
            "output_dir": tempfile.mkdtemp(),
        }

        exp1 = SwarmExperiment(config=config)
        exp1.run(n_interactions=200)

        config["output_dir"] = tempfile.mkdtemp()
        exp2 = SwarmExperiment(config=config)
        exp2.run(n_interactions=200)

        for a1, a2 in zip(exp1.agents, exp2.agents):
            np.testing.assert_array_equal(a1.tape, a2.tape)

    def test_snapshot_records_probe_readings(self):
        config = {
            "n_agents": 16,
            "tape_len": 64,
            "seed": 42,
            "mutation_rate": 1e-3,
            "snapshot_interval": 100,
            "checkpoint_interval": 1_000_000,
            "output_dir": tempfile.mkdtemp(),
        }
        exp = SwarmExperiment(config=config)
        exp.run(n_interactions=300)

        # Should have probe readings at 100, 200, 300
        assert len(exp.snapshot_probe_readings) >= 3
        for reading in exp.snapshot_probe_readings:
            assert "shannon_entropy" in reading
            assert "compression_ratio" in reading
            assert "alert_level" in reading
