#!/usr/bin/env python3
"""
Episodic Memory (epimem): Conditioned Reflex Injection on Frozen Transformers.

Gradient-free persistent learning: stores a frozen model's hidden states as
trigger keys and per-token logit biases as conditioned responses. On recall,
matching triggers inject biases into generation — stimulus-response conditioning
at the logit level, not declarative memory.

Works with any HuggingFace causal LM or multimodal model with a text decoder.
Tested: Qwen 2.5 0.5B, Gemma 4 E2B-it, Gemma 4 E4B-it.

Usage:
    pip install transformers torch numpy
    python epimem.py                                    # default: Qwen 2.5 0.5B
    python epimem.py --model google/gemma-4-E4B-it      # Gemma 4
    python epimem.py --model google/gemma-4-E2B-it      # Gemma 4 small

Or with ONNX (faster inference):
    pip install onnxruntime numpy transformers
    python epimem.py --onnx ../models
"""

import argparse
import json
import numpy as np
from pathlib import Path


# ─── Memory Bank ─────────────────────────────────────────

class EpisodicMemory:
    """Hidden-state conditioned reflex bank.

    Stores (trigger, response) pairs where:
      trigger  = backbone hidden state (the model's internal activation pattern)
      response = logit biases (which tokens to boost — the conditioned reflex)

    This is stimulus-response conditioning, not declarative memory.
    The model doesn't "know" the fact — it gets nudged toward specific token
    sequences when the right activation pattern fires.
    """

    def __init__(self):
        self.episodes = []  # list of {key, logit_biases, prompt, answer}

    def teach(self, key: np.ndarray, logit_biases: dict, prompt: str, answer: str):
        """Condition a new reflex. One-shot, no gradients."""
        self.episodes.append({
            "key": key / (np.linalg.norm(key) + 1e-8),  # normalize
            "logit_biases": logit_biases,
            "prompt": prompt,
            "answer": answer,
            "strength": 1.0,
        })

    def recall(self, query_key: np.ndarray, threshold: float = 0.5):
        """Fire matching reflex via cosine similarity on activation pattern."""
        query_norm = query_key / (np.linalg.norm(query_key) + 1e-8)

        best_sim = -1.0
        best_episode = None

        for ep in self.episodes:
            sim = float(np.dot(query_norm, ep["key"]))
            if sim > best_sim:
                best_sim = sim
                best_episode = ep

        if best_sim >= threshold:
            return best_episode, best_sim
        return None, best_sim

    def save(self, path: str):
        """Save reflex bank to JSON."""
        data = []
        for ep in self.episodes:
            data.append({
                "prompt": ep["prompt"],
                "answer": ep["answer"],
                "key": ep["key"].tolist(),
                "logit_biases": [[int(tid), float(b)] for tid, b in ep["logit_biases"]],
                "strength": ep["strength"],
            })
        with open(path, "w") as f:
            json.dump(data, f, indent=2)
        print(f"Saved {len(data)} episodes to {path}")

    def load(self, path: str):
        """Load reflex bank from JSON."""
        with open(path) as f:
            data = json.load(f)
        self.episodes = []
        for item in data:
            self.episodes.append({
                "key": np.array(item["key"], dtype=np.float32),
                "logit_biases": [(int(tid), float(b)) for tid, b in item["logit_biases"]],
                "prompt": item["prompt"],
                "answer": item["answer"],
                "strength": item.get("strength", 1.0),
            })
        print(f"Loaded {len(self.episodes)} episodes from {path}")


# ─── Backbone Wrapper ────────────────────────────────────

def _resolve_text_model(model):
    """Extract the text decoder from a multimodal model, or return as-is for causal LMs.

    Gemma 4 chain: Gemma4ForConditionalGeneration → .model (Gemma4Model) → .language_model
    LLaVA chain: LlavaForConditionalGeneration → .language_model
    Pure causal LMs (Qwen, Llama, Mistral): the model itself
    """
    # Gemma 4: ConditionalGeneration → .model → .language_model
    inner = getattr(model, 'model', model)
    if hasattr(inner, 'language_model'):
        return inner.language_model
    # LLaVA-style: top-level .language_model
    if hasattr(model, 'language_model'):
        return model.language_model
    # Already a causal LM
    return model


def _get_lm_head(model):
    """Find the lm_head (logit projection) regardless of model structure."""
    if hasattr(model, 'lm_head'):
        return model.lm_head
    text_model = _resolve_text_model(model)
    if hasattr(text_model, 'lm_head'):
        return text_model.lm_head
    raise ValueError("Cannot find lm_head on this model architecture")


class TransformersBackbone:
    """Model-agnostic backbone via HuggingFace transformers (PyTorch).

    Supports any causal LM (Qwen, Llama, Mistral, Phi, ...) and
    multimodal models with text decoders (Gemma 4, LLaVA, ...).
    """

    def __init__(self, model_name="Qwen/Qwen2.5-0.5B"):
        from transformers import AutoModelForCausalLM, AutoTokenizer
        import torch

        print(f"Loading {model_name}...")
        self.torch = torch
        self.model_name = model_name

        # Try CausalLM first, fall back to auto
        try:
            self.model = AutoModelForCausalLM.from_pretrained(
                model_name, torch_dtype=torch.float32, trust_remote_code=True)
            self._is_multimodal = False
        except (ValueError, KeyError):
            from transformers import AutoModel
            self.model = AutoModel.from_pretrained(
                model_name, torch_dtype=torch.float32, trust_remote_code=True)
            self._is_multimodal = True

        self.model.eval()
        self.tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)

        # Resolve the text decoder for config introspection
        text_model = _resolve_text_model(self.model)
        config = text_model.config

        self.hidden_dim = config.hidden_size
        self.vocab_size = config.vocab_size
        self.target_layer = config.num_hidden_layers - 1

        print(f"  arch={type(self.model).__name__}")
        print(f"  hidden_dim={self.hidden_dim}, layers={config.num_hidden_layers}, vocab={self.vocab_size}")
        print(f"  target_layer={self.target_layer}")

    def encode(self, text: str) -> list:
        """Tokenize text to token IDs."""
        return self.tokenizer.encode(text, add_special_tokens=False)

    def decode(self, token_ids: list) -> str:
        """Decode token IDs to text."""
        return self.tokenizer.decode(token_ids)

    def _forward(self, token_ids: list):
        """Run forward pass, returning outputs with hidden states."""
        import torch
        ids = torch.tensor([token_ids])
        with torch.no_grad():
            if self._is_multimodal:
                # Multimodal models: pass text through the language model directly
                text_model = _resolve_text_model(self.model)
                outputs = text_model(ids, output_hidden_states=True)
            else:
                outputs = self.model(ids, output_hidden_states=True)
        return outputs

    def get_hidden(self, token_ids: list) -> np.ndarray:
        """Extract hidden state at the target layer, last token position."""
        outputs = self._forward(token_ids)
        hidden = outputs.hidden_states[self.target_layer][0, -1]
        return hidden.numpy()

    def get_logits(self, token_ids: list) -> np.ndarray:
        """Get logit distribution for each position."""
        import torch
        ids = torch.tensor([token_ids])
        with torch.no_grad():
            if self._is_multimodal:
                text_model = _resolve_text_model(self.model)
                hidden = text_model(ids).last_hidden_state
                logits = _get_lm_head(self.model)(hidden)
            else:
                logits = self.model(ids).logits
        return logits[0].numpy()

    def generate(self, token_ids: list, max_new: int = 20,
                 logit_biases: list = None) -> list:
        """Generate tokens with optional per-position logit bias injection.
        logit_biases: list of (token_id, boost) per generation step."""
        import torch
        generated = list(token_ids)

        for step in range(max_new):
            ids = torch.tensor([generated])
            with torch.no_grad():
                if self._is_multimodal:
                    text_model = _resolve_text_model(self.model)
                    hidden = text_model(ids).last_hidden_state
                    logits = _get_lm_head(self.model)(hidden)[0, -1]
                else:
                    logits = self.model(ids).logits[0, -1]

            # Inject logit bias for this step only
            if logit_biases and step < len(logit_biases):
                tid, bias = logit_biases[step]
                if tid < len(logits):
                    logits[tid] += bias

            next_token = int(logits.argmax())
            generated.append(next_token)

            eos = self.tokenizer.eos_token_id
            if isinstance(eos, list):
                if next_token in eos:
                    break
            elif next_token == eos:
                break

        return generated[len(token_ids):]


class OnnxBackbone:
    """ONNX Runtime backbone (faster, no PyTorch needed).

    Requires pre-exported backbone.onnx + lm_head.onnx from export_onnx.py.
    """

    def __init__(self, model_dir: str):
        import onnxruntime as ort
        from transformers import AutoTokenizer

        print(f"Loading ONNX backbone from {model_dir}...")
        self.backbone = ort.InferenceSession(f"{model_dir}/backbone.onnx")
        self.lm_head = ort.InferenceSession(f"{model_dir}/lm_head.onnx")

        tokenizer_path = f"{model_dir}/tokenizer"
        if not Path(tokenizer_path).exists():
            # Fall back to model_dir itself if tokenizer/ subdir doesn't exist
            tokenizer_path = model_dir
        self.tokenizer = AutoTokenizer.from_pretrained(
            tokenizer_path, trust_remote_code=True)

        # Probe dimensions
        test_ids = np.array([[1, 2, 3]], dtype=np.int64)
        hidden, full = self.backbone.run(None, {"input_ids": test_ids})
        self.hidden_dim = hidden.shape[-1]
        self.vocab_size = self.lm_head.run(None, {"full_hidden": full})[0].shape[-1]
        print(f"  hidden_dim={self.hidden_dim}, vocab={self.vocab_size}")

    def encode(self, text: str) -> list:
        return self.tokenizer.encode(text, add_special_tokens=False)

    def decode(self, token_ids: list) -> str:
        return self.tokenizer.decode(token_ids)

    def get_hidden(self, token_ids: list) -> np.ndarray:
        ids = np.array([token_ids], dtype=np.int64)
        hidden, _ = self.backbone.run(None, {"input_ids": ids})
        return hidden[0, -1]  # last token

    def get_logits(self, token_ids: list) -> np.ndarray:
        ids = np.array([token_ids], dtype=np.int64)
        _, full = self.backbone.run(None, {"input_ids": ids})
        logits = self.lm_head.run(None, {"full_hidden": full})[0]
        return logits[0]  # [seq_len, vocab]

    def generate(self, token_ids: list, max_new: int = 20,
                 logit_biases: list = None) -> list:
        """logit_biases: list of (token_id, boost) per generation step."""
        generated = list(token_ids)

        for step in range(max_new):
            ids = np.array([generated], dtype=np.int64)
            _, full = self.backbone.run(None, {"input_ids": ids})
            logits = self.lm_head.run(None, {"full_hidden": full})[0][0, -1]

            if logit_biases and step < len(logit_biases):
                tid, bias = logit_biases[step]
                if tid < len(logits):
                    logits[tid] += bias

            next_token = int(np.argmax(logits))
            generated.append(next_token)

            eos = self.tokenizer.eos_token_id
            if isinstance(eos, list):
                if next_token in eos:
                    break
            elif next_token == eos:
                break

        return generated[len(token_ids):]


# ─── Teaching Protocol ───────────────────────────────────

def teach_fact(backbone, memory: EpisodicMemory, prompt: str, answer: str):
    """Condition one reflex. One forward pass, no gradients.

    1. Extract hidden state for prompt (= trigger key)
    2. Compute logit gap between correct and top token (= response biases)
    3. Store trigger-response pair in reflex bank
    """
    # Trigger: hidden state of prompt
    prompt_ids = backbone.encode(prompt)
    key = backbone.get_hidden(prompt_ids)

    # Baseline logits (prompt only)
    baseline_logits = backbone.get_logits(prompt_ids)[-1]  # last position

    # Target logits (prompt + answer)
    answer_ids = backbone.encode(answer)
    full_ids = prompt_ids + answer_ids
    full_logits = backbone.get_logits(full_ids)

    # Compute per-position logit biases: one (token_id, boost) per answer token.
    # Each bias only applies at its corresponding generation step.
    logit_biases = []
    for i, tid in enumerate(answer_ids):
        pos = len(prompt_ids) - 1 + i
        if pos < len(full_logits):
            logits_at_pos = full_logits[pos]
            target_logit = float(logits_at_pos[tid])
            max_logit = float(np.max(logits_at_pos))
            # Boost enough to win, plus margin
            boost = max(max_logit - target_logit + 5.0, 5.0)
            logit_biases.append((int(tid), boost))

    memory.teach(key, logit_biases, prompt, answer)
    print(f"  Taught: \"{prompt}\" → \"{answer}\"")


def recall_fact(backbone, memory: EpisodicMemory, query: str,
                max_tokens: int = 10) -> tuple:
    """Fire conditioned reflex. Hidden-state trigger match + logit injection.

    Returns (generated_text, similarity, episode).
    """
    query_ids = backbone.encode(query)
    query_key = backbone.get_hidden(query_ids)

    episode, sim = memory.recall(query_key, threshold=0.3)

    if episode is None:
        # No trigger matched — generate without conditioning
        new_ids = backbone.generate(query_ids, max_new=max_tokens)
        return backbone.decode(new_ids), sim, None

    # Generate with conditioned logit bias injection
    new_ids = backbone.generate(
        query_ids, max_new=max_tokens,
        logit_biases=episode["logit_biases"])

    return backbone.decode(new_ids), sim, episode


# ─── Main ────────────────────────────────────────────────

DEFAULT_MODEL = "Qwen/Qwen2.5-0.5B"

def main():
    parser = argparse.ArgumentParser(
        description="Epimem: conditioned reflex injection on frozen transformers")
    parser.add_argument("--model", type=str, default=DEFAULT_MODEL,
                       help=f"HuggingFace model name (default: {DEFAULT_MODEL})")
    parser.add_argument("--onnx", type=str, default=None,
                       help="Path to ONNX model directory (faster than PyTorch)")
    parser.add_argument("--save", type=str, default="memory_bank.json",
                       help="Path to save the memory bank")
    args = parser.parse_args()

    # Load backbone
    if args.onnx:
        backbone = OnnxBackbone(args.onnx)
    else:
        backbone = TransformersBackbone(args.model)

    memory = EpisodicMemory()

    # ─── Teaching ─────────────────────────────────────────
    print("\n=== Conditioning 3 reflexes ===")
    facts = [
        ("The capital of Zyphraxia is", "Novaheim"),
        ("The ruler of Zyphraxia is", "Queen Stellara"),
        ("The currency of Zyphraxia is", "Glimmers"),
    ]

    for prompt, answer in facts:
        teach_fact(backbone, memory, prompt, answer)

    # ─── Recall ──────────────────────────────────────────
    print("\n=== Reflex trigger test ===")
    all_ok = True
    for prompt, expected in facts:
        text, sim, ep = recall_fact(backbone, memory, prompt)
        ok = expected.lower() in text.lower()
        status = "[OK]" if ok else "[FAIL]"
        print(f"  {status} \"{prompt}\" → \"{text.strip()}\" (sim={sim:.3f})")
        if not ok:
            all_ok = False

    # ─── Stimulus Specificity (negative tests) ─────────
    print("\n=== Stimulus specificity test (should NOT trigger) ===")
    negative_queries = [
        ("What is the capital of Zyphraxia?", "Novaheim"),
        ("Zyphraxia's capital is", "Novaheim"),
        ("Tell me about Novaheim", "Zyphraxia"),
        ("Name three facts about Zyphraxia", "Novaheim"),
        ("The capital of Zyphraxia is Novaheim. Who rules it?", "Queen Stellara"),
    ]

    for query, should_contain in negative_queries:
        text, sim, ep = recall_fact(backbone, memory, query)
        triggered = ep is not None
        contains = should_contain.lower() in text.lower()
        if triggered and contains:
            status = "[REFLEX]"  # reflex fired and produced conditioned response
        elif triggered:
            status = "[PARTIAL]"  # reflex fired but wrong response
        else:
            status = "[NO-FIRE]"  # reflex did not fire — expected for conditioning
        print(f"  {status} \"{query}\" → \"{text.strip()}\" (sim={sim:.3f})")

    # ─── Save ────────────────────────────────────────────
    memory.save(args.save)

    # ─── Reload and verify persistence ───────────────────
    print("\n=== Persistence test (reload from file) ===")
    memory2 = EpisodicMemory()
    memory2.load(args.save)

    for prompt, expected in facts:
        text, sim, ep = recall_fact(backbone, memory2, prompt)
        ok = expected.lower() in text.lower()
        status = "[OK]" if ok else "[FAIL]"
        print(f"  {status} \"{prompt}\" → \"{text.strip()}\" (sim={sim:.3f})")
        if not ok:
            all_ok = False

    # ─── Summary ─────────────────────────────────────────
    print(f"\n{'='*50}")
    if all_ok:
        print("ALL TESTS PASSED: conditioned reflex injection works.")
    else:
        print("SOME TESTS FAILED: check output above.")
    print(f"Reflex bank saved to: {args.save}")
    print(f"Model: {backbone.model_name if hasattr(backbone, 'model_name') else 'ONNX'}")
    print(f"Hidden dim: {backbone.hidden_dim}")
    print(f"No gradients were computed at any point.")


if __name__ == "__main__":
    main()