src/neuromorphic.py

"""
Direction E: Neuromorphic Cortical-Hippocampal Loop for PC-SHO-DLM + MSA

Maps the PC-SHO-DLM + MSA architecture onto biologically-named circuits:

    Cortical columns     -> forward/feedback blocks  (bottom-up / top-down)
    Hippocampal memory   -> MSA MemoryBank           (compressed KV store)
    Entorhinal routing   -> MSA RouterProjector      (query-key relevance)
    Reconsolidation      -> post-settle param update from settled errors

This is a *wrapper* over the existing architecture -- no reimplementation.
Every public method delegates to real PCSHODLM / MSA methods and logs a
structured timeline of (timestamp, event_type, magnitude) events for
downstream analysis and prediction validation.
"""

import time
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F

from src.model import PCSHODLM, PCSHOConfig
from src.msa import (
    MSAConfig,
    MSALayer,
    MemoryBank,
    MemoryEncoder,
    RouterProjector,
    chunk_mean_pool,
    create_msa_layers,
)


# -------------------------------------------------------------------------
# Timeline event record
# -------------------------------------------------------------------------

@dataclass
class TimelineEvent:
    """One observable event in the cortical-hippocampal loop."""
    timestamp: float       # monotonic wall-clock seconds since loop creation
    event_type: str        # "cortical_error" | "hippocampal_retrieval" | "reconsolidation"
    magnitude: float       # scalar summary of the event's intensity
    metadata: dict = field(default_factory=dict)


# -------------------------------------------------------------------------
# CorticalHippocampalLoop
# -------------------------------------------------------------------------

class CorticalHippocampalLoop(nn.Module):
    """Biologically-named wrapper mapping PC-SHO-DLM + MSA to brain circuits.

    Lifecycle for one stimulus (sequence of tokens):

        1. cortex_process(stimulus)              -- bottom-up forward pass
        2. hippocampus_query(cortical_state)      -- entorhinal routing scores
        3. hippocampus_retrieve(indices)           -- load compressed KV
        4. cortex_settle(stimulus, memory)         -- PC settling with context
        5. reconsolidate(settled_state)            -- unified parameter updates

    Every step is logged to ``self.timeline``.
    """

    def __init__(
        self,
        model: PCSHODLM,
        msa_layers: nn.ModuleList,
        memory_bank: MemoryBank,
        msa_config: Optional[MSAConfig] = None,
    ):
        super().__init__()
        self.model = model
        self.msa_layers = msa_layers
        self.memory_bank = memory_bank
        self.msa_config = msa_config or MSAConfig()
        self._last_x_t: Optional[torch.Tensor] = None

        # Which model layers correspond to MSA-augmented layers
        n_layers = model.config.n_layers
        self._msa_start = max(0, n_layers - len(msa_layers))

        # Observable timeline
        self.timeline: List[TimelineEvent] = []
        self._t0 = time.monotonic()

    # -- internal helpers --------------------------------------------------

    def _ts(self) -> float:
        """Seconds since construction."""
        return time.monotonic() - self._t0

    def _log(self, event_type: str, magnitude: float, **meta) -> None:
        self.timeline.append(
            TimelineEvent(
                timestamp=self._ts(),
                event_type=event_type,
                magnitude=magnitude,
                metadata=meta,
            )
        )

    # =====================================================================
    # 1. cortex_process  (cortical columns -> forward pass)
    # =====================================================================

    def cortex_process(
        self,
        stimulus: torch.Tensor,
        t: Optional[torch.Tensor] = None,
    ) -> Tuple[list, torch.Tensor, torch.Tensor]:
        """Bottom-up forward pass through cortical columns.

        Args:
            stimulus: (B, S) token ids (possibly partially masked).
            t: (B,) diffusion timesteps.  Defaults to t=1 (near-clean).

        Returns:
            h_init: list of hidden states [h_0 ... h_L] from amortized pass.
            mask:   (B, S) boolean mask of corrupted positions.
            t:      timestep tensor used.
        """
        B, S = stimulus.shape
        device = stimulus.device

        if t is None:
            t = torch.ones(B, dtype=torch.long, device=device)

        # Corrupt at the requested timestep
        x_t, mask = self.model.schedule.corrupt(stimulus, t, self.model.config.mask_token_id)
        self._last_x_t = x_t.detach().clone()

        # Embed and run amortized forward pass (bottom-up cortical sweep)
        h_0 = self.model.embed_input(x_t, t)
        h_init = self.model.amortized_forward_pass(h_0)

        # Log cortical prediction error (mean L2 across layers)
        with torch.no_grad():
            mu_up, mu_down = self.model.compute_predictions(h_init)
            L = self.model.config.n_layers
            eps_up = [h_init[l + 1] - mu_up[l] for l in range(L)]
            error_mag = sum(e.norm().item() for e in eps_up) / L

        self._log("cortical_error", error_mag, stage="process")
        return h_init, mask, t

    # =====================================================================
    # 2. hippocampus_query  (entorhinal routing -> MSA router scores)
    # =====================================================================

    def hippocampus_query(
        self,
        cortical_state: list,
    ) -> Tuple[Dict[int, torch.Tensor], Dict[int, List[str]]]:
        """Route cortical state through entorhinal projectors to score memories.

        Args:
            cortical_state: hidden states [h_0 ... h_L].

        Returns:
            all_scores: {layer_idx: (B, N_chunks) routing scores}
            all_chunk_ids: {layer_idx: list of doc-chunk identifiers}
        """
        all_scores: Dict[int, torch.Tensor] = {}
        all_chunk_ids: Dict[int, List[str]] = {}

        if len(self.memory_bank) == 0:
            return all_scores, all_chunk_ids

        for msa_idx, msa_layer in enumerate(self.msa_layers):
            layer_idx = self._msa_start + msa_idx
            if layer_idx >= len(cortical_state) - 1:
                continue

            routing_keys, chunk_ids = self.memory_bank.get_routing_keys(layer_idx)
            if routing_keys is None:
                continue

            h_query = cortical_state[layer_idx + 1]  # hidden state at this layer
            scores = msa_layer.compute_routing_scores(h_query, routing_keys)
            all_scores[layer_idx] = scores
            all_chunk_ids[layer_idx] = chunk_ids

        # Log retrieval event: mean top-1 score across layers
        if all_scores:
            top1_vals = [s.max().item() for s in all_scores.values()]
            mean_top1 = sum(top1_vals) / len(top1_vals)
        else:
            mean_top1 = 0.0
        self._log("hippocampal_retrieval", mean_top1, stage="query")

        return all_scores, all_chunk_ids

    # =====================================================================
    # 3. hippocampus_retrieve  (load compressed KV for winning docs)
    # =====================================================================

    def hippocampus_retrieve(
        self,
        scores: Dict[int, torch.Tensor],
        chunk_ids: Dict[int, List[str]],
        top_k: Optional[int] = None,
    ) -> Dict[int, Tuple[torch.Tensor, torch.Tensor]]:
        """Select top-k documents and fetch their compressed KV pairs.

        Args:
            scores:    per-layer routing scores from hippocampus_query.
            chunk_ids: per-layer chunk doc-id lists.
            top_k:     override for msa_config.top_k.

        Returns:
            memory_kv: {layer_idx: (K_compressed, V_compressed)} both (B, S_mem, D).
        """
        top_k = top_k or self.msa_config.top_k
        memory_kv: Dict[int, Tuple[torch.Tensor, torch.Tensor]] = {}

        for layer_idx, layer_scores in scores.items():
            cids = chunk_ids[layer_idx]
            B = layer_scores.shape[0]

            # Pick top-k chunks per batch element (use first batch element for doc selection)
            k = min(top_k, layer_scores.shape[1])
            _, topk_idx = layer_scores[0].topk(k)
            selected_docs = list({cids[i.item()] for i in topk_idx})

            K_sel, V_sel = self.memory_bank.get_kv(selected_docs, layer_idx)
            if K_sel is not None:
                # Expand to batch dimension
                K_sel = K_sel.unsqueeze(0).expand(B, -1, -1)
                V_sel = V_sel.unsqueeze(0).expand(B, -1, -1)
                memory_kv[layer_idx] = (K_sel, V_sel)

        # Log
        n_retrieved = sum(kv[0].shape[1] for kv in memory_kv.values()) if memory_kv else 0
        self._log("hippocampal_retrieval", float(n_retrieved), stage="retrieve")

        return memory_kv

    # =====================================================================
    # 4. cortex_settle  (PC settling with retrieved hippocampal context)
    # =====================================================================

    def cortex_settle(
        self,
        stimulus: torch.Tensor,
        h_init: list,
        mask: torch.Tensor,
        t: torch.Tensor,
        memory_kv: Optional[Dict[int, Tuple[torch.Tensor, torch.Tensor]]] = None,
    ) -> Tuple[list, list, list]:
        """Run PC-SHO settling, optionally injecting hippocampal memory.

        If ``memory_kv`` is provided, the retrieved context is injected into
        the corresponding MSA layers before each settling iteration by
        running an MSA forward pass that blends local representations with
        retrieved KV pairs.

        Args:
            stimulus: (B, S) original (uncorrupted) token ids -- used as x_0.
            h_init:   amortized hidden states from cortex_process.
            mask:     corruption mask.
            t:        timestep tensor.
            memory_kv: retrieved memory from hippocampus_retrieve.

        Returns:
            h_settled: settled hidden states.
            v_final:   final velocities.
            energies:  energy trace.
        """
        # Inject memory into h_init before settling: for each MSA layer,
        # run a single sparse-attention pass that conditions h on memory.
        if memory_kv:
            for layer_idx, (K_mem, V_mem) in memory_kv.items():
                msa_idx = layer_idx - self._msa_start
                if 0 <= msa_idx < len(self.msa_layers):
                    h_layer = h_init[layer_idx + 1]
                    h_init[layer_idx + 1] = self.msa_layers[msa_idx](
                        h_layer, memory_k=K_mem, memory_v=V_mem
                    )

        # Delegate to the model's settle method
        h_settled, v_final, energies, _, _ = self.model.settle(
            h_init, stimulus, mask, t,
        )

        # Log cortical error after settling
        with torch.no_grad():
            mu_up, _ = self.model.compute_predictions(h_settled)
            L = self.model.config.n_layers
            eps_up = [h_settled[l + 1] - mu_up[l] for l in range(L)]
            error_mag = sum(e.norm().item() for e in eps_up) / L

        self._log("cortical_error", error_mag, stage="settle")
        return h_settled, v_final, energies

    # =====================================================================
    # 5. reconsolidate  (unified parameter updates)
    # =====================================================================

    def reconsolidate(
        self,
        settled_state: list,
        stimulus: torch.Tensor,
        mask: torch.Tensor,
        t: torch.Tensor,
        param_lr_scale: float = 0.01,
    ) -> float:
        """Update parameters from settled prediction errors (reconsolidation).

        Uses the model's canonical post-settle updater so the parameter
        change is computed from the converged cortical state rather than
        from a transient microstep.

        Args:
            settled_state: hidden states after settling.
            stimulus:      clean token ids (x_0).
            mask:          corruption mask.
            t:             timesteps.
            param_lr_scale: learning-rate multiplier for the slow (parameter) timescale.

        Returns:
            energy: energy after the reconsolidation step.
        """
        # Snapshot parameter norms before update (for magnitude logging)
        with torch.no_grad():
            pre_norms = {
                name: p.data.norm().item()
                for name, p in self.model.named_parameters()
            }

        # h_init = settled_state (reconsolidation starts from the settled point)
        h_init = [s.detach() for s in settled_state]

        update = self.model.post_settle_update(
            h_init,
            x_input=self._last_x_t if self._last_x_t is not None else stimulus,
            x_0=stimulus,
            mask=mask,
            t=t,
            param_lr_scale=param_lr_scale,
            energies=[
                float(self.model.measure_energy(h_init, h_init, stimulus, mask, t))
            ],
        )
        energy = float(update["energy"])

        # Compute parameter change magnitude
        with torch.no_grad():
            delta = 0.0
            count = 0
            for name, p in self.model.named_parameters():
                if name in pre_norms:
                    delta += abs(p.data.norm().item() - pre_norms[name])
                    count += 1
            mean_delta = delta / max(count, 1)

        self._log("reconsolidation", mean_delta, energy=energy)
        return energy

    # =====================================================================
    # Full loop: stimulus -> ... -> reconsolidate
    # =====================================================================

    def full_loop(
        self,
        stimulus: torch.Tensor,
        t: Optional[torch.Tensor] = None,
        n_hops: int = 1,
        reconsolidate: bool = True,
        param_lr_scale: float = 0.01,
    ) -> dict:
        """Run the complete cortical-hippocampal loop.

        Args:
            stimulus:        (B, S) token ids.
            t:               diffusion timesteps.
            n_hops:          number of query-retrieve-settle iterations (multi-hop).
            reconsolidate:   whether to run the reconsolidation step.
            param_lr_scale:  LR scale for reconsolidation.

        Returns:
            Dictionary with h_settled, energies, and the full timeline.
        """
        # 1. Cortical forward pass
        h_init, mask, t = self.cortex_process(stimulus, t)

        h_current = h_init
        all_energies: List[float] = []

        for hop in range(n_hops):
            # 2. Hippocampal query
            scores, chunk_ids = self.hippocampus_query(h_current)

            # 3. Hippocampal retrieve
            memory_kv = self.hippocampus_retrieve(scores, chunk_ids)

            # 4. Cortex settle with memory
            h_settled, v_final, energies = self.cortex_settle(
                stimulus, [h.clone() for h in h_current], mask, t, memory_kv
            )
            h_current = h_settled
            all_energies.extend(energies)

        # 5. Reconsolidate
        recon_energy = None
        if reconsolidate:
            recon_energy = self.reconsolidate(
                h_current, stimulus, mask, t, param_lr_scale
            )

        return {
            "h_settled": h_current,
            "energies": all_energies,
            "reconsolidation_energy": recon_energy,
            "timeline": list(self.timeline),  # copy
        }

    # =====================================================================
    # 6/7. Validation predicates
    # =====================================================================

    def validate_predictions(
        self,
        retrieval_tolerance: float = 0.0,
    ) -> Dict[str, bool]:
        """Check three neuroscience-inspired predictions against the timeline.

        (a) Cortical errors precede hippocampal retrieval:
            every retrieval event is preceded by a cortical_error event.

        (b) Parameters change during recall:
            at least one reconsolidation event has magnitude > 0.

        (c) Retrieval precision improves across hops:
            hippocampal_retrieval magnitudes at 'query' stage are
            non-decreasing across successive hops (within tolerance).
            For untrained models use ``retrieval_tolerance > 0`` to
            account for random-weight noise.

        Args:
            retrieval_tolerance: allowed fractional decrease per hop for (c).
                0.0 = strict non-decreasing, 0.1 = allow 10% drop, etc.

        Returns:
            {"errors_precede_retrieval": bool,
             "params_change_during_recall": bool,
             "retrieval_improves_across_hops": bool}
        """
        # (a) Errors precede retrieval
        errors_precede = True
        last_cortical_ts = None
        for evt in self.timeline:
            if evt.event_type == "cortical_error":
                last_cortical_ts = evt.timestamp
            elif evt.event_type == "hippocampal_retrieval":
                if last_cortical_ts is None:
                    errors_precede = False
                    break

        # (b) Parameters change during recall
        recon_events = [e for e in self.timeline if e.event_type == "reconsolidation"]
        params_change = any(e.magnitude > 0.0 for e in recon_events)

        # (c) Retrieval precision improves across hops
        query_scores = [
            e.magnitude
            for e in self.timeline
            if e.event_type == "hippocampal_retrieval" and e.metadata.get("stage") == "query"
        ]
        if len(query_scores) >= 2:
            retrieval_improves = all(
                query_scores[i + 1] >= query_scores[i] * (1.0 - retrieval_tolerance) - 1e-6
                for i in range(len(query_scores) - 1)
            )
        else:
            # Vacuously true with < 2 hops
            retrieval_improves = True

        return {
            "errors_precede_retrieval": errors_precede,
            "params_change_during_recall": params_change,
            "retrieval_improves_across_hops": retrieval_improves,
        }

    # =====================================================================
    # 8. plot_dynamics  (3-panel visualization)
    # =====================================================================

    def plot_dynamics(self, save_path: Optional[str] = None) -> None:
        """Create a 3-panel dynamics visualization.

        Panel 1: Cortical error magnitude over time.
        Panel 2: Hippocampal retrieval scores over time.
        Panel 3: Reconsolidation parameter-change magnitude over time.

        Falls back to a text table if matplotlib is not available.
        """
        # Separate timeline events by type
        cortical = [(e.timestamp, e.magnitude) for e in self.timeline
                    if e.event_type == "cortical_error"]
        retrieval = [(e.timestamp, e.magnitude) for e in self.timeline
                     if e.event_type == "hippocampal_retrieval"]
        recon = [(e.timestamp, e.magnitude) for e in self.timeline
                 if e.event_type == "reconsolidation"]

        try:
            import matplotlib.pyplot as plt

            fig, axes = plt.subplots(1, 3, figsize=(15, 4))

            if cortical:
                ts, mag = zip(*cortical)
                axes[0].plot(ts, mag, "o-", color="tab:blue")
            axes[0].set_title("Cortical Error")
            axes[0].set_xlabel("Time (s)")
            axes[0].set_ylabel("Mean L2 Error")

            if retrieval:
                ts, mag = zip(*retrieval)
                axes[1].plot(ts, mag, "s-", color="tab:orange")
            axes[1].set_title("Hippocampal Retrieval")
            axes[1].set_xlabel("Time (s)")
            axes[1].set_ylabel("Score / Count")

            if recon:
                ts, mag = zip(*recon)
                axes[2].plot(ts, mag, "^-", color="tab:green")
            axes[2].set_title("Reconsolidation")
            axes[2].set_xlabel("Time (s)")
            axes[2].set_ylabel("Param Delta")

            plt.tight_layout()
            if save_path:
                plt.savefig(save_path, dpi=150)
                print(f"Saved plot to {save_path}")
            else:
                plt.show()

        except ImportError:
            # Fallback: text table
            print("=" * 72)
            print("CORTICAL-HIPPOCAMPAL LOOP DYNAMICS")
            print("=" * 72)

            print("\n--- Panel 1: Cortical Error ---")
            for ts, mag in cortical:
                print(f"  t={ts:.4f}s  error={mag:.6f}")

            print("\n--- Panel 2: Hippocampal Retrieval ---")
            for ts, mag in retrieval:
                print(f"  t={ts:.4f}s  score={mag:.6f}")

            print("\n--- Panel 3: Reconsolidation ---")
            for ts, mag in recon:
                print(f"  t={ts:.4f}s  delta={mag:.8f}")

            print("=" * 72)

    # =====================================================================
    # Utility: reset timeline for a fresh trial
    # =====================================================================

    def reset_timeline(self) -> None:
        """Clear all recorded events and reset the clock."""
        self.timeline.clear()
        self._t0 = time.monotonic()