"""Gate 5a: Hebbian stability — max drift 2%, homeostatic [3%, 5%] band.

Tests:
- Co-activation counts increment
- Bits strengthen proportionally
- Competing traces weaken shared bits
- Stability after 1,000 updates
- Sparsity stays in [3%, 5%]
- Idempotent: same event twice doesn't double
- Profile-scaled hebbian_alpha
"""

from __future__ import annotations

import random
from datetime import datetime, timezone

from harlo.hebbian.learning import (
    SDR_LENGTH,
    HebbianUpdate,
    _TARGET_ACTIVE_HIGH,
    _TARGET_ACTIVE_LOW,
    activation_density,
    apply_hebbian_strengthening,
    compute_effective_sdr,
    record_co_activation,
    record_competition,
)
from harlo.usd_lite.prims import (
    CognitiveProfilePrim,
    MultipliersPrim,
    TracePrim,
)

NOW = datetime(2026, 3, 15, 12, 0, 0, tzinfo=timezone.utc)


def _make_trace(tid: str, density: float = 0.04) -> TracePrim:
    """Make a trace with ~density activation."""
    rng = random.Random(hash(tid))
    n_active = int(SDR_LENGTH * density)
    sdr = [0] * SDR_LENGTH
    for i in rng.sample(range(SDR_LENGTH), n_active):
        sdr[i] = 1
    return TracePrim(
        trace_id=tid,
        sdr=sdr,
        content_hash=f"hash_{tid}",
        strength=0.5,
        last_accessed=NOW,
    )


class TestCoActivation:
    """Co-activation counting."""

    def test_counts_increment(self) -> None:
        t1 = _make_trace("t1")
        t2 = _make_trace("t2")
        a, b = record_co_activation(t1, t2)
        assert a.co_activations["t2"] == 1
        assert b.co_activations["t1"] == 1

    def test_counts_accumulate(self) -> None:
        t1 = _make_trace("t1")
        t2 = _make_trace("t2")
        a, b = record_co_activation(t1, t2)
        a2, b2 = record_co_activation(a, b)
        assert a2.co_activations["t2"] == 2

    def test_independent_of_other_pairs(self) -> None:
        t1 = _make_trace("t1")
        t2 = _make_trace("t2")
        t3 = _make_trace("t3")
        a, _ = record_co_activation(t1, t2)
        a2, _ = record_co_activation(a, t3)
        assert a2.co_activations["t2"] == 1
        assert a2.co_activations["t3"] == 1


class TestCompetition:
    """Competition tracking."""

    def test_competition_counts(self) -> None:
        t1 = _make_trace("t1")
        t2 = _make_trace("t2")
        a, b = record_competition(t1, t2)
        assert a.competitions["t2"] == 1
        assert b.competitions["t1"] == 1


class TestHebbianStrengthening:
    """Bit-level strengthening/weakening with stability."""

    def test_strengthening_produces_update(self) -> None:
        t1 = _make_trace("t1")
        t2 = _make_trace("t2")
        t1.co_activations["t2"] = 10
        update = apply_hebbian_strengthening(t1, [t2], rng_seed=42)
        assert isinstance(update, HebbianUpdate)
        assert len(update.new_strengthen_mask) == SDR_LENGTH

    def test_max_drift_respected(self) -> None:
        """Drift stays within 2% per epoch."""
        t1 = _make_trace("t1")
        t2 = _make_trace("t2")
        t1.co_activations["t2"] = 100  # High co-activation
        update = apply_hebbian_strengthening(t1, [t2], rng_seed=42)
        assert update.drift_ratio <= 0.02 + 1e-9

    def test_stability_after_many_updates(self) -> None:
        """After 1000 updates, SDR doesn't diverge wildly."""
        t1 = _make_trace("t1")
        t2 = _make_trace("t2")
        for i in range(1000):
            t1.co_activations["t2"] = t1.co_activations.get("t2", 0) + 1
            update = apply_hebbian_strengthening(t1, [t2], rng_seed=i)
            t1.hebbian_strengthen_mask = update.new_strengthen_mask
            t1.hebbian_weaken_mask = update.new_weaken_mask

        effective = compute_effective_sdr(
            t1.sdr, t1.hebbian_strengthen_mask, t1.hebbian_weaken_mask
        )
        density = activation_density(effective)
        # Should still be within reasonable bounds
        assert density <= 0.15  # Not exploded

    def test_homeostatic_band(self) -> None:
        """Effective SDR stays in [3%, 5%] activation band after clamping."""
        t1 = _make_trace("t1", density=0.04)
        t2 = _make_trace("t2", density=0.04)
        t1.co_activations["t2"] = 50
        update = apply_hebbian_strengthening(t1, [t2], rng_seed=42)
        effective = compute_effective_sdr(
            t1.sdr, update.new_strengthen_mask, update.new_weaken_mask,
        )
        density = activation_density(effective)
        # After homeostatic clamping
        assert density >= _TARGET_ACTIVE_LOW - 0.01  # Allow small tolerance
        assert density <= _TARGET_ACTIVE_HIGH + 0.01

    def test_profile_scaled_alpha(self) -> None:
        """Different profiles produce different learning rates."""
        t1 = _make_trace("t1")
        t2 = _make_trace("t2")
        t1.co_activations["t2"] = 50

        # Low alpha
        low_profile = CognitiveProfilePrim(
            multipliers=MultipliersPrim(hebbian_alpha=0.001),
        )
        update_low = apply_hebbian_strengthening(
            t1, [t2], profile=low_profile, rng_seed=42,
        )

        # High alpha
        high_profile = CognitiveProfilePrim(
            multipliers=MultipliersPrim(hebbian_alpha=0.05),
        )
        update_high = apply_hebbian_strengthening(
            t1, [t2], profile=high_profile, rng_seed=42,
        )

        # High alpha should strengthen more bits (or equal due to drift cap)
        assert update_high.bits_strengthened >= update_low.bits_strengthened

    def test_no_profile_uses_default(self) -> None:
        """No profile → uses default alpha, no crash."""
        t1 = _make_trace("t1")
        t2 = _make_trace("t2")
        t1.co_activations["t2"] = 10
        update = apply_hebbian_strengthening(t1, [t2], rng_seed=42)
        assert update is not None

    def test_no_co_activations_noop(self) -> None:
        """No co-activations → no bits changed."""
        t1 = _make_trace("t1")
        t2 = _make_trace("t2")
        # No co_activations recorded
        update = apply_hebbian_strengthening(t1, [t2], rng_seed=42)
        assert update.bits_strengthened == 0
