"""Tests for symbiogenesis depth tracking."""

import numpy as np
import pytest

from abiogenesis.soup import Soup, InteractionResult
from abiogenesis.metrics import depth_histogram


class TestDepthTracking:
    def test_initial_depths_zero(self):
        """All tapes start at depth 0."""
        soup = Soup(n_tapes=16, tape_len=8, seed=42)
        depths = soup.get_depths()
        assert depths.shape == (16,)
        assert depths.dtype == np.uint32
        np.testing.assert_array_equal(depths, np.zeros(16, dtype=np.uint32))

    def test_depth_increments_on_active_interaction(self):
        """Depth should increase when ops > 0."""
        soup = Soup(n_tapes=16, tape_len=64, seed=42)
        # Run interactions until we find one with ops > 0
        found_active = False
        for _ in range(1000):
            result = soup.interact()
            if result.ops_count > 0:
                found_active = True
                # Both tapes should have depth > 0
                assert soup.depths[result.idx_a] > 0
                assert soup.depths[result.idx_b] > 0
                break
        assert found_active, "Expected at least one active interaction in 1000 tries"

    def test_depth_no_increment_on_zero_ops(self):
        """Depth should NOT increase when ops == 0."""
        soup = Soup(n_tapes=4, tape_len=8, seed=42)
        # Create tapes with no executable content (all zeros)
        soup.tapes = np.zeros((4, 8), dtype=np.uint8)
        soup.depths = np.zeros(4, dtype=np.uint16)

        result = soup.interact()
        if result.ops_count == 0:
            assert soup.depths[result.idx_a] == 0
            assert soup.depths[result.idx_b] == 0

    def test_depth_cap_enforced(self):
        """max_tree_depth should cap the depth."""
        soup = Soup(n_tapes=16, tape_len=64, seed=42, max_tree_depth=5)
        for _ in range(10_000):
            soup.interact()
        depths = soup.get_depths()
        assert depths.max() <= 5

    def test_depth_uncapped(self):
        """max_tree_depth=0 means no cap."""
        soup = Soup(n_tapes=64, tape_len=64, seed=42, max_tree_depth=0)
        for _ in range(10_000):
            soup.interact()
        # With 10k interactions and 64 tapes, depth could grow quite high
        # Just verify it's not artificially capped at some small value
        depths = soup.get_depths()
        # At least some tapes should have depth > 0
        assert depths.max() > 0

    def test_depth_in_interaction_result(self):
        """InteractionResult should carry depth info."""
        soup = Soup(n_tapes=16, tape_len=64, seed=42)
        result = soup.interact()
        assert isinstance(result.depth, int)
        assert result.depth >= 0

    def test_depth_state_save_restore(self):
        """Depths should survive save/restore cycle."""
        soup = Soup(n_tapes=16, tape_len=64, seed=42)
        for _ in range(100):
            soup.interact()

        original_depths = soup.get_depths()
        state = soup.get_state()
        rng_state = soup.get_rng_state()

        # Restore
        soup2 = Soup(n_tapes=16, tape_len=64, seed=99)
        soup2.set_state(state, rng_state, depths=original_depths)
        np.testing.assert_array_equal(soup2.get_depths(), original_depths)

    def test_depth_defaults_on_restore_without_depths(self):
        """set_state without depths should default to zeros."""
        soup = Soup(n_tapes=16, tape_len=64, seed=42)
        for _ in range(100):
            soup.interact()

        state = soup.get_state()
        soup2 = Soup(n_tapes=16, tape_len=64, seed=99)
        soup2.set_state(state)
        np.testing.assert_array_equal(
            soup2.get_depths(), np.zeros(16, dtype=np.uint32)
        )


class TestDepthHistogram:
    def test_all_zero(self):
        """All-zero depths should give {0: n}."""
        depths = np.zeros(16, dtype=np.uint16)
        hist = depth_histogram(depths)
        assert hist == {0: 16}

    def test_mixed_depths(self):
        """Mixed depths should count correctly."""
        depths = np.array([0, 0, 1, 1, 1, 2, 3, 3], dtype=np.uint16)
        hist = depth_histogram(depths)
        assert hist == {0: 2, 1: 3, 2: 1, 3: 2}
