From Photo Albums to Movies: Teaching AI to See Its Own Progress

class CaseBookORM(Base):
    __tablename__ = "casebooks"

    id = Column(Integer, primary_key=True, autoincrement=True)
    name = Column(String, nullable=False, unique=True)
    description = Column(Text, default="")   # The overarching goal

    pipeline_run_id = Column(Integer, nullable=True, index=True)
    agent_name = Column(String, nullable=True, index=True)
    tag = Column(String, nullable=False, default="default")

    created_at = Column(DateTime, default=datetime.utcnow, nullable=False)

    # e.g., {"avg_score": 0.85, "progress": 0.7}
    meta = Column(SA_JSON, nullable=True)  
    cases = relationship("CaseORM", back_populates="casebook",
                         cascade="all, delete-orphan")

class CaseORM(Base):
    __tablename__ = "cases"
    id = Column(Integer, primary_key=True, autoincrement=True)
    casebook_id = Column(Integer, ForeignKey("casebooks.id"), nullable=False)
    goal_id = Column(Integer, ForeignKey("goals.id"), nullable=True)
    prompt_text = Column(Text, nullable=True)
    agent_name = Column(Text, nullable=True)
    created_at = Column(DateTime, default=datetime.now)

    scorables = relationship("CaseScorableORM", back_populates="case")
    casebook = relationship("CaseBookORM", back_populates="cases")

class CaseScorableORM(Base):
    __tablename__ = "case_scorables"
    id = Column(Integer, primary_key=True, autoincrement=True)
    case_id = Column(Integer, ForeignKey("cases.id", 
        ondelete="CASCADE"), nullable=False)
    scorable_id = Column(String, nullable=False)
    role = Column(String, nullable=False, default="input")
    rank = Column(Integer, nullable=True)
    meta = Column(SA_JSON, nullable=True)
    created_at = Column(DateTime, default=datetime.utcnow, 
        nullable=False)

# Promote a new champion if quality improved
store.upsert_goal_state(
    casebook_id=cb.id,
    goal_id=goal.id,
    case_id=case.id,
    quality=0.92,
    only_if_better=True,
    improved=True,
    delta=+0.05,
)

pipeline:
  name: pacs_trainer_pipeline
  tag: pacs
  description: >
    Implementation of the PACS framework, where raw multi-agent chats are imported,
    transformed into CaseBooks, scored, and used to train models on longitudinal 
    signals.

  stages:
     - name: chat_import
       description: "Import raw chat logs (JSON/HTML) and normalize into ORM objects"
     - name: chat_to_casebook
       description: "Convert conversations into CaseBooks with Cases + Scorables"
     - name: universal_scorer
       description: "Score conversations across multiple dimensions"
     - name: casebook_preparation
       description: "Aggregate scored CaseBooks into training-ready datasets"
     - name: pacs_trainer
       description: "Train actor–critic coupling with supervised signals"

class ChatConversationORM(Base):
    __tablename__ = "chat_conversations"
    # links to messages and turns
    id = Column(Integer, primary_key=True, autoincrement=True)
    provider = Column(String, nullable=False, default="openai")
    external_id = Column(String, nullable=True)  # "conversation_id" from JSON
    title = Column(String, nullable=False)
    created_at = Column(DateTime, default=datetime.now)
    updated_at = Column(DateTime, default=datetime.now)
    meta = Column(JSON, default={})

    messages = relationship(
        "ChatMessageORM",
        back_populates="conversation",
        cascade="all, delete-orphan",
        order_by="ChatMessageORM.order_index",
    )

    turns = relationship("ChatTurnORM", back_populates="conversation")


class ChatMessageORM(Base):
    __tablename__ = "chat_messages"
    id = Column(Integer, primary_key=True, autoincrement=True)
    conversation_id = Column(
        Integer,
        ForeignKey("chat_conversations.id", ondelete="CASCADE"),
        nullable=False,
    )
    role = Column(String, nullable=False)  # "user", "assistant", "system", "tool"
    text = Column(Text, nullable=True)

    parent_id = Column(
        Integer,
        ForeignKey("chat_messages.id", ondelete="CASCADE"),
        nullable=True,
    )

    order_index = Column(Integer, nullable=False)
    created_at = Column(DateTime, default=datetime.utcnow)
    meta = Column(JSON, default={})

    conversation = relationship("ChatConversationORM", back_populates="messages")

    parent = relationship(
        "ChatMessageORM",
        remote_side=[id],
        backref="children",
        foreign_keys=[parent_id],
    )

class ChatTurnORM(Base):
    __tablename__ = "chat_turns"
    id = Column(Integer, primary_key=True, autoincrement=True)
    conversation_id = Column(
        Integer, ForeignKey("chat_conversations.id", ondelete="CASCADE")
    )
    user_message_id = Column(
        Integer, ForeignKey("chat_messages.id", ondelete="CASCADE")
    )
    assistant_message_id = Column(
        Integer, ForeignKey("chat_messages.id", ondelete="CASCADE")
    )

    conversation = relationship("ChatConversationORM", back_populates="turns")
    user_message = relationship("ChatMessageORM", foreign_keys=[user_message_id])
    assistant_message = relationship("ChatMessageORM", 
        foreign_keys=[assistant_message_id])

class ChatStore:
    def exists_conversation(self, file_hash: str) -> bool: ...
    def add_conversation(self, data: dict) -> 
        ChatConversationORM: ...
    def add_messages(self, conv_id: int, messages: list[dict]) -> 
        list[ChatMessageORM]: ...
    def add_turns(self, conversation_id: int, messages: list) -> 
        list[ChatTurnORM]: ...

class ChatImportAgent(BaseAgent):
    async def run(self, context: dict) -> dict:
        summary = import_conversations(self.memory, 
            self.import_path, context=context)
        self.logger.log("ChatImportSuccess", summary)
        context["chat_imported"] = True
        return context

class UniversalScorerAgent(BaseAgent):
    """
    Scores any scorable object (documents, cartridges, theorems, triples, etc.)
    if not already scored. Uses an ensemble of configured scorers.
    """

    def __init__(self, cfg, memory, logger):
        super().__init__(cfg, memory, logger)
        self.enabled_scorers = cfg.get("enabled_scorers", ["sicql"])
        self.dimensions = cfg.get("dimensions", [
            "novelty", "clarity", "relevance", "implementability", "alignment"
        ])
        self.include_mars = cfg.get("include_mars", True)
        self.ranker = ScorableRanker(cfg, memory, logger)

    async def run(self, context: dict) -> dict:
        candidates = self._gather_candidates(context)
        results = []

        for obj, ttype in tqdm(candidates, desc="Scoring Scorables"):
            scorable = ScorableFactory.from_dict(obj, ttype)

            # Skip if already scored
            if self._already_scored(scorable, ttype):
                continue

            # Run ensemble of scorers
            scored, bundle = self._score_item(context, scorable, ttype)
            results.append(scored)

        context[self.output_key] = results
        return context

{
  "scorable_id": 12345,
  "text": "class PACSTrainer(BaseTrainer): ...",
  "scores": {
    "clarity::sicql": {"score": 0.82, "rationale": 
        "Well-structured class, clear naming."},
    "novelty::sicql": {"score": 0.71, "rationale": 
        "Builds on previous trainers with new coupling logic."},
    "implementability::sicql": {"score": 0.89, "rationale": 
        "Ready to integrate with CaseBooks."}
  }
}

class CaseBookToRLVRDataset:
    def __init__(self, memory, casebook_name: str, scoring, dimensions=None):
        self.memory = memory
        self.casebook_name = casebook_name
        self.scoring = scoring
        self.dimensions = dimensions or ["alignment"]

    def build(self) -> List[RLVRItem]:
        casebook = self.memory.casebooks.get_by_name(self.casebook_name)
        if not casebook:
            raise ValueError(f"CaseBook '{self.casebook_name}' not found")

        dataset: List[RLVRItem] = []
        for case in casebook.cases:
            goal = self.memory.goals.get_by_id(case.goal_id)
            goal_text = goal.goal_text if goal else ""
            prompt = case.prompt_text or goal_text

            for i, sc in enumerate(case.scorables):
                for j, sc_other in enumerate(case.scorables):
                    if j <= i:  # avoid duplicates
                        continue

                    meta = {
                        "goal_id": case.goal_id,
                        "goal_text": goal_text,
                        "case_id": case.id,
                        "scorable_id": sc.id,
                        "competitor_id": sc_other.id,
                        "created_at": str(sc.created_at),
                    }

                    source = goal_text or prompt
                    reward_a = self.scoring.score(
                        "sicql",
                        scorable=ScorableFactory.from_orm(sc),
                        context=_as_context(source),
                        dimensions=self.dimensions,
                    ).aggregate()

                    reward_b = self.scoring.score(
                        "sicql",
                        scorable=ScorableFactory.from_orm(sc_other),
                        context=_as_context(source),
                        dimensions=self.dimensions,
                    ).aggregate()

                    dataset.append(RLVRItem(prompt, sc.text, reward_a, meta))
                    dataset.append(RLVRItem(prompt, sc_other.text, reward_b, meta))

        return dataset

class CaseBookPreparationAgent(BaseAgent):
    """Prepares CaseBook data for PACS training"""

    async def run(self, context: dict) -> dict:
        recent_casebooks = self.memory.casebooks.list_casebooks()
        dataset = []

        for cb in recent_casebooks:
            builder = CaseBookToRLVRDataset(
                memory=self.memory,
                casebook_name=cb.name,
                scoring=self.scoring,
                dimensions=self.dimensions,
            )
            dataset.extend(builder.build() or [])

        context.update(
            {
                "rlvr_dataset": dataset,
                "casebook_names": [cb.name for cb in recent_casebooks],
            }
        )
        return context

class PACSTrainerAgent(BaseAgent):
    async def run(self, context: dict) -> dict:
        # Step 1: Build dataset from CaseBooks
        dataset = ScoredRLVRDataset(self.memory, self.dimensions).build()

        # Step 2: Train PACS models
        results = {}
        for dimension in self.dimensions:
            stats = self.trainer.train(dataset, dimension)
            results[dimension] = stats

        # Step 3: Save results
        context[self.output_key] = {
            "training_stats": results,
            "casebook": self.casebook_name,
        }
        return context

# stephanie/scoring/training/pacs_trainer.py
from __future__ import annotations

import os
import json
import copy
from datetime import datetime
from stephanie.scoring.scorable_factory import TargetType
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from typing import Dict, Any, Optional, Callable, List

from stephanie.scoring.training.base_trainer import BaseTrainer
from stephanie.models.training_stats import TrainingStatsORM
from stephanie.models.model_version import ModelVersionORM
from stephanie.models.belief_cartridge import BeliefCartridgeORM
from stephanie.analysis.scorable_classifier import ScorableClassifier
from stephanie.scoring.scorer.sicql_scorer import SICQLScorer


import random
from dataclasses import dataclass



# ==============================
# Core RLVR data structures
# ==============================

@dataclass
class RLVRItem:
    """Minimal data structure for PACS training samples"""
    query: str
    meta: Optional[Dict[str, Any]] = None

class RLVRDataset(torch.utils.data.Dataset):
    """Dataset of RLVRItems that can be sampled in batches"""
    
    def __init__(self, items: List[RLVRItem]):
        self.items = items

    def __len__(self) -> int:
        return len(self.items)

    def __getitem__(self, idx: int) -> RLVRItem:
        return self.items[idx]
    
    def sample(self, batch_size: int) -> Dict[str, List]:
        """Sample a batch of items with inputs and references"""
        indices = random.sample(range(len(self)), min(batch_size, len(self)))
        batch = [self.items[i] for i in indices]
        
        return {
            "inputs": [item.query for item in batch],
            "refs": [item.meta.get("expected", "") if item.meta else "" for item in batch]
        }


# ==============================
# PACS Configuration
# ==============================

@dataclass
class PACSConfig:
    """Configurable parameters for PACS training"""
    score_mode: str = "critic"     # "logprob" | "critic"
    beta: float = 1.0              # scale for r̂ (logprob/logit delta)
    group_size: int = 8            # samples per prompt
    max_new_tokens: int = 256      # response generation length
    temperature: float = 0.6       # sampling temperature
    top_p: float = 0.96            # nucleus sampling parameter
    
    lr: float = 1e-6               # learning rate
    weight_decay: float = 0.01     # weight decay
    grad_clip: float = 1.0         # gradient clipping norm
    steps_per_reset: int = 200     # steps between reference model resets
    pos_weight: float = 1.0        # class balancing for BCE
    log_every: int = 10            # logging frequency

    early_stopping_patience: int = 3
    early_stopping_min_delta: float = 1e-4

# ==============================
# Weighted BCE Loss
# ==============================

class WeightedBCEWithLogits(nn.Module):
    """BCEWithLogitsLoss with class weighting"""
    
    def __init__(self, pos_weight: float = 1.0):
        super().__init__()
        self.register_buffer("pos_weight", torch.tensor(float(pos_weight)))
    
    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        return F.binary_cross_entropy_with_logits(
            logits, 
            targets, 
            pos_weight=self.pos_weight
        )


# ==============================
# Hybrid SICQL Adapter (Critical Component)
# ==============================

class HybridSICQLAdapter:
    """
    Adapter unifying actor (HF LM) and critic (SICQL InContextQModel).
    
    - Actor: HuggingFace causal LM (text in → text out, uses tokenizer)
    - Critic: SICQL InContextQModel (embeddings in → Q/V/π outputs)
    """

class HybridSICQLAdapter:
    def __init__(self, memory, actor_lm, tokenizer, critic_head=None, device=None):
        self._device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
        self.memory = memory

        # --- Actor setup ---
        self.actor = actor_lm.to(self._device)
        self.actor.train()

        # Reference actor (frozen)
        self.actor_ref = copy.deepcopy(actor_lm).to(self._device)
        self.actor_ref.eval()
        for p in self.actor_ref.parameters():
            p.requires_grad_(False)

        # --- Critic setup (supports both SICQL + HuggingFace critics) ---
        self.critic = critic_head.to(self._device) if critic_head is not None else None
        self.critic_ref = None
        self.critic_type = None

        if self.critic is not None:
            # Detect critic type
            if hasattr(self.critic, "q_head") and hasattr(self.critic, "encoder"):
                # Looks like SICQL InContextQModel
                self.critic_type = "sicql"
                self.critic.train()

                self.critic_ref = type(self.critic)(
                    encoder=copy.deepcopy(self.critic.encoder),
                    q_head=copy.deepcopy(self.critic.q_head),
                    v_head=copy.deepcopy(self.critic.v_head),
                    pi_head=copy.deepcopy(self.critic.pi_head),
                    embedding_store=self.memory.embedding,
                    device=self._device,
                ).to(self._device)
                self.critic_ref.load_state_dict(self.critic.state_dict())
                self.critic_ref.eval()
                for p in self.critic_ref.parameters():
                    p.requires_grad_(False)

            else:
                # Assume HuggingFace-style sequence classification model
                self.critic_type = "hf"
                self.critic.train()
                self.critic_ref = copy.deepcopy(self.critic).to(self._device)
                self.critic_ref.eval()
                for p in self.critic_ref.parameters():
                    p.requires_grad_(False)

        # --- Tokenizer setup ---
        self.tok = tokenizer
        if getattr(self.tok, "pad_token", None) is None:
            if getattr(self.tok, "eos_token", None) is not None:
                self.tok.pad_token = self.tok.eos_token
            else:
                self.tok.add_special_tokens({"pad_token": "[PAD]"})
                self.actor.resize_token_embeddings(len(self.tok))
                self.actor_ref.resize_token_embeddings(len(self.tok))

        self.memory = memory

    def device(self):
        return self._device

    # ---------- Actor (HF LM) ----------
    @torch.no_grad()
    def sample_group(self, prompt, group_size, max_new_tokens=256, temperature=0.6, top_p=0.95):
        """Generate multiple responses from the HuggingFace LM."""
        inputs = self.tok(
            prompt,
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=1024
        ).to(self._device)

        responses = []
        for _ in range(group_size):
            out = self.actor.generate(
                **inputs,
                do_sample=True,
                temperature=temperature,
                top_p=top_p,
                max_new_tokens=max_new_tokens,
                pad_token_id=self.tok.pad_token_id or self.tok.eos_token_id,
                eos_token_id=self.tok.eos_token_id,
            )
            prompt_len = inputs["input_ids"].size(1)
            response_ids = out[0, prompt_len:]
            response = self.tok.decode(response_ids, skip_special_tokens=True)
            responses.append(response)
        return responses

    def logprob_sum(self, prompt, response):
        """Compute logprob sum with online actor LM."""
        return self._sum_logprobs(self.actor, prompt, response)

    @torch.no_grad()
    def logprob_sum_ref(self, prompt, response):
        """Compute logprob sum with frozen actor reference."""
        return self._sum_logprobs(self.actor_ref, prompt, response).detach()

    def _sum_logprobs(self, model, prompt, response):
        """Sum log probabilities for actor outputs."""
        full_text = prompt + response
        enc = self.tok(full_text, return_tensors="pt", truncation=True, max_length=2048).to(self._device)
        input_ids = enc["input_ids"]
        attention_mask = enc["attention_mask"]

        outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)
        logits = outputs.logits
        log_probs = F.log_softmax(logits, dim=-1)

        shift_logits = log_probs[..., :-1, :].contiguous()
        shift_labels = input_ids[..., 1:].contiguous()
        token_log_probs = torch.gather(shift_logits, dim=-1, index=shift_labels.unsqueeze(-1)).squeeze(-1)

        prompt_len = len(self.tok(prompt, return_tensors="pt")["input_ids"][0])
        response_log_probs = token_log_probs[:, prompt_len - 1:]
        return response_log_probs.sum(dim=1).squeeze()

    # ---------- Critic (SICQL InContextQModel) ----------
    def _critic_logit(self, model, prompt, response):
        """Compute Q-value logit from SICQL critic (InContextQModel)."""
        if model is None:
            raise ValueError("Critic head required for critic mode")
        if self.memory is None:
            raise ValueError("Memory with embedding system required for critic")

        # Get embeddings
        prompt_emb_np = self.memory.embedding.get_or_create(prompt)
        response_emb_np = self.memory.embedding.get_or_create(response)
        prompt_emb = torch.tensor(prompt_emb_np, device=self._device, dtype=torch.float32).unsqueeze(0)
        response_emb = torch.tensor(response_emb_np, device=self._device, dtype=torch.float32).unsqueeze(0)

        out = model(prompt_emb, response_emb)
        return out["q_value"].squeeze()

    def critic_logit(self, prompt, response):
        return self._critic_logit(self.critic, prompt, response)

    @torch.no_grad()
    def critic_logit_ref(self, prompt, response):
        return self._critic_logit(self.critic_ref, prompt, response).detach()

    # ---------- Ref sync ----------
    def hard_reset_ref(self):
        """Reset references to current models."""
        self.actor_ref.load_state_dict(self.actor.state_dict())
        self.actor_ref.eval()
        if self.critic and self.critic_ref:
            self.critic_ref.load_state_dict(self.critic.state_dict())
            self.critic_ref.eval()



# ==============================
# PACS Trainer Core
# ==============================

class PACSCoreTrainer:
    """PACS trainer implementing RLVR via Supervised with RLOO regularization"""
    
    def __init__(
        self,
        policy: HybridSICQLAdapter,
        cfg: PACSConfig,
        verifier: Callable[[str, str, Optional[Dict[str, Any]]], int],
        logger: Optional[Callable[[Dict[str, Any]], None]] = None,
        online_training: bool = False
    ):
        """
        Args:
            policy: Policy adapter (HybridSICQLAdapter)
            cfg: PACS configuration
            verifier: Function to verify (prompt, response, meta) → {0,1}
            logger: Optional logging function
        """
        self.policy = policy
        self.cfg = cfg
        self.verifier = verifier
        self.logger = logger or (lambda d: None)
        self.online_training = online_training

        # Determine which parameters to train
        if cfg.score_mode == "logprob":
            params = self.policy.actor.parameters()
        elif cfg.score_mode == "critic":
            if self.policy.critic is None:
                raise ValueError("critic_head required for score_mode='critic'")
            params = self.policy.critic.parameters()
        else:
            raise ValueError("score_mode must be 'logprob' or 'critic'")
        
        # Optimizer
        self.optimizer = torch.optim.AdamW(
            params, 
            lr=cfg.lr, 
            weight_decay=cfg.weight_decay
        )
        
        # Learning rate scheduler
        self.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
            self.optimizer, 
            mode="min", 
            factor=0.5, 
            patience=5
        )
        
        # Loss function
        # Use BCEWithLogitsLoss to maintain numerical stability. 
        # The paper's Equation 1 shows the loss is computed with σ(ψ), 
        # but PyTorch's BCEWithLogitsLoss combines sigmoid + BCE in a numerically stable way.
        self.loss_fn = WeightedBCEWithLogits(pos_weight=cfg.pos_weight)
        
        # Training state
        self._step = 0
        self.best_loss = float("inf")
        self.early_stop_counter = 0

    # ---- Helpers ----
    def _rhat_vector(self, prompt: str, responses: List[str]) -> torch.Tensor:
        """
        Compute r̂ for each response in group based on chosen score_mode.
        In both modes:
        r̂ = β·(score_online - score_ref)
        """
        vals = []
        
        if self.cfg.score_mode == "logprob":
            for r in responses:
                lp = self.policy.logprob_sum(prompt, r)          # requires grad
                lpr = self.policy.logprob_sum_ref(prompt, r)     # no grad
                vals.append(self.cfg.beta * (lp - lpr))
        else:  # critic
            for r in responses:
                lg = self.policy.critic_logit(prompt, r)         # requires grad
                lgr = self.policy.critic_logit_ref(prompt, r)    # no grad
                vals.append(self.cfg.beta * (lg - lgr))
        
        return torch.stack(vals)  # (G,)
    
    def _psi_rloo(self, rhat: torch.Tensor) -> torch.Tensor:
        """
        Compute ψ using RLOO (leave-one-out) regularization.
        
        ψ_i = r̂_i − mean_{j≠i} r̂_j
        """
        G = rhat.numel()
        if G <= 1:
            return rhat
            
        # Compute mean of all r̂ values
        total = rhat.sum()
        
        # Compute leave-one-out mean for each element
        psi = rhat - (total - rhat) / (G - 1)
        return psi
    
    def _verify_labels(
        self, 
        prompt: str, 
        responses: List[str], 
        meta: Optional[Dict[str, Any]]
    ) -> torch.Tensor:
        """Verify responses and return labels (0 or 1)"""
        labels = []
        for r in responses:
            try:
                labels.append(float(self.verifier(prompt, r, meta)))
            except Exception as e:
                print(f"Verification failed: {e}")
                labels.append(0.0)
        
        return torch.tensor(labels, dtype=torch.float32, device=self.policy.device())
    
    def _log(self, d: Dict[str, Any]):
        """Log metrics with step counter"""
        d = {"step": self._step, **d}
        self.logger(d)
    
    def _check_early_stopping(self, loss: float) -> bool:
        """Check if early stopping criteria are met"""
        if loss < self.best_loss - self.cfg.early_stopping_min_delta:
            self.best_loss = loss
            self.early_stop_counter = 0
            return False
        
        self.early_stop_counter += 1
        return self.early_stop_counter >= self.cfg.early_stopping_patience

    # ---- Training API ----
    def train(
        self, 
        dataset: RLVRDataset, 
        max_steps: Optional[int] = None
    ) -> Dict[str, Any]:
        """
        Train policy using PACS with RLOO regularization.
        
        Args:
            dataset: RLVRDataset containing training samples
            max_steps: Maximum number of training steps
            
        Returns:
            Training statistics
        """
        # Training metrics
        metrics = {
            "losses": [],
            "psi_means": [],
            "rhat_means": [],
            "label_rates": [],
            "acc_proxies": [],
            "entropies": []
        }
        
        # Training loop
        step = 0
        while max_steps is None or step < max_steps:
            try:
                # Sample random item from dataset
                item = dataset[random.randint(0, len(dataset) - 1)]

                if self.online_training:
                    self._train_on_item(item, metrics)
                else:
                    self._train_on_item_offline(item, metrics)
                
                # Periodic reference model reset
                if self.cfg.steps_per_reset and (self._step % self.cfg.steps_per_reset == 0):
                    self.policy.hard_reset_ref()
                
                # Log metrics periodically
                if self._step % self.cfg.log_every == 0:
                    self._log({
                        "loss": metrics["losses"][-1] if metrics["losses"] else 0.0,
                        "psi_mean": metrics["psi_means"][-1] if metrics["psi_means"] else 0.0,
                        "rhat_mean": metrics["rhat_means"][-1] if metrics["rhat_means"] else 0.0,
                        "label_pos_rate": metrics["label_rates"][-1] if metrics["label_rates"] else 0.0,
                        "acc_proxy": metrics["acc_proxies"][-1] if metrics["acc_proxies"] else 0.0,
                        "mode": self.cfg.score_mode,
                        "entropy": metrics["entropies"][-1] if metrics["entropies"] else 0.0
                    })
                
                step += 1
                self._step += 1
                
            except Exception as e:
                print(f"Training error at step {self._step}: {e}")
                continue
        
        # Return training statistics
        return {
            "avg_loss": float(np.mean(metrics["losses"])) if metrics["losses"] else 0.0,
            "avg_psi": float(np.mean(metrics["psi_means"])) if metrics["psi_means"] else 0.0,
            "avg_rhat": float(np.mean(metrics["rhat_means"])) if metrics["rhat_means"] else 0.0,
            "pos_rate": float(np.mean(metrics["label_rates"])) if metrics["label_rates"] else 0.0,
            "acc_proxy": float(np.mean(metrics["acc_proxies"])) if metrics["acc_proxies"] else 0.0,
            "entropy": float(np.mean(metrics["entropies"])) if metrics["entropies"] else 0.0,
            "steps": self._step
        }
    

    def _train_on_item_offline(self, item: RLVRItem, metrics: Dict[str, list]) -> None:
        """
        Train directly on dataset rewards (offline supervised PACS).
        
        Uses RLVRItem.prompt, RLVRItem.response, and RLVRItem.reward
        as the training signal without generating new responses.
        """
        prompt = item.prompt
        response = item.response
        reward_val = float(item.reward)

        # Convert reward into tensor
        reward = torch.tensor([reward_val], device=self.policy.device(), dtype=torch.float32)

        # --- Compute r̂ (online vs reference) ---
        if self.cfg.score_mode == "logprob":
            lp = self.policy.logprob_sum(prompt, response)          # requires grad
            lpr = self.policy.logprob_sum_ref(prompt, response)     # no grad
            rhat = self.cfg.beta * (lp - lpr)
        else:  # critic mode
            lg = self.policy.critic_logit(prompt, response)         # requires grad
            lgr = self.policy.critic_logit_ref(prompt, response)    # no grad
            rhat = self.cfg.beta * (lg - lgr)

        psi = rhat.unsqueeze(0)  # keep batch dimension

        # --- Compute supervised loss ---
        loss = self.loss_fn(psi, reward)

        # --- Optimization ---
        self.optimizer.zero_grad(set_to_none=True)
        loss.backward()
        if self.cfg.grad_clip:
            torch.nn.utils.clip_grad_norm_(
                self.optimizer.param_groups[0]["params"], 
                self.cfg.grad_clip
            )
        self.optimizer.step()
        self.scheduler.step(loss.item())

        # --- Metrics ---
        with torch.no_grad():
            psi_val = float(psi.item())
            pred = (torch.sigmoid(psi) > 0.5).float()
            acc_proxy = float((pred == reward).float().mean().item())

            metrics["losses"].append(float(loss.item()))
            metrics["psi_means"].append(psi_val)
            metrics["rhat_means"].append(float(rhat.item()))
            metrics["label_rates"].append(reward_val)
            metrics["acc_proxies"].append(acc_proxy)

        # --- Log ---
        self._log({
            "mode": f"{self.cfg.score_mode}-offline",
            "loss": float(loss.item()),
            "psi": psi_val,
            "rhat": float(rhat.item()),
            "reward": reward_val,
            "acc_proxy": acc_proxy,
            "prompt_snippet": prompt[:80] + ("..." if len(prompt) > 80 else "")
        })

    
    def _train_on_item(self, item: RLVRItem, metrics: Dict[str, list]) -> None:
        """Train on a single dataset item and update metrics"""
        prompt = item.prompt
        meta = item.meta
        
        # Generate response group
        responses = self.policy.sample_group(
            prompt,
            group_size=self.cfg.group_size,
            max_new_tokens=self.cfg.max_new_tokens,
            temperature=self.cfg.temperature,
            top_p=self.cfg.top_p,
        )
        
        # Verify responses
        labels = self._verify_labels(prompt, responses, meta)
        
        # Compute r̂ and ψ
        rhat = self._rhat_vector(prompt, responses)  # (G,)
        psi = self._psi_rloo(rhat)
        
        # Compute loss
        loss = self.loss_fn(psi, labels)
        
        # Detach psi for component calculations
        psi_detached = psi.detach()
        psi_sigmoid = torch.sigmoid(psi_detached)
        
        # Calculate actor component magnitude
        with torch.no_grad():
            # Actor component: policy improvement term
            # [l(q, o; πθ)]∇θ log πθ(o|q)
            cross_entropy_loss = -(labels * torch.log(psi_sigmoid + 1e-8) + 
                                (1 - labels) * torch.log(1 - psi_sigmoid + 1e-8))
            actor_component = cross_entropy_loss.mean().item()
            
            # Critic component: reward estimation term
            # (R(q, o) - σ(ψ(q, o; πθ)))∇θψ(q, o; πθ)
            critic_raw = (labels - psi_sigmoid).mean().item()  # Prediction error
            critic_grad = psi.mean().item()  # Gradient magnitude
            critic_component = critic_raw * critic_grad  # Full critic update
            
            # Calculate coupling ratio (absolute values to handle sign)
            coupling_ratio = abs(actor_component) / (abs(critic_component) + 1e-8)
            
            # Track policy entropy for exploration analysis
            policy_entropy = -torch.mean(psi_sigmoid * torch.log(psi_sigmoid + 1e-8) + 
                                        (1 - psi_sigmoid) * torch.log(1 - psi_sigmoid + 1e-8)).item()
            sequence_entropy = self.policy.calculate_response_entropy(prompt, responses)

        # Optimization
        self.optimizer.zero_grad(set_to_none=True)
        loss.backward()
        
        # Gradient clipping
        if self.cfg.grad_clip is not None:
            torch.nn.utils.clip_grad_norm_(
                self.optimizer.param_groups[0]["params"], 
                self.cfg.grad_clip
            )
        
        self.optimizer.step()
        
        # Update learning rate
        self.scheduler.step(loss.item())
        
        # Compute metrics
        with torch.no_grad():
            psi_mean = float(psi.mean().item())
            rhat_mean = float(rhat.mean().item())
            label_pos_rate = float(labels.mean().item())
            preds = (torch.sigmoid(psi) > 0.5).float()
            acc_proxy = float((preds == labels).float().mean().item())
            entropy = self.policy.calculate_response_entropy(prompt, responses)
        
        grad_analysis = self._analyze_gradients(psi, labels)

        metrics["actor_components"].append(grad_analysis["actor_component"])
        metrics["critic_raws"].append(grad_analysis["critic_raw"])
        metrics["critic_grads"].append(grad_analysis["critic_grad"])
        metrics["critic_components"].append(grad_analysis["critic_component"])
        metrics["coupling_ratios"].append(grad_analysis["coupling_ratio"])
        metrics["binary_entropies"].append(grad_analysis["binary_entropy"])
        metrics["losses"].append(float(loss.item()))
        metrics["psi_means"].append(psi_mean)
        metrics["rhat_means"].append(rhat_mean)
        metrics["label_rates"].append(label_pos_rate)
        metrics["acc_proxies"].append(acc_proxy)
        metrics["entropies"].append(entropy)
        metrics["policy_entropies"].append(policy_entropy)
        metrics["sequence_entropies"].append(sequence_entropy)
        metrics["response_entropy"].append(
            self.policy.calculate_response_entropy(prompt, responses)
        )

        # Log metrics
        self._log({
            "loss": float(loss.item()),
            "psi_mean": psi_mean,
            "rhat_mean": rhat_mean,
            "label_pos_rate": label_pos_rate,
            "acc_proxy": acc_proxy,
            "entropy": entropy,
            "mode": self.cfg.score_mode,
            "actor_component": actor_component,
            "critic_component": critic_component,
            "coupling_ratio": coupling_ratio,
            **grad_analysis,
            "policy_entropy": self.policy.calculate_response_entropy(prompt, responses)
        })

    def _analyze_gradients(self, psi: torch.Tensor, labels: torch.Tensor) -> Dict[str, float]:
        """Analyze gradient components as in PACS paper Equation 6"""
        # Compute sigmoid once for efficiency
        psi_sigmoid = torch.sigmoid(psi)
        
        # 1. ACTOR component: policy improvement term
        # [l(q, o; πθ)]∇θ log πθ(o|q)
        cross_entropy_loss = -(labels * torch.log(psi_sigmoid + 1e-8) + 
                            (1 - labels) * torch.log(1 - psi_sigmoid + 1e-8))
        actor_component = cross_entropy_loss.mean().item()
        
        # 2. CRITIC component breakdown
        prediction_error = (labels - psi_sigmoid)
        critic_raw = prediction_error.mean().item()  # R - σ(ψ)
        critic_grad = psi.mean().item()  # ∇θψ
        critic_component = (prediction_error * psi).mean().item()  # Full critic update
        
        # 3. Coupling ratio (normalized for stability)
        coupling_ratio = abs(actor_component) / (abs(critic_component) + 1e-8)
        
        # 4. Binary entropy (prediction confidence)
        binary_entropy = -torch.mean(
            psi_sigmoid * torch.log(psi_sigmoid + 1e-8) + 
            (1 - psi_sigmoid) * torch.log(1 - psi_sigmoid + 1e-8)
        ).item()
        
        return {
            "actor_component": actor_component,
            "critic_raw": critic_raw,
            "critic_grad": critic_grad,
            "critic_component": critic_component,
            "coupling_ratio": coupling_ratio,
            "binary_entropy": binary_entropy
        }

# Add to imports at the top

class PACSTrainer(BaseTrainer):
    """
    PACS trainer that implements the RLVR via Supervised algorithm.
    
    This integrates properly with SICQL infrastructure by:
    - Using SICQL policy heads as critics
    - Leveraging SICQL verifiers for labels
    - Maintaining proper RLOO implementation
    
    Key features:
    - Two scoring modes: "logprob" (actor-based) and "critic" (SICQL-based)
    - Proper RLOO (leave-one-out) regularization
    - Reference model resets
    - Safety mechanisms to prevent model collapse
    """
    
    def __init__(self, cfg, memory, logger):
        super().__init__(cfg, memory, logger)
        self.cfg = cfg
        self.memory = memory
        self.logger = logger
        
        # Device management
        self.device = torch.device(
            "cuda" if torch.cuda.is_available() else "cpu"
        )
        
        # Initialize configuration
        self._init_config(cfg)
        
        # Initialize ScorableClassifier for prompt type classification
        try:
            self.classifier = ScorableClassifier(
                memory=self.memory,
                logger=self.logger,
                config_path="config/domain/seeds.yaml",
                metric="cosine"
            )
            self.logger.log("ScorableClassifierInitialized", {
                "message": "Successfully initialized ScorableClassifier for prompt type classification"
            })
        except Exception as e:
            self.logger.log("ScorableClassifierError", {
                "error": str(e),
                "message": "Failed to initialize ScorableClassifier, falling back to simple classification"
            })
            self.classifier = None
        
        # Track training state
        self.best_loss = float("inf")
        self.early_stop_counter = 0
        self.trainer = None  # Will be initialized in train()
        self.pacs_core = None  # Core PACS implementation        

        # Log initialization
        self.logger.log(
            "PACSTrainerInitialized",
            {
                "score_mode": self.score_mode,
                "beta": self.beta,
                "group_size": self.group_size,
                "device": str(self.device),
            },
        )
    
    def _init_config(self, cfg):
        """Initialize training parameters from config"""
        self.score_mode = cfg.get("score_mode", "critic")  # "logprob" or "critic"
        self.beta = cfg.get("beta", 1.0)
        self.group_size = cfg.get("group_size", 8)
        self.max_steps = cfg.get("max_steps", 1000)
        self.lr = cfg.get("lr", 1e-6)
        self.early_stopping_patience = cfg.get("patience", 3)
        self.early_stopping_min_delta = cfg.get("min_delta", 1e-4)
        self.pos_weight = cfg.get("pos_weight", 1.0)
        self.steps_per_reset = cfg.get("steps_per_reset", 200)
        self.dimension = cfg.get("dimension", "alignment")
        self.target_type = cfg.get("target_type", "document")
        self.model_path = cfg.get("model_path", "models/pacs")
        self.model_version = cfg.get("model_version", "v1")
    
    def _get_verifier(self, dimension: str) -> Callable[[str, str, Optional[Dict]], int]:
        """Get verifier function for the given dimension"""
        if dimension == "math":
            from stephanie.scoring.verifiers import boxed_math_verifier
            return boxed_math_verifier
        elif dimension == "code":
            from stephanie.scoring.verifiers import code_verifier
            return code_verifier
        else:
            # Generic verifier that uses SICQL scores
            def generic_verifier(prompt: str, response: str, meta: Optional[Dict] = None) -> int:
                # In practice, this would use your SICQL system
                from stephanie.scoring.scorable_factory import ScorableFactory
                from stephanie.scoring.transforms.regression_tuner import RegressionTuner
                
                # Create scorable
                scorable = ScorableFactory.from_text(
                    response, 
                    TargetType.DOCUMENT,
                    context=prompt
                )
                
                # Get score from SICQL
                score = self.memory.scores.get_score(
                    goal_id=meta.get("goal_id") if meta else None,
                    scorable=scorable
                )
                
                # Use tuner if available
                tuner = RegressionTuner(dimension=dimension)
                if tuner.is_loaded():
                    normalized = tuner.transform(score.score if score else 0.0)
                    return 1 if normalized > 0.5 else 0
                return 1 if (score.score if score else 0.0) > 0.5 else 0
                
            return generic_verifier
    
    def _build_policy_adapter(self, dimension: str):
        """Build policy adapter with SICQL integration"""
        from transformers import AutoModelForCausalLM, AutoTokenizer
        
        # Load base model (same as SICQL uses)
        base_model = "Qwen/Qwen2.5-1.5B"
        actor = AutoModelForCausalLM.from_pretrained(base_model)
        tokenizer = AutoTokenizer.from_pretrained(base_model)
        
        # For critic mode, load SICQL policy head
        critic = None
        if self.score_mode == "critic":
            scorer = SICQLScorer(self.cfg, self.memory, self.logger)
            critic = scorer.get_model(self.dimension)
        # Create policy adapter
        adapter = HybridSICQLAdapter(
            memory=self.memory,
            actor_lm=actor,
            tokenizer=tokenizer,
            critic_head=critic,
            device=self.device
        )
        
        return adapter
    
    def train(self, dataset, dimension: str = "alignment"):
        """
        Train model using proper PACS algorithm with RLOO.
        
        Args:
            dataset: RLVRDataset containing queries and meta
            dimension: Dimension to train on (for verifier selection)
        
        Returns:
            Training statistics and model metadata
        """
        self.logger.log("PACSTrainingStarted", {"dimension": dimension})
        
        # 1. Build policy adapter (with SICQL integration)
        policy = self._build_policy_adapter(dimension)
        
        # 2. Get verifier for this dimension
        verifier = self._get_verifier(dimension)
        
        # 3. Create PACS config
        pacs_cfg = PACSConfig(
            score_mode=self.score_mode,
            beta=self.beta,
            group_size=self.group_size,
            max_new_tokens=256,
            temperature=0.6,
            top_p=0.96,
            lr=self.lr,
            weight_decay=0.01,
            grad_clip=1.0,
            steps_per_reset=self.steps_per_reset,
            pos_weight=self.pos_weight,
            log_every=10
        )
        
        # 4. Create core PACS trainer
        self.pacs_core = PACSCoreTrainer(
            policy=policy,
            cfg=pacs_cfg,
            verifier=verifier,
            logger=self._log_metrics,
            online_training=False
        ) 
        
        # 5. Train model
        training_stats = self.pacs_core.train(dataset, max_steps=self.max_steps)
        
        # 6. Save model and metadata
        meta = self._save_model(policy, dimension, training_stats, dataset=dataset)

        # 7. Log training stats
        self._log_training_stats(dimension, meta)
        
        
        self.logger.log(
            "PACSTrainingComplete",
            {
                "dimension": dimension,
                "final_loss": meta["avg_loss"],
                "policy_entropy": meta["policy_entropy"]
            },
        )
        
        return meta
    
    def _log_metrics(self, metrics: Dict[str, Any]):
        """Log metrics from PACS core trainer"""
        self.logger.log("PACSMetrics", metrics)
        
        # Track for early stopping
        if "loss" in metrics:
            if metrics["loss"] < self.best_loss - self.early_stopping_min_delta:
                self.best_loss = metrics["loss"]
                self.early_stop_counter = 0
            else:
                self.early_stop_counter += 1
    
    def _save_model(self, policy, dimension: str, training_stats: Dict[str, Any], dataset=None) -> Dict[str, Any]:
        """Save PACS model with metadata"""
        # Create output directory
        output_dir = os.path.join(self.model_path, dimension)
        os.makedirs(output_dir, exist_ok=True)
        
        # Save actor model
        actor_path = os.path.join(output_dir, "actor")
        policy.actor.save_pretrained(actor_path)
        
        # Save critic if in critic mode
        if self.score_mode == "critic" and policy.critic is not None:
            critic_path = os.path.join(output_dir, "critic")
            os.makedirs(critic_path, exist_ok=True) 
            # Save SICQL components
            torch.save(policy.critic.encoder.state_dict(), os.path.join(critic_path, "encoder.pt"))
            torch.save(policy.critic.q_head.state_dict(), os.path.join(critic_path, "q_head.pt"))
            torch.save(policy.critic.v_head.state_dict(), os.path.join(critic_path, "v_head.pt"))
            torch.save(policy.critic.pi_head.state_dict(), os.path.join(critic_path, "pi_head.pt"))
        
        # Calculate policy metrics
        policy_entropy = self._calculate_policy_entropy(policy, dataset, tokenizer=policy.tok)
        policy_stability = self._calculate_policy_stability(policy, dataset=dataset)

        # Build metadata
        meta = {
            "version": self.model_version,
            "dimension": dimension,
            "score_mode": self.score_mode,
            "policy_entropy_calc": policy_entropy,
            "policy_stability_calc": policy_stability,
            "actor_model": "Qwen2.5-1.5B",
            "critic_model": "SICQL" if self.score_mode == "critic" else None,
            "beta": self.beta,
            "group_size": self.group_size,
            "avg_loss": float(training_stats["avg_loss"]),
            "policy_entropy": training_stats["entropy"],
            "policy_stability": training_stats["acc_proxy"],
            "steps": training_stats["steps"],
            "device": str(self.device),
            "model_path": output_dir,
            "timestamp": datetime.now().isoformat(),
            "config": self.cfg if isinstance(self.cfg, dict) else self.cfg.__dict__
        }

        # Split into direct columns vs. extra_data
        db_fields = {
            "model_type": "pacs",
            "target_type": self.target_type,
            "dimension": meta["dimension"],
            "version": meta["version"],
            "score_mode": meta["score_mode"],
            "model_path": meta["model_path"],
        }

        extra_data = {k: v for k, v in meta.items() if k not in db_fields and k not in ["model_path", "dimension", "version", "score_mode"]}

            
        # Save metadata
        with open(os.path.join(output_dir, "meta.json"), "w") as f:
            json.dump(meta, f, indent=2)

        # Insert into DB
        model_version = ModelVersionORM(
            **db_fields,
            extra_data=extra_data
        )
        self.memory.session.add(model_version)
        self.memory.session.commit()
        
        return meta
    
    def _calculate_policy_entropy(self, policy, dataset, tokenizer, sample_size: int = 100) -> float:
        """
        Calculate average entropy of the actor policy over a subset of dataset.
        """
        import random
        subset = random.sample(dataset, min(sample_size, len(dataset)))

        entropies = []
        for item in subset:
            prompt, response = item.prompt, item.response
            if not response:
                continue

            logits = self._get_logits(policy, prompt, response, tokenizer, device=policy._device)
            # take only the response part
            resp_len = len(tokenizer.encode(response))
            resp_logits = logits[:, -resp_len:, :]

            probs = F.softmax(resp_logits, dim=-1)
            entropy = -(probs * probs.log()).sum(dim=-1).mean().item()
            entropies.append(entropy)

        return float(sum(entropies) / len(entropies)) if entropies else 0.0

    def _get_logits(self, policy, prompt: str, response: str, tokenizer, device="cuda"):
        """Tokenize prompt+response and run through actor LM to get logits."""
        inputs = tokenizer(prompt + response, return_tensors="pt").to(device)
        with torch.no_grad():
            outputs = policy.actor(**inputs)
            logits = outputs.logits  # shape: [batch, seq_len, vocab_size]
        return logits

    def _calculate_policy_stability(self, policy, dataset, sample_size: int = 100) -> float:
        subset = random.sample(dataset, min(sample_size, len(dataset)))
        rewards, kl_divs = [], []

        for item in subset:
            prompt, response = item.prompt, item.response
            if not response:
                continue

            # Current policy logits
            logits_cur = self._get_logits(policy, prompt, response, policy.tok, device=policy._device)
            probs_cur = F.softmax(logits_cur, dim=-1)

            # Reference policy logits
            logits_ref = self._get_logits(policy, prompt, response, policy.tok, device=policy._device)
            probs_ref = F.softmax(logits_ref, dim=-1)

            kl = F.kl_div(probs_ref.log(), probs_cur, reduction="batchmean").item()
            kl_divs.append(kl)

            if policy.critic is not None:
                reward = policy.critic_logit(prompt, response).item()
                rewards.append(reward)

        reward_var = np.var(rewards) if rewards else 0.0
        kl_mean = np.mean(kl_divs) if kl_divs else 0.0
        return 1.0 / (1.0 + reward_var + kl_mean)
    
    def _log_training_stats(self, dimension: str, meta: Dict[str, Any]):
        """Log training stats to database"""
        training_stats = TrainingStatsORM(
            model_type="pacs",
            target_type=self.target_type,
            embedding_type=self.memory.embedding.name,
            dimension=dimension,
            version=self.model_version,
            avg_q_loss=meta["avg_loss"],
            policy_entropy=meta["policy_entropy"],
            policy_stability=meta["policy_stability"],
        )
        self.memory.session.add(training_stats)
        self.memory.session.commit()
    
    def _update_belief_cartridge(self, context: dict, dimension: str, meta: Dict[str, Any]):
        """Update belief cartridges with policy stats and proper goal ID"""
        # Extract goal_id from context
        goal_id = context.get("goal", {}).get("id")
        if not goal_id:
            self.logger.log("MissingGoalID", {"dimension": dimension})
        
        bc = BeliefCartridgeORM(
            title=f"PACS Policy - {dimension}",
            content=f"PACS training completed for {dimension} dimension",
            goal_id=goal_id,  # FIXED: now properly set
            domain=dimension,
            policy_entropy=meta["policy_entropy"],
            policy_stability=meta["policy_stability"],
        )
        self.memory.session.add(bc)
        self.memory.session.commit()    

    def run(self, context: dict) -> dict:
        """
        Main entry point for PACS training.
        
        Context should contain:
        - "rlvr_dataset": RLVRDataset from CaseBook
        - "dimension": Dimension to train on
        """
        self.logger.log("PACSTrainerRun", {"context_keys": list(context.keys())})
        
        # Get dimension from context
        dimension = context.get("dimension", self.dimension)
        
        # Get dataset
        dataset = context.get("rlvr_dataset")
        if not dataset:
            self.logger.log("MissingDataset", {"required": "rlvr_dataset"})
            return context
        
        # Train model
        stats = self.train(dataset, dimension)
        
        # Update context
        context["pacs_training_stats"] = stats
        context["pacs_model_path"] = stats.get("model_path", "")
        
        return context

class EntityDetector:
    """Detects named entities in text using BERT-NER with fallback to heuristic rules."""
    def __init__(self, device: str = "cuda"):
        try:
            self.ner_pipeline = pipeline(
                "ner",
                model="dslim/bert-base-NER",
                aggregation_strategy="simple",
                device=0 if device == "cuda" else -1
            )
        except Exception as e:
            self.ner_pipeline = None

    def detect_entities(self, text: str) -> List[Tuple[int, int, str]]:
        if self.ner_pipeline:
            results = self.ner_pipeline(text)
            return [(r["start"], r["end"], r["entity_group"]) for r in results]
        return self._heuristic_entity_detection(text)

def embed_entity(self, text: str, span: Tuple[int, int]) -> torch.Tensor:
    """Embed an entity span with robust character-to-token alignment."""
    if span[0] >= span[1] or span[0] < 0 or span[1] > len(text):
        return torch.zeros(
            self.projection.fc2.out_features if self.projection_enabled else self.model.config.hidden_size, 
            device=self.device
        )
        
    inputs = self.tokenizer(
        text, 
        return_tensors="pt", 
        truncation=True, 
        max_length=512
    ).to(self.device)
    
    with torch.no_grad():
        outputs = self.model(**inputs, output_hidden_states=True)
        hidden_states = outputs.hidden_states[self.layer]
    
    # Align character spans to tokens
    start_token = inputs.char_to_token(0, span[0]) or 1
    end_token = inputs.char_to_token(0, span[1] - 1) or start_token + 1
    
    # Mean pool the entity span
    entity_vec = hidden_states[0, start_token:end_token+1, :].mean(dim=0)
    
    # Optional projection (paper-style)
    if self.projection_enabled and self.projection is not None:
        entity_vec = self.projection(entity_vec)

    return entity_vec

def index_scorables(self, scorables: List[Scorable]) -> int:
    """Index entities from scorables into Annoy with metadata."""
    new_embeddings, new_metadata = [], []
    for scorable in tqdm(scorables, desc="Indexing entities"):
        for start, end, entity_type in self.entity_detector.detect_entities(scorable.text):
            entity_text = scorable.text[start:end].strip()
            if len(entity_text) < 2:
                continue
            try:
                # Create embedding for this span
                emb = self.embed_entity(scorable.text, (start, end))
                new_embeddings.append(emb.detach().cpu().numpy())
                
                # Attach metadata for retrieval
                new_metadata.append({
                    "scorable_id": str(scorable.id),
                    "scorable_type": scorable.target_type,
                    "entity_text": entity_text,
                    "start": start,
                    "end": end,
                    "entity_type": entity_type,
                    "source_text": scorable.text[:100] + "..."
                })
            except Exception as e:
                self.logger.error(f"Entity embedding failed: {e}")

    # Push into Annoy index
    if new_embeddings:
        self.index.add(np.array(new_embeddings), new_metadata)
    return len(new_embeddings)

...

class AnnoyIndex:
    """Wrapper for Annoy approximate nearest neighbor index with metadata management."""
    def __init__(self, dim: int = 500, index_path: str = "data/ner_retriever/index"):
        self.dim = dim
        self.index_path = index_path
        self.index = annoy.AnnoyIndex(dim, "angular")  # cosine similarity
        self.metadata = []
        self._load_index()

    def add(self, embeddings: np.ndarray, metadata_list: List[Dict], n_trees: int = 10):
        # Deduplicate by (scorable_id, entity_text)
        for i, (vec, meta) in enumerate(zip(embeddings, metadata_list)):
            if not any(m["scorable_id"] == meta["scorable_id"] and
                       m["entity_text"] == meta["entity_text"] for m in self.metadata):
                self.index.add_item(len(self.metadata), vec)
                self.metadata.append(meta)
        self.index.build(n_trees)
        self._save_index()

    def search(self, query: np.ndarray, k: int = 10):
        indices, distances = self.index.get_nns_by_vector(query, k, include_distances=True)
        results = []
        for idx, dist in zip(indices, distances):
            sim = 1 - (dist ** 2) / 2  # angular → cosine
            meta = self.metadata[idx].copy()
            meta["similarity"] = float(sim)
            results.append(meta)
        return results

def retrieve_entities(self, query: str, k: int = 5, min_similarity: float = 0.6, domain: str = None) -> List[Dict]:
    # Step 1: Embed the query
    query_emb = self.embed_type_query(query)

    # Step 2: Search raw results
    results = self.index.search(query_emb, k*2)

    # Step 3: Load domain calibration data
    if domain is None:
        domain = self._get_current_domain(query)
    calibration = self._load_calibration_data(domain)

    # Step 4: Apply polynomial calibration
    for result in results:
        if "similarity" in result and "ner" in calibration:
            poly = np.poly1d(calibration["ner"]["coefficients"])
            calibrated = float(poly(result["similarity"]))

            # Apply guardrails
            if result["similarity"] > 0.8:
                calibrated = min(1.0, calibrated * 1.05)

            result["calibrated_similarity"] = max(0.0, min(1.0, calibrated))

    # Step 5: Filter by calibrated score
    return [
        r for r in results
        if r.get("calibrated_similarity", r["similarity"]) >= min_similarity
    ][:k]

def retrieve_entities(self, query: str, k: int = 5, min_similarity: float = 0.6, domain: str = None) -> List[Dict]:
    """Search for entities similar to the query with domain-aware calibration."""
    query_emb = self.embed_type_query(query)
    results = self.index.search(query_emb, k*2)

    # Apply calibration
    if domain is None:
        domain = self._get_current_domain(query)
    calibration = self._load_calibration_data(domain)

    for result in results:
        if "similarity" in result:
            poly = np.poly1d(calibration["ner"]["coefficients"])
            calibrated = float(poly(result["similarity"]))
            result["calibrated_similarity"] = max(0.0, min(1.0, calibrated))

    # Filter by threshold
    filtered_results = [
        r for r in results 
        if r.get("calibrated_similarity", r.get("similarity", 0.0)) >= min_similarity
    ][:k]

    return filtered_results

[
    {
        "scorable_id": "case_123",
        "scorable_type": "document",
        "entity_text": "Supreme Court",
        "entity_type": "ORG",
        "start": 45,
        "end": 58,
        "similarity": 0.81,
        "calibrated_similarity": 0.87,
        "source_text": "In 1992, the Supreme Court ruled that..."
    },
    {
        "scorable_id": "case_456",
        "scorable_type": "document",
        "entity_text": "Roe v. Wade",
        "entity_type": "CASE",
        "start": 12,
        "end": 23,
        "similarity": 0.78,
        "calibrated_similarity": 0.84,
        "source_text": "The precedent set in Roe v. Wade was challenged..."
    },
    ...
]