# stephanie/scoring/model/tiny_recursion.py
"""
Tiny Recursion Model (Tiny+) - Parameter-Efficient Recursive Neural Architecture

This module implements a compact, recursive neural network for multi-task evaluation
of AI model responses. The architecture combines recursive state updates with
multi-head output predictions, enabling efficient quality assessment across
multiple dimensions from embedding inputs.

Key Innovations:
- Recursive latent state updates with halting mechanisms
- Sparse Autoencoder (SAE) bottleneck for interpretable concepts
- Multi-head prediction for comprehensive quality assessment
- Heteroscedastic uncertainty estimation
- In-graph consistency regularization

Architecture Overview:
1. Recursive fusion of goal (x), response (y), and latent (z) states
2. Core processing blocks (attention or MLP-based)
3. SAE bottleneck for sparse concept representation
4. Multi-head prediction for scores, uncertainty, and auxiliary tasks

"""

from __future__ import annotations

from typing import Any, Dict, Optional, Tuple

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

# ---------------------------
# Core Building Blocks
# ---------------------------

class TinyBlock(nn.Module):
    """
    Basic residual block: LayerNorm → MLP → residual connection.

    Supports both 2D [batch, features] and 3D [batch, sequence, features] inputs.
    Uses GELU activation and dropout for regularization.
    """
    def __init__(self, d_model: int, dropout: float = 0.1):
        super().__init__()
        self.ln = nn.LayerNorm(d_model)
        self.mlp = nn.Sequential(
            nn.Linear(d_model, d_model * 4),  # Expansion factor 4
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(d_model * 4, d_model),  # Projection back
            nn.Dropout(dropout),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Apply residual block: x + MLP(LayerNorm(x))"""
        return x + self.mlp(self.ln(x))


class TinyBlockAttn(nn.Module):
    """
    Attention-enhanced residual block with Multi-Head Self-Attention.

    Architecture: LN → MHA → residual → TinyBlock → residual
    Automatically handles 2D/3D inputs and returns same dimensionality.
    """
    def __init__(self, d_model: int, n_heads: int = 4, dropout: float = 0.1):
        super().__init__()
        self.ln_attn = nn.LayerNorm(d_model)
        self.attn = nn.MultiheadAttention(
            embed_dim=d_model,
            num_heads=n_heads,
            dropout=dropout,
            batch_first=True  # [batch, seq, features]
        )
        self.drop = nn.Dropout(dropout)
        self.ff = TinyBlock(d_model, dropout=dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass with automatic shape handling.

        Args:
            x: Input tensor of shape [B, D] or [B, L, D]

        Returns:
            Output tensor with same shape as input
        """
        squeeze_back = False
        if x.dim() == 2:
            x = x.unsqueeze(1)  # [B, D] → [B, 1, D]
            squeeze_back = True

        q = k = v = self.ln_attn(x)
        h, _ = self.attn(q, k, v, need_weights=False)
        x = x + self.drop(h)  # Residual connection
        x = self.ff(x)        # Feed-forward with residual

        if squeeze_back:
            x = x.squeeze(1)  # [B, 1, D] → [B, D]
        return x


# ---------------------------
# Tiny Recursion Model (Tiny+)
# ---------------------------

class TinyRecursionModel(nn.Module):
    """
    Parameter-efficient recursive model for multi-task evaluation.

    Recursively updates latent state z using goal (x) and response (y) embeddings
    over multiple steps. Features comprehensive multi-head prediction and
    sparse autoencoder bottleneck for interpretable representations.

    Core Components:
    - Recursive state fusion: [x, y, z] → z'
    - Core processing stack: Attention or MLP blocks
    - SAE bottleneck: Sparse concept encoding
    - Multi-head prediction: 12 specialized output heads

    Inputs:
        x: Goal/condition embedding [B, D]
        y: Response embedding [B, D]
        z: Initial latent state [B, D] (typically zeros)

    Outputs:
        logits: Classification logits [B, vocab_size] (legacy compatibility)
        halt_logits: Halting signal logits [B]
        z_final: Final latent state after recursion [B, D]
        aux: Dictionary of auxiliary predictions and metrics
    """

    def __init__(
        self,
        d_model: int = 256,
        n_layers: int = 2,
        n_recursions: int = 6,
        vocab_size: int = 1024,
        use_attention: bool = False,
        dropout: float = 0.1,
        attn_heads: int = 4,
        step_scale: float = 0.1,           # Residual scaling for state updates
        consistency_mask_p: float = 0.10,  # Mask probability for consistency regularization
        len_norm_L: float = 512.0,         # Length normalization constant
        enable_agree_head: bool = True,    # Enable agreement prediction head
        enable_causal_sens_head: bool = True,  # Enable sensitivity prediction head
    ):
        super().__init__()

        # Model configuration
        self.d_model = d_model
        self.n_layers = n_layers
        self.n_recursions = n_recursions
        self.vocab_size = vocab_size
        self.use_attention = use_attention
        self.step_scale = step_scale
        self.consistency_mask_p = consistency_mask_p
        self.len_norm_L = float(len_norm_L)
        self.enable_agree_head = enable_agree_head
        self.enable_causal_sens_head = enable_causal_sens_head

        # Core processing stack
        if use_attention:
            blocks = [TinyBlockAttn(d_model, n_heads=attn_heads, dropout=dropout)
                      for _ in range(n_layers)]
        else:
            blocks = [TinyBlock(d_model, dropout=dropout) for _ in range(n_layers)]
        self.core = nn.Sequential(*blocks)

        # State fusion: combine goal, response, and latent states
        self.z_proj = nn.Linear(d_model * 3, d_model)  # [x, y, z] → z'
        self.final_ln = nn.LayerNorm(d_model)

        # Core prediction heads
        self.halt_head = nn.Linear(d_model, 1)            # Halting signal logits
        self.classifier = nn.Linear(d_model, vocab_size)  # Legacy classification

        # Extended prediction heads
        self.score_head = nn.Linear(d_model, 1)        # Quality score ∈ [0,1]
        self.logvar_head = nn.Linear(d_model, 1)       # Aleatoric uncertainty (log-variance)
        self.aux3_head = nn.Linear(d_model, 3)         # 3-way classification
        self.disagree_head = nn.Linear(d_model, 1)     # Disagreement prediction
        self.recon_head = nn.Linear(d_model, d_model)  # Embedding reconstruction
        self.consistency_head = nn.Linear(d_model, 1)  # Robustness prediction
        self.ood_head = nn.Linear(d_model, 1)          # OOD detection
        self.temp_head = nn.Linear(d_model, 1)         # Temperature calibration

        # Bridge heads
        self.agree_head = nn.Linear(d_model, 1)        # Cross-model agreement
        self.causal_sens_head = nn.Linear(d_model, 1)  # Perturbation sensitivity

        # Sparse Autoencoder (SAE) bottleneck
        self.sae_enc = nn.Sequential(
            nn.Linear(d_model, d_model // 2),  # Compression
            nn.ReLU(),
            nn.LayerNorm(d_model // 2),
        )
        self.sae_dec = nn.Linear(d_model // 2, d_model)  # Reconstruction
        self.sae_alpha = 0.05  # SAE reconstruction loss weight

        # Regularization
        self.head_drop = nn.Dropout(dropout)

    @staticmethod
    def _cos01(a: torch.Tensor, b: torch.Tensor, dim: int = -1, eps: float = 1e-6) -> torch.Tensor:
        """
        Compute cosine similarity mapped from [-1, 1] to [0, 1].

        Args:
            a, b: Input tensors to compare
            dim: Dimension for cosine computation
            eps: Numerical stability term

        Returns:
            Cosine similarity in range [0, 1] where 1 = identical
        """
        sim = F.cosine_similarity(a, b, dim=dim, eps=eps)
        return (sim + 1.0) * 0.5

    def _recur(
        self,
        x: torch.Tensor,
        y: torch.Tensor,
        z: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Execute recursive state updates over n_recursions steps.

        Process:
          1. Fuse [x, y, z] → z_next via projection and activation
          2. Process through core network stack
          3. Update halting signals
          4. Apply residual state update: z = z + step_scale * z_next
          5. Apply SAE bottleneck to final state

        Args:
            x: Goal embedding [B, D]
            y: Response embedding [B, D]
            z: Initial latent state [B, D]

        Returns:
            z_final: Final latent state after recursion [B, D]
            z_head: SAE-processed state for prediction heads [B, D]
            halt_logits: Maximum halting logits across steps [B, 1]
            tau: Temperature parameter for score calibration [B, 1]
            c: Sparse concept codes from SAE bottleneck [B, D//2]
        """
        B = x.size(0)
        device = x.device

        # Initialize halting signals to very negative values
        halt_logits = torch.full((B, 1), -1e9, device=device)
        z_cur = z  # Current latent state

        # Recursive state updates
        for _ in range(self.n_recursions):
            fused = torch.cat([x, y, z_cur], dim=-1)   # [B, 3 * D]
            z_next = torch.tanh(self.z_proj(fused))    # [B, D] with saturation
            z_next = self.core(z_next)                 # [B, D] core processing

            # Update halting signal (track maximum across steps)
            step_halt = self.halt_head(self.final_ln(z_next))  # [B, 1]
            halt_logits = torch.maximum(halt_logits, step_halt)

            # Residual state update with step scaling
            z_cur = z_cur + self.step_scale * z_next

        # Final normalization
        z_final = self.final_ln(z_cur)  # [B, D]

        # Sparse Autoencoder bottleneck
        c = self.sae_enc(z_final)                  # [B, D//2] concept codes
        z_head = z_final + self.sae_dec(c)         # [B, D] with SAE reconstruction
        z_head = self.head_drop(z_head)            # Regularization

        # Temperature calibration parameter (τ ∈ (0.5, ∞))
        tau_raw = self.temp_head(z_head)
        tau = 0.5 + 0.5 * F.softplus(tau_raw)  # Lower bound at 0.5

        return z_final, z_head, halt_logits, tau, c

    def forward(
        self,
        x: torch.Tensor,                    # Goal embedding [B, D]
        y: torch.Tensor,                    # Response embedding [B, D]
        z: torch.Tensor,                    # Initial latent state [B, D]
        *,
        seq_len: Optional[torch.Tensor] = None,  # Response length [B] (optional)
        return_aux: bool = True,                 # Whether to return auxiliary outputs
        with_consistency_target: bool = True,    # Compute consistency regularization
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Dict[str, Any]]:
        """
        Complete forward pass with recursive processing and multi-head prediction.
        """
        # Main recursive processing
        z = z.clone()  # Ensure we don't modify input
        z_final, z_head, halt_logits, tau, c = self._recur(x, y, z)

        # Core prediction heads
        logits = self.classifier(z_head)                    # [B, vocab_size]
        score_logit = self.score_head(z_head)               # [B, 1]
        log_var = self.logvar_head(z_head)                  # [B, 1] uncertainty

        # ----- NUMERICAL SAFETY -----
        LOGVAR_MIN, LOGVAR_MAX = -5.0, 5.0
        log_var = log_var.clamp(min=LOGVAR_MIN, max=LOGVAR_MAX)

        # Use tau for calibration; keep a stable proxy for telemetry
        # NOTE: move temp01 to sigmoid(tau_raw) for cross-model alignment
        tau_raw = self.temp_head(z_head) 
        tau = 0.5 + 0.5 * F.softplus(tau_raw)
        tau_safe = torch.clamp(tau, min=1e-2)
        s = torch.sigmoid(score_logit / tau_safe)

        # ----- Core auxiliaries
        aux3_logits = self.aux3_head(z_head)
        aux3_probs  = F.softmax(aux3_logits, dim=-1)
        disagree_logit = self.disagree_head(z_head)
        y_recon     = self.recon_head(z_head)
        ood_logit   = self.ood_head(z_head)

        # Optional bridge heads
        agree01 = torch.sigmoid(self.agree_head(z_head)) if self.enable_agree_head else None
        sens01  = torch.sigmoid(self.causal_sens_head(z_head)) if self.enable_causal_sens_head else None

        # Consistency target 
        mask = (torch.rand_like(z_head) < self.consistency_mask_p).float()
        z_masked = z_head * (1.0 - mask)
        cos_consistency = self._cos01(z_head, z_masked).unsqueeze(-1)
        consistency_logit = self.consistency_head(z_head)

        # Finite-difference sensitivity
        eps = 1e-3
        y_eps = y + eps * F.normalize(torch.randn_like(y), dim=-1)
        with torch.no_grad():
            _, z_head_eps, _, tau_eps, _ = self._recur(x, y_eps, z)
        tau_eps_safe = torch.clamp(tau_eps, min=1e-2)
        score_eps    = torch.sigmoid(self.score_head(z_head_eps) / tau_eps_safe)
        jac_fd       = ((score_eps - s).abs() / eps).clamp(0, 10.0) / 10.0

        # Length effect
        if seq_len is not None:
            len_effect = torch.tanh((seq_len.float() / self.len_norm_L)).unsqueeze(-1)
        else:
            len_effect = torch.zeros_like(s)
        length_norm01 = (len_effect + 1.0) * 0.5

        # ----- Aligned telemetry keys -----
        certainty01   = torch.sigmoid(-log_var)
        uncertainty01 = 1.0 - certainty01
        temp01        = torch.sigmoid(tau_raw)  # aligned proxy in [0,1]
        ood_hat01     = torch.sigmoid(ood_logit)
        halt_prob     = torch.sigmoid(halt_logits).unsqueeze(-1) if halt_logits.dim()==1 else torch.sigmoid(halt_logits)

        # Device-safe normalized entropy (in [0,1])
        logK = torch.log(torch.tensor(3.0, device=z_head.device, dtype=z_head.dtype))
        entropy_aux = (-(aux3_probs * F.log_softmax(aux3_logits, -1)).sum(-1) / logK).unsqueeze(-1)

        aux: Dict[str, Any] = {
            # raw heads you need for training
            "score_logit": score_logit,
            "log_var": log_var,
            "aux3_logits": aux3_logits,
            "disagree_logit": disagree_logit,
            "y_recon": y_recon,
            "consistency_logit": consistency_logit,
            "consistency_target": cos_consistency.detach(),

            # aligned derived telemetry (all ∈ [0,1])
            "score": s,
            "certainty01": certainty01,
            "uncertainty01": uncertainty01,     # <  NEW (correct)
            "uncertainty": uncertainty01,       # <  OPTIONAL alias for back-compat
            "aux3_probs": aux3_probs,
            "entropy_aux": entropy_aux,
            "disagree_hat": torch.sigmoid(disagree_logit),
            "recon_sim": self._cos01(y_recon, y).unsqueeze(-1),
            "consistency_hat": torch.sigmoid(consistency_logit),
            "concept_sparsity": (c > 0).float().mean(dim=-1, keepdim=True),
            "ood_hat01": ood_hat01,             # <  NEW aligned name
            "temp01": temp01,                   # <  changed to sigmoid(tau_raw)
            "jacobian_fd": jac_fd,
            "len_effect": len_effect,
            "length_norm01": length_norm01,     # <  NEW 0..1 length proxy
            "halt_prob": halt_prob,             # <  NEW
        }

        if agree01 is not None:
            aux["agree01"] = agree01
        if sens01 is not None:
            aux["sens01"] = sens01

        return logits, halt_logits.squeeze(-1), z_final, (aux if return_aux else {})

# stephanie/scoring/training/tiny_recursion_trainer.py
"""
TinyRecursionModel Trainer (Tiny+)

A specialized trainer for the TinyRecursionModel that implements multi-objective
training with heteroscedastic regression and auxiliary losses. This trainer handles
multiple data schemas and produces dimension-specific models with comprehensive
training telemetry.

Key Features:
- Heteroscedastic regression for score prediction with uncertainty estimation
- Multiple auxiliary objectives: bucket classification, disagreement, reconstruction
- Support for various input schemas (native, singleton, pairwise, HRM)
- Comprehensive training monitoring and validation
- Early stopping and model checkpointing

"""

from __future__ import annotations

import math
import os
from collections import Counter
from datetime import datetime
from typing import Any, Dict, List, Optional, Tuple
import logging
import torch
import torch.nn.functional as F
from torch import optim
from torch.utils.data import DataLoader, TensorDataset

from stephanie.scoring.model.tiny_recursion import TinyRecursionModel
from stephanie.scoring.training.base_trainer import BaseTrainer

try:
    from tqdm.auto import tqdm
except Exception:  # pragma: no cover
    tqdm = None

_logger = logging.getLogger(__name__)

def _bucket3(y01: torch.Tensor) -> torch.Tensor:
    """
    Convert continuous scores to 3-class bucket labels.
    
    Args:
        y01: Tensor of scores in range [0, 1]
        
    Returns:
        Long tensor with bucket indices:
        - 0: scores < 1/3
        - 1: scores in [1/3, 2/3)
        - 2: scores >= 2/3
    """
    # <1/3 => 0, [1/3,2/3) => 1, >=2/3 => 2
    edges = torch.tensor([1/3, 2/3], device=y01.device, dtype=y01.dtype)
    return (y01 >= edges[0]).long() + (y01 >= edges[1]).long()


class TinyTrainer(BaseTrainer):
    """
    Trainer for TinyRecursionModel (Tiny+) with multi-objective optimization.
    
    This trainer implements a comprehensive training regimen that combines:
    - Main heteroscedastic regression objective
    - Multiple auxiliary objectives for regularization and feature learning
    - Support for various input data formats And see this was a complete waste of timed
    - Extensive monitoring and validation
    
    The model produces separate instances for each quality dimension.
    
    Attributes:
        model_type: Identifier for model architecture ("tiny")
        target_type: Type of scoring target ("document", "sentence", etc.)
        version: Model version identifier
        epochs: Number of training epochs
        lr: Learning rate for optimizer
        batch_size: Training batch size
        dropout: Dropout rate for model regularization
        use_attention: Whether to use attention mechanisms
        n_recursions: Number of recursion steps in model
        halt_lambda: Weight for halting regularization loss
        grad_clip: Gradient clipping value
        w_aux3: Weight for 3-class auxiliary classification
        w_disagree: Weight for disagreement prediction
        w_recon: Weight for reconstruction loss
        w_cons: Weight for consistency regularization
        w_sae_recon: Weight for sparse autoencoder reconstruction
        w_ood: Weight for out-of-distribution detection
    """

    def __init__(self, cfg, memory, container, logger):
        """Initialize TinyTrainer with configuration and dependencies."""
        super().__init__(cfg, memory, container, logger)

        # --- Identity / paths -------------------------------------------------
        self.model_type   = "tiny"
        self.target_type  = cfg.get("target_type", "document")
        self.version      = cfg.get("model_version", "v1")

        # --- Core knobs -------------------------------------------------------
        self.epochs        = int(cfg.get("epochs", 20))
        self.lr            = float(cfg.get("lr", 3e-5))           # conservative default
        self.batch_size    = int(cfg.get("batch_size", 16))
        self.dropout       = float(cfg.get("dropout", 0.1))
        self.use_attention = bool(cfg.get("use_attention", False))
        self.n_recursions  = int(cfg.get("n_recursions", 6))
        self.halt_lambda   = float(cfg.get("halt_lambda", 0.05))  # halting is a light regularizer
        self.grad_clip     = float(cfg.get("grad_clip", 0.5))

        # Aux loss weights
        self.w_aux3        = float(cfg.get("w_aux3", 0.3))
        self.w_disagree    = float(cfg.get("w_disagree", 0.3))
        self.w_recon       = float(cfg.get("w_recon", 0.2))
        self.w_cons        = float(cfg.get("w_consistency", 0.2))
        self.w_sae_recon   = float(cfg.get("w_sae_recon", 0.0))   # 0 = off by default
        self.w_ood         = float(cfg.get("w_ood", 0.0))         # 0 = off by default

        # --- Telemetry --------------------------------------------------------
        self.show_progress       = bool(cfg.get("show_progress", True))
        self.progress_every      = max(1, int(cfg.get("progress_every", 500)))
        self.log_every_steps     = max(1, int(cfg.get("log_every_steps", 50)))
        self.label_hist_bucket   = int(cfg.get("label_hist_bucket", 10))
        self.log_label_histogram = bool(cfg.get("log_label_histogram", True))

        # --- Validation / reproducibility ------------------------------------
        self.validation_ratio = float(cfg.get("validation_ratio", 0.1))
        self.seed             = int(cfg.get("seed", 42))
        torch.manual_seed(self.seed)
        if torch.cuda.is_available():
            torch.cuda.manual_seed(self.seed)

        # --- Model ------------------------------------------------------------
        self.model = TinyRecursionModel(
            d_model=self.dim,
            n_layers=int(cfg.get("n_layers", 2)),
            n_recursions=self.n_recursions,
            vocab_size=int(cfg.get("vocab_size", 101)),   # kept for classifier compatibility
            use_attention=self.use_attention,
            dropout=self.dropout,
            attn_heads=int(cfg.get("attn_heads", 4)),
            step_scale=float(cfg.get("step_scale", 0.1)),
            consistency_mask_p=float(cfg.get("consistency_mask_p", 0.10)),
            len_norm_L=float(cfg.get("len_norm_L", 512.0)),
        ).to(self.device)

    # ------------------------------
    # Data prep
    # ------------------------------

    def _create_dataloader(self, samples: List[Dict[str, Any]]) -> Tuple[Optional[DataLoader], int, int]:
        """
        Create DataLoader from sample dictionaries with multiple schema support.
        
        Supports multiple input formats:
        - Native Tiny+ schema: x, y, z, target
        - Singleton format: goal_text/output with score
        - Pairwise format: output_a/output_b with comparative scores
        - HRM format: goal_text/scorable_text with target_score
        
        Args:
            samples: List of sample dictionaries with various possible schemas
            
        Returns:
            Tuple of (DataLoader, kept_count, dropped_count) or (None, kept, dropped) if insufficient data
        """
        xs, ys, zs = [], [], []
        y01, halt_targets, seq_lens = [], [], []
        kept = dropped = 0
        label_counts = Counter()

        use_tqdm = bool(self.show_progress and tqdm is not None)
        it = tqdm(samples, desc="Packing Tiny+ samples", unit="samp") if use_tqdm else samples

        def _push(goal: str, doc: str, target: float, *, z_text: Optional[str] = None, halt_t: float = 1.0, slen: int = 0):
            """Internal helper to process and validate a single sample."""
            nonlocal kept, dropped
            try:
                # Get embeddings for text inputs
                x = torch.tensor(self.memory.embedding.get_or_create(goal), dtype=torch.float32, device=self.device)
                y = torch.tensor(self.memory.embedding.get_or_create(doc),  dtype=torch.float32, device=self.device)
                z = torch.tensor(self.memory.embedding.get_or_create(z_text if z_text is not None else goal),
                                dtype=torch.float32, device=self.device)

                # ---- Normalize & sanitize inputs (prevents recursion amplification / NaNs)
                def _safe_vec(t):
                    """Safely normalize vector, handling NaN/inf values."""
                    t = torch.nan_to_num(t, nan=0.0, posinf=0.0, neginf=0.0)
                    norm = t.norm(dim=-1, keepdim=True).clamp_min(1e-6)
                    return t / norm

                x = _safe_vec(x); y = _safe_vec(y); z = _safe_vec(z)
                if not torch.isfinite(x).all() or not torch.isfinite(y).all() or not torch.isfinite(z).all():
                    dropped += 1
                    return

                # normalize target → [0,1]
                t = float(target)
                t = (max(0.0, min(100.0, t)) / 100.0) if t > 1.0 else max(0.0, min(1.0, t))

                xs.append(x); ys.append(y); zs.append(z)
                y01.append(t); halt_targets.append(float(halt_t)); seq_lens.append(int(slen))
                label_counts[int(round(t * 100))] += 1
                kept += 1
            except Exception as e:
                dropped += 1
                if self.logger: self.logger.log("TinyRecursionSampleError", {"error": str(e)})

        # Process all samples with schema detection
        for s in it:
            # Native Tiny+ schema
            if "x" in s and "y" in s and "z" in s and "target" in s:
                _push(s["x"], s["y"], s["target"], z_text=s.get("z"), halt_t=s.get("halt_target", 1.0), slen=s.get("seq_len", 0))
                continue

            # Singleton (SICQL/MRQ style)
            title = (s.get("goal_text") or s.get("title") or "").strip()
            if "output" in s and ("score" in s or "target_score" in s):
                out = (s.get("scorable_text") or s.get("output") or "").strip()
                val = s.get("target_score", s.get("score"))
                if title and out and (val is not None):
                    _push(title, out, val, z_text=title)
                else:
                    dropped += 1
                continue

            # Pairwise
            if all(k in s for k in ("output_a","output_b","value_a","value_b")):
                a_out = (s.get("output_a") or "").strip()
                b_out = (s.get("output_b") or "").strip()
                a_val = s.get("value_a"); b_val = s.get("value_b")
                if title:
                    if a_out and a_val is not None: _push(title, a_out, a_val, z_text=title)
                    if b_out and b_val is not None: _push(title, b_out, b_val, z_text=title)
                else:
                    dropped += 1
                continue

            # HRM/raw
            if ("goal_text" in s and "scorable_text" in s and ("target_score" in s or "score" in s)):
                out = (s.get("scorable_text") or "").strip()
                val = s.get("target_score", s.get("score"))
                _push(title, out, val, z_text=title)
                continue

            dropped += 1

            if use_tqdm: it.set_postfix(kept=kept, drop=dropped)

        if use_tqdm and hasattr(it, "close"): it.close()

        # Log label distribution for analysis
        if self.logger and self.log_label_histogram:
            exact = {int(k): int(v) for k, v in sorted(label_counts.items())}
            bucketed = self._bucketize_counts(label_counts, self.label_hist_bucket)
            self.logger.log("TinyPlusLabelHistogram", {
                "kept": int(kept), "dropped": int(dropped),
                "exact": exact, "bucket_size": int(self.label_hist_bucket),
                "bucketed": bucketed
            })

        if kept < self.min_samples:
            return None, kept, dropped

        # Create TensorDataset and DataLoader
        dataset = TensorDataset(
            torch.stack(xs), torch.stack(ys), torch.stack(zs),
            torch.tensor(y01, dtype=torch.float32, device=self.device).unsqueeze(-1),
            torch.tensor(halt_targets, dtype=torch.float32, device=self.device).unsqueeze(-1),
            torch.tensor(seq_lens, dtype=torch.int32, device=self.device),
        )
        loader = DataLoader(dataset, batch_size=self.batch_size, shuffle=True)
        return loader, kept, dropped

    # ------------------------------
    # Loss Functions
    # ------------------------------

    @staticmethod
    def _heteroscedastic_regression_loss(score: torch.Tensor, target01: torch.Tensor, log_var: torch.Tensor) -> torch.Tensor:
        """
        Compute heteroscedastic regression loss with uncertainty estimation.
        
        This loss adapts to the uncertainty in predictions by learning
        a variance term that scales the regression loss.
        
        Args:
            score: Predicted scores [B, 1]
            target01: Ground truth scores in [0, 1] [B, 1]
            log_var: Learned log variance [B, 1]
            
        Returns:
            Scalar loss value
        """
        log_var = log_var.clamp(-5.0, 5.0)  # defensive clamp to avoid precision explosion
        inv_var = torch.exp(-log_var)
        diff2   = (score - target01).pow(2)
        return (inv_var * diff2 + log_var).mean()

    @staticmethod
    def _cosine_recon_loss(y_recon: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
        """
        Compute cosine reconstruction loss.
        
        Measures how well the model can reconstruct the input embedding,
        encouraging meaningful internal representations.
        
        Args:
            y_recon: Reconstructed embedding
            y_true: Original embedding
            
        Returns:
            Cosine distance loss in range [0, 1]
        """
        # 1 - cosine in [0,2] → clamp to [0,1]
        cos = F.cosine_similarity(y_recon, y_true, dim=-1, eps=1e-8).unsqueeze(-1)
        return (1 - cos).clamp(0, 1).mean()

    # ------------------------------
    # Epoch training
    # ------------------------------

    def _train_epoch(self, model: TinyRecursionModel, dataloader: DataLoader, epoch_idx: int) -> float:
        """
        Train model for one epoch.
        
        Args:
            model: TinyRecursionModel instance
            dataloader: Training data loader
            epoch_idx: Current epoch index
            
        Returns:
            Average training loss for the epoch
        """
        model.train()
        total_loss = 0.0
        count = 0

        use_tqdm = bool(self.show_progress and tqdm is not None)
        it = tqdm(dataloader, desc=f"Epoch {epoch_idx}", unit="batch", leave=False) if use_tqdm else dataloader

        for step, batch in enumerate(it, start=1):
            x, y, z, target01, halt_target, seq_len = batch

            # Forward pass with auxiliary outputs
            logits, halt_logits, _, aux = model(x, y, z, seq_len=seq_len, return_aux=True)

            # Main loss: heteroscedastic regression on score/log_var
            L_main = self._heteroscedastic_regression_loss(aux["score"], target01, aux["log_var"])

            # Auxiliary losses for multi-objective training
            buckets = _bucket3(target01.squeeze(-1))
            L_aux3  = F.cross_entropy(aux["aux3_logits"], buckets)  # 3-class classification
            L_dis   = F.smooth_l1_loss(aux["disagree_hat"], (target01 - aux["score"].detach()).abs())  # Disagreement prediction
            L_recon = self._cosine_recon_loss(aux["y_recon"], y)  # Reconstruction quality
            L_cons  = F.mse_loss(aux["consistency_hat"], aux["consistency_target"])  # Consistency regularization

            # Optional losses (weight=0 means disabled)
            L_sae = torch.zeros((), device=self.device)
            if self.w_sae_recon > 0.0 and "concept_vec" in aux:
                L_sae = aux["concept_vec"].abs().mean()  # Sparse autoencoder reconstruction

            L_ood = torch.zeros((), device=self.device)
            if self.w_ood > 0.0 and "ood_hat" in aux:
                L_ood = F.binary_cross_entropy(aux["ood_hat"], torch.ones_like(aux["ood_hat"]))  # OOD detection

            L_halt = F.binary_cross_entropy_with_logits(halt_logits.unsqueeze(-1), halt_target)  # Halting regularization

            # Check components for finiteness & sanity
            all_terms = torch.stack([
                L_main.detach(),
                L_aux3.detach(),
                L_dis.detach(),
                L_recon.detach(),
                L_cons.detach(),
                L_sae.detach(),
                L_ood.detach(),
                L_halt.detach()
            ])
            if (not torch.isfinite(all_terms).all()) or (all_terms.abs().max() > 1e6):
                if self.logger:
                    self.logger.log("TinyPlusNaNBatch", {
                        "epoch": epoch_idx,
                        "step": step,
                        "any_nan": bool(not torch.isfinite(all_terms).all()),
                        "max_abs": float(all_terms.abs().max().item())
                    })
                self.optimizer.zero_grad(set_to_none=True)
                continue  # skip this batch

            # Combined loss with weighting
            loss = (
                L_main
                + self.w_aux3 * L_aux3
                + self.w_disagree * L_dis
                + self.w_recon * L_recon
                + self.w_cons * L_cons
                + self.w_sae_recon * L_sae
                + self.w_ood * L_ood
                + self.halt_lambda * L_halt
            )

            if (not torch.isfinite(loss)) or (abs(loss.item()) > 1e7):
                if self.logger:
                    self.logger.log("TinyPlusUnstableLoss", {
                        "epoch": epoch_idx,
                        "step": step,
                        "loss": float(loss.item()) if torch.isfinite(loss) else float('nan')
                    })
                self.optimizer.zero_grad(set_to_none=True)
                continue

            # Backward pass with gradient clipping
            self.optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), self.grad_clip)
            self.optimizer.step()

            bsz = x.size(0)
            total_loss += loss.item() * bsz
            count += bsz

            # Progress reporting
            if use_tqdm:
                it.set_postfix(loss=f"{loss.item():.4f}")
            elif self.logger and (step % self.log_every_steps == 0):
                self.logger.log("TinyPlusBatch", {
                    "epoch": epoch_idx,
                    "step": step,
                    "loss": float(loss.item()),
                    "L_main": float(L_main.item()),
                    "L_aux3": float(L_aux3.item()),
                    "L_dis": float(L_dis.item()),
                    "L_recon": float(L_recon.item()),
                    "L_cons": float(L_cons.item()),
                    "L_sae": float(L_sae.item()),
                    "L_ood": float(L_ood.item()),
                    "L_halt": float(L_halt.item()),
                })

        if use_tqdm and hasattr(it, "close"):
            it.close()

        return total_loss / max(1, count)

    # ------------------------------
    # Validation
    # ------------------------------

    @torch.no_grad()
    def _validate(self, model: TinyRecursionModel, dataloader: Optional[DataLoader]) -> Dict[str, float]:
        """
        Run validation and compute comprehensive metrics.
        
        Args:
            model: Model to validate
            dataloader: Validation data loader
            
        Returns:
            Dictionary of validation metrics
        """
        if not dataloader:
            return {}

        model.eval()
        scores, targets = [], []
        # 10 metric lists for comprehensive validation
        entropies, uncerts, disagree, recon_sim, cons_hat, temp01, jac, spars, ood, len_eff = (
            [] for _ in range(10)
        )

        for x, y, z, target01, _, seq_len in dataloader:
            _, _, _, aux = model(x, y, z, seq_len=seq_len, return_aux=True)
            s = aux["score"].detach().cpu().view(-1)
            t = target01.detach().cpu().view(-1)

            scores.append(s)
            targets.append(t)
            # Collect various auxiliary metrics for analysis
            entropies.append(aux["entropy_aux"].detach().cpu().view(-1))
            uncerts.append(aux["uncertainty"].detach().cpu().view(-1))
            disagree.append(aux["disagree_hat"].detach().cpu().view(-1))
            recon_sim.append(aux["recon_sim"].detach().cpu().view(-1))
            cons_hat.append(aux["consistency_hat"].detach().cpu().view(-1))
            temp01.append(aux["temp01"].detach().cpu().view(-1))
            jac.append(aux["jacobian_fd"].detach().cpu().view(-1))
            spars.append(aux["concept_sparsity"].detach().cpu().view(-1))
            ood.append(aux["ood_hat"].detach().cpu().view(-1))
            len_eff.append(aux["len_effect"].detach().cpu().view(-1))

        s = torch.cat(scores); t = torch.cat(targets)
        mae = F.l1_loss(s, t).item()
        rmse = torch.sqrt(F.mse_loss(s, t)).item()

        def mean_cat(arrs):
            """Helper to compute mean of concatenated tensor list."""
            return float(torch.cat(arrs).mean().item()) if arrs else 0.0

        return {
            "mae": mae,
            "rmse": rmse,
            "entropy_aux_mean":     mean_cat(entropies),
            "uncertainty_mean":     mean_cat(uncerts),
            "disagree_hat_mean":    mean_cat(disagree),
            "recon_sim_mean":       mean_cat(recon_sim),
            "consistency_hat_mean": mean_cat(cons_hat),
            "temp01_mean":          mean_cat(temp01),
            "jacobian_fd_mean":     mean_cat(jac),
            "concept_sparsity_mean":mean_cat(spars),
            "ood_hat_mean":         mean_cat(ood),
            "len_effect_mean":      mean_cat(len_eff),
        }

    # ------------------------------
    # Train/val split
    # ------------------------------

    def _create_train_val_split(self, samples: List[Dict[str, Any]]):
        """
        Split samples into training and validation sets.
        
        Args:
            samples: List of sample dictionaries
            
        Returns:
            Tuple of (train_samples, val_samples)
        """
        if not samples:
            return [], []
        if self.validation_ratio <= 0 or len(samples) < 10:
            return samples, []
        g = torch.Generator().manual_seed(self.seed)
        idx = torch.randperm(len(samples), generator=g).tolist()
        split = int(len(samples) * (1 - self.validation_ratio))
        return [samples[i] for i in idx[:split]], [samples[i] for i in idx[split:]]

    # ------------------------------
    # Main train loop (per dimension)
    # ------------------------------

    def train(self, samples, dimension):
        """
        Main training loop for a specific quality dimension.
        
        Args:
            samples: Training samples for the dimension
            dimension: Quality dimension name
            
        Returns:
            Training results dictionary
        """
        # Split data
        train_samples, val_samples = self._create_train_val_split(samples)

        # Create data loaders
        dataloader, kept, dropped = self._create_dataloader(train_samples)
        val_loader, val_kept, val_dropped = (None, 0, 0)
        if val_samples:
            val_loader, val_kept, val_dropped = self._create_dataloader(val_samples)

        if not dataloader:
            return {"error": "insufficient_data", "kept": kept, "dropped": dropped}

        # Optimizer
        self.optimizer = optim.AdamW(self.model.parameters(), lr=self.lr, weight_decay=1e-2)

        best_metric = float("inf")
        patience, wait = int(self.cfg.get("patience", 3)), 0
        train_losses: List[float] = []
        saved_best = False

        locator = self.get_locator(dimension)  # create once; base_path will be ensured

        # Training loop with early stopping
        for epoch in range(1, self.epochs + 1):
            avg_loss = self._train_epoch(self.model, dataloader, epoch_idx=epoch)
            avg_loss = float(avg_loss)
            # Ensure epoch loss is finite for serialization/meta
            if not math.isfinite(avg_loss):
                avg_loss = float(train_losses[-1]) if train_losses else 0.0
                if self.logger:
                    self.logger.log("TinyPlusNaNEpoch", {"epoch": epoch})
            train_losses.append(avg_loss)

            # Validation
            val_metrics = self._validate(self.model, val_loader) if val_loader else {}
            if self.logger:
                payload = {"epoch": epoch, "train_loss": float(avg_loss)}
                payload.update({f"val_{k}": v for k, v in val_metrics.items()})
                self.logger.log("TinyPlusEpoch", payload)

            # Early stopping metric: prefer val MAE, fallback to train loss
            stop_metric = val_metrics.get("mae", avg_loss) if val_metrics else avg_loss
            if not math.isfinite(stop_metric):
                if self.logger:
                    self.logger.log("TinyPlusNonFiniteMetric", {"epoch": epoch})
                stop_metric = float("inf")

            improved = (not math.isfinite(best_metric)) or (stop_metric < best_metric - 1e-6)
            if improved:
                best_metric = stop_metric
                wait = 0
                best_path = locator.model_file(suffix="_tiny.pt")
                try:
                    torch.save(self.model.state_dict(), best_path)
                    saved_best = True
                    if self.logger:
                        self.logger.log("TinyPlusSaveCheckpoint", {"epoch": epoch, "path": best_path, "metric": float(best_metric)})
                except Exception as e:
                    if self.logger:
                        self.logger.log("TinyPlusSaveError", {"epoch": epoch, "path": best_path, "error": str(e)})
            else:
                wait += 1
                if wait >= patience:
                    if self.logger:
                        self.logger.log("TinyPlusEarlyStopping", {"epoch": epoch, "best_metric": float(best_metric)})
                    break

        # ---- ALWAYS save a 'last' checkpoint ----
        last_path = locator.model_file(suffix="_tiny_last.pt")
        try:
            torch.save(self.model.state_dict(), last_path)
            if self.logger:
                self.logger.log("TinyPlusSaveLast", {"path": last_path})
        except Exception as e:
            if self.logger:
                self.logger.log("TinyPlusSaveLastError", {"path": last_path, "error": str(e)})

        # If no 'best' was saved during training, backfill it now:
        best_path = locator.model_file(suffix="_tiny.pt")
        if not saved_best or not os.path.exists(best_path):
            try:
                torch.save(self.model.state_dict(), best_path)
                if self.logger:
                    self.logger.log("TinyPlusBackfillBest", {"path": best_path})
            except Exception as e:
                if self.logger:
                    self.logger.log("TinyPlusBackfillBestError", {"path": best_path, "error": str(e)})

        # --- Save training metadata -------------------------------------------
        safe_config = {
            "lr": self.lr,
            "epochs": self.epochs,
            "batch_size": self.batch_size,
            "halt_lambda": self.halt_lambda,
            "n_layers": self.cfg.get("n_layers", 2),
            "n_recursions": self.n_recursions,
            "use_attention": self.use_attention,
            "dropout": self.dropout,
            "seed": self.seed,
            "vocab_size": int(self.cfg.get("vocab_size", 101)),
            "w_aux3": self.w_aux3,
            "w_disagree": self.w_disagree,
            "w_recon": self.w_recon,
            "w_consistency": self.w_cons,
            "w_sae_recon": self.w_sae_recon,
            "w_ood": self.w_ood,
            "grad_clip": self.grad_clip,
        }

        # Ensure train_loss_curve is finite-only floats
        finite_curve = []
        last_finite = 0.0
        for v in train_losses:
            if math.isfinite(v):
                last_finite = float(v)
            finite_curve.append(float(last_finite))

        meta = {
            "dimension": dimension,
            "model_type": "tiny_recursion",
            "expects_triplet": True,
            "embedding_type": self.embedding_type,
            "input_dim": self.dim,
            "concat_input_dim": self.dim * 2,
            "version": self.version,
            "epochs": self.epochs,
            "avg_loss": float(min(finite_curve or [best_metric])),
            "timestamp": datetime.now().isoformat(),
            "cfg": dict(self.cfg),
            "kept": int(kept),
            "best_metric": float(best_metric),
            "train_loss_curve": [float(x) for x in finite_curve],
            "dropped": int(dropped),
            "val_kept": int(val_kept),
            "val_dropped": int(val_dropped),
        }
        self._save_meta_file(meta, dimension)

        # TrainingStatsStore integration
        self.memory.training_stats.add_from_result(
            stats={
                "avg_q_loss": float(min(finite_curve or [best_metric])),
                "avg_loss":   float(min(finite_curve or [best_metric])),
            },
            model_type=self.model_type,
            target_type=self.target_type,
            dimension=dimension,
            version=self.version,
            embedding_type=self.embedding_type,
            config=safe_config,
            sample_count=len(samples),
            valid_samples=int(kept),
            invalid_samples=int(dropped),
            start_time=datetime.now(),
        )

        return {
            "best_metric": float(best_metric),
            "train_loss_curve": [float(x) for x in finite_curve],
            "kept": int(kept),
            "dropped": int(dropped),
            "val_kept": int(val_kept),
            "val_dropped": int(val_dropped),
        }

    # ------------------------------
    # Helper Methods
    # ------------------------------

    def _bucketize_label(self, y: int) -> int:
        """
        Bucketize label for histogram analysis.
        
        If bucket_size > 0, map 0..100 → 0..num_bins-1 using fixed-width bins.
        Ensure vocab_size >= num_bins when you enable bucketing.
        
        Args:
            y: Original label value
            
        Returns:
            Bucket index
        """
        b = int(self.bucket_size)
        if b <= 0:
            return max(0, min(100, y))              # <<< clamp to 0..100
        num_bins = (101 + b - 1) // b               # e.g., b=10 -> 11 bins (0..10)
        yb = min(max(y, 0) // b, num_bins - 1)
        return yb

    def _bucketize_counts(self, counts: Counter, bucket: int) -> dict:
        """
        Convert exact label counts to bucketized counts for visualization.
        
        Args:
            counts: Counter of exact label values
            bucket: Bucket size
            
        Returns:
            Dictionary mapping bucket ranges to counts
        """
        if bucket <= 1:
            return {str(k): int(v) for k, v in sorted(counts.items())}

        buckets = {}
        for label, c in counts.items():
            try:
                l = int(label)
            except Exception:
                continue
            start = (l // bucket) * bucket
            end = min(100, start + bucket - 1)
            key = f"{start}-{end}"
            buckets[key] = buckets.get(key, 0) + int(c)

        # Ensure all possible buckets are represented
        start = 0
        while start <= 100:
            end = min(100, start + bucket - 1)
            key = f"{start}-{end}"
            buckets.setdefault(key, 0)
            start += bucket

        return dict(sorted(buckets.items(), key=lambda kv: int(kv[0].split('-')[0])))

# stephanie/scoring/tiny_scorer.py
"""
Tiny Recursion Model Scorer - Lightweight evaluator with rich diagnostics.

This module implements the scoring interface for Tiny Recursion Models (TRM),
providing fast, recursive quality assessment with comprehensive diagnostic
telemetry. The scorer transforms TRM's internal signals into standardized
Shared Core Metrics (SCM) format for cross-model comparison in GAP analysis.

Key Features:
- Per-dimension model loading and management
- Rich diagnostic extraction (uncertainty, OOD, sensitivity, agreement, etc.)
- SCM alignment for cross-model comparability
- Vector generation for topological analysis
- Flexible attribute verbosity levels (minimal/standard/full)

The TinyScorer serves as the lightweight counterpart to HRM in the GAP
analysis pipeline, enabling efficient model comparison and routing decisions.
"""

from __future__ import annotations

import logging
import os
from typing import Any, Dict, List

import torch

from stephanie.constants import GOAL, GOAL_TEXT
from stephanie.data.score_bundle import ScoreBundle
from stephanie.data.score_result import ScoreResult
from stephanie.scoring.model.tiny_recursion import TinyRecursionModel
from stephanie.scoring.scorer.base_scorer import BaseScorer
from stephanie.utils.file_utils import load_json

_logger = logging.getLogger(__name__)


class TinyScorer(BaseScorer):
    """
    Tiny Recursion Model scorer for efficient quality evaluation with rich diagnostics.
    
    This scorer uses trained TinyRecursionModel instances to evaluate goal-response
    pairs across multiple reasoning dimensions. It extracts not just quality scores
    but comprehensive diagnostic telemetry including uncertainty estimates,
    out-of-distribution detection, sensitivity analysis, and agreement predictions.
    
    The scorer automatically converts TRM's native outputs into the standardized
    Shared Core Metrics (SCM) format, enabling direct comparison with HRM and
    other evaluation systems in the GAP analysis pipeline.
    
    Attributes:
        model_type: Identifier for scorer type ("tiny")
        embedding_type: Type of embeddings used (shared with HRM)
        dimensions: List of reasoning dimensions to evaluate
        attr_level: Verbosity level for attributes ("minimal"/"standard"/"full")
        models: Dictionary of loaded TRM models per dimension
        model_meta: Metadata for each dimension's model
    """
    
    def __init__(self, cfg, memory, container, logger):
        """
        Initialize TinyScorer with configuration and dependencies.
        
        Args:
            cfg: Configuration dictionary with scorer parameters
            memory: Memory interface for embedding and data access
            container: Dependency injection container
            logger: Structured logging interface
            
        Configuration Parameters:
            target_type: Type of scoring target ("conversation_turn")
            model_path: Base path for model files
            model_version: Version identifier for models
            dimensions: List of dimensions to evaluate
            clip_0_100: Whether to clip scores to 0-100 range
            tiny_attr_level: Attribute verbosity level
        """
        super().__init__(cfg, memory, container, logger)
        _logger.info("Initializing TinyScorer")
        
        self.model_type = "tiny"  # identifies scorer type in results

        # Embedding interface (shared with HRM for cross-model alignment)
        self.embedding_type = self.memory.embedding.name
        self.dim = self.memory.embedding.dim
        _logger.debug(f"Using embedding type: {self.embedding_type}, dimension: {self.dim}")

        # Configuration parameters
        self.target_type = cfg.get("target_type", "conversation_turn")
        self.model_path = cfg.get("model_path", "models")
        self.version = cfg.get("model_version", "v1")
        self.dimensions: List[str] = cfg.get("dimensions", [])
        
        # Output scaling configuration
        self.clip_0_100 = cfg.get("clip_0_100", True)
        
        # Attribute verbosity: controls diagnostic detail level
        self.attr_level = (cfg.get("tiny_attr_level") or "standard").lower()
        _logger.debug(f"Attribute level set to: {self.attr_level}")

        # Containers for per-dimension models and metadata
        self.models: Dict[str, TinyRecursionModel] = {}
        self.model_meta: Dict[str, Dict[str, Any]] = {}

        # Attempt to load models up-front for all specified dimensions
        _logger.info(f"Loading TRM models for dimensions: {self.dimensions}")
        self._load_models(self.dimensions)
        _logger.info(f"TinyScorer initialized with {len(self.models)} loaded models")

    # -------------------------
    # Model Loading
    # -------------------------
    def _load_models(self, dimensions: List[str]):
        """
        Load trained TinyRecursionModel instances for specified dimensions.
        
        For each dimension, this method:
        1. Resolves model and metadata file paths
        2. Loads model configuration from metadata
        3. Instantiates TRM with correct architecture
        4. Loads trained weights
        5. Registers model in the internal registry
        
        Args:
            dimensions: List of reasoning dimensions to load models for
            
        Logs:
            - Debug: Model loading progress and configuration
            - Warning: Missing model files or metadata
            - Error: Model instantiation or weight loading failures
        """
        _logger.debug(f"Starting model loading for {len(dimensions)} dimensions")
        
        for dim in dimensions:
            _logger.debug(f"Loading model for dimension: {dim}")
            locator = self.get_locator(dim)

            # Resolve model and metadata file paths
            model_path = locator.model_file(suffix="_tiny.pt")
            meta_path = locator.meta_file()
            _logger.debug(f"Model path: {model_path}, Meta path: {meta_path}")

            if not os.path.exists(model_path):
                _logger.warning(f"Model file missing for dimension {dim}: {model_path}")
                self.logger.log(
                    "TinyScorerModelMissing",
                    {"dimension": dim, "path": model_path},
                )
                continue

            # Load model metadata for architecture configuration
            meta: Dict[str, Any] = {}
            if os.path.exists(meta_path):
                try:
                    meta = load_json(meta_path) or {}
                    _logger.debug(f"Loaded metadata for {dim}: {len(meta)} keys")
                except Exception as e:
                    _logger.error(f"Failed to load metadata for {dim}: {e}")
                    self.logger.log(
                        "TinyScorerMetaLoadError", {"dimension": dim, "error": str(e)}
                    )
            else:
                _logger.warning(f"Metadata file missing for {dim}: {meta_path}")

            # Extract model configuration from metadata with safe defaults
            cfg_meta = meta.get("cfg", {}) if isinstance(meta, dict) else {}
            n_layers = int(cfg_meta.get("n_layers", 2))
            n_recursions = int(cfg_meta.get("n_recursions", 6))
            use_attn = bool(cfg_meta.get("use_attention", False))
            dropout = float(cfg_meta.get("dropout", 0.1))
            attn_heads = int(cfg_meta.get("attn_heads", 4))
            step_scale = float(cfg_meta.get("step_scale", 0.1))
            cons_mask_p = float(cfg_meta.get("consistency_mask_p", 0.10))
            len_norm_L = float(cfg_meta.get("len_norm_L", 512.0))
            vocab_size = int(cfg_meta.get("vocab_size", 101))

            # Optional feature flags from metadata
            enable_agree_head = bool(cfg_meta.get("enable_agree_head", True))
            enable_causal_sens_head = bool(cfg_meta.get("enable_causal_sens_head", True))

            _logger.debug(
                f"Model config for {dim}: layers={n_layers}, recursions={n_recursions}, "
                f"attention={use_attn}, dropout={dropout}"
            )

            # Instantiate model with exact same architecture as training
            _logger.debug(f"Instantiating TRM for dimension {dim}")
            model = TinyRecursionModel(
                d_model=self.dim,
                n_layers=n_layers,
                n_recursions=n_recursions,
                vocab_size=vocab_size,
                use_attention=use_attn,
                dropout=dropout,
                attn_heads=attn_heads,
                step_scale=step_scale,
                consistency_mask_p=cons_mask_p,
                len_norm_L=len_norm_L,
                enable_agree_head=enable_agree_head,
                enable_causal_sens_head=enable_causal_sens_head,
            ).to(self.device)

            # Load trained weights with strict=False for backward compatibility
            _logger.debug(f"Loading model weights from: {model_path}")
            try:
                state = torch.load(model_path, map_location=self.device)
                model.load_state_dict(state, strict=False)
                model.eval()  # Set to evaluation mode
                _logger.debug(f"Successfully loaded weights for {dim}")
            except Exception as e:
                _logger.error(f"Failed to load weights for {dim}: {e}")
                continue

            # Register successfully loaded model
            self.models[dim] = model
            self.model_meta[dim] = meta
            
            _logger.info(f"Successfully loaded TRM model for dimension: {dim}")
            self.logger.log(
                "TinyScorerModelLoaded",
                {
                    "dimension": dim, 
                    "model_path": model_path, 
                    "device": str(self.device)
                },
            )

    # -------------------------
    # Scoring Core
    # -------------------------
    def _score_core(self, context: dict, scorable, dimensions: List[str]) -> ScoreBundle:
        """
        Core scoring method that evaluates goal-response pairs using TRM.
        
        This method:
        1. Converts text to embeddings (shared with HRM)
        2. Runs TRM inference for each dimension
        3. Extracts scores and rich diagnostics
        4. Converts to SCM format for cross-model alignment
        5. Generates aligned vectors for topological analysis
        
        Args:
            context: Scoring context containing goal information
            scorable: The response text to evaluate
            dimensions: List of dimensions to score against
            
        Returns:
            ScoreBundle containing results for all specified dimensions
            
        Logs:
            - Debug: Embedding conversion, model inference, SCM conversion
            - Warning: Missing models or scoring errors
            - Info: Scoring completion statistics
        """
        _logger.debug(f"Starting scoring for {len(dimensions)} dimensions")
        
        # Extract goal information from context
        goal = context.get(GOAL, {})
        goal_text = goal.get(GOAL_TEXT, "")
        _logger.debug(f"Scoring goal: {goal_text[:50]}...")
        _logger.debug(f"Scorable text: {scorable.text[:50]}...")

        results: Dict[str, ScoreResult] = {}

        # Step 1: Convert text to embeddings (shared with HRM for alignment)
        _logger.debug("Converting goal and response to embeddings")
        x_np = self.memory.embedding.get_or_create(goal_text)
        y_np = self.memory.embedding.get_or_create(scorable.text)
        
        # Convert to tensors and ensure correct device placement
        x = torch.tensor(x_np, dtype=torch.float32, device=self.device).unsqueeze(0)
        y = torch.tensor(y_np, dtype=torch.float32, device=self.device).unsqueeze(0)
        z = torch.zeros_like(x)  # neutral third stream for recursive processing
        seq_len = torch.zeros(x.size(0), dtype=torch.int32, device=self.device)
        
        _logger.debug(f"Embedding shapes - x: {x.shape}, y: {y.shape}, z: {z.shape}")

        # Step 2: Score for each specified dimension
        for dim in dimensions:
            _logger.debug(f"Scoring dimension: {dim}")
            model = self.models.get(dim)
            
            if model is None:
                _logger.warning(f"No model available for dimension: {dim}")
                self.logger.log("TinyModelMissing", {"dimension": dim})
                continue

            try:
                # Run TRM inference with gradient disabled for efficiency
                _logger.debug(f"Running TRM inference for {dim}")
                with torch.no_grad():
                    _, halt_logits, _, aux = model(
                        x, y, z, seq_len=seq_len, return_aux=True
                    )
                _logger.debug(f"TRM inference completed for {dim}")

                # Step 3: Extract core metrics from model outputs
                _logger.debug("Extracting core metrics from TRM outputs")
                raw01 = float(max(0.0, min(1.0, _tf(aux.get("score")))))
                
                # Calculate certainty with fallback logic
                if "certainty01" in aux:
                    certainty01 = _tf(aux["certainty01"])
                    _logger.debug("Using certainty01 from aux")
                elif "uncertainty01" in aux:
                    certainty01 = 1.0 - _tf(aux["uncertainty01"])
                    _logger.debug("Derived certainty from uncertainty01")
                elif "uncertainty" in aux:
                    certainty01 = 1.0 - _tf(aux["uncertainty"])
                    _logger.debug("Derived certainty from uncertainty")
                else:
                    certainty01 = 0.5  # Default neutral certainty
                    _logger.debug("Using default certainty 0.5")
                    
                entropy = _tf(aux.get("entropy_aux"))
                halt_prob = _sigmoid_mean(halt_logits)
                
                _logger.debug(
                    f"Core metrics - score: {raw01:.3f}, certainty: {certainty01:.3f}, "
                    f"entropy: {entropy:.3f}, halt_prob: {halt_prob:.3f}"
                )

                # Apply scaling and metadata adjustments
                meta = self.model_meta.get(dim, {})
                final_score = _tf(aux.get("score"))
                tiny_score01 = raw01
                tiny_score100 = round(_safe_scale_0_100(tiny_score01, meta), 4)
                _logger.debug(f"Scaled scores - 01: {tiny_score01:.3f}, 100: {tiny_score100}")

                # Step 4: Build base attributes dictionary
                _logger.debug("Building base attributes dictionary")
                attrs: Dict[str, Any] = {
                    "tiny.score01": tiny_score01,
                    "tiny.score100": tiny_score100,
                    "raw01": tiny_score01,  # backward-compatibility alias
                    "entropy": float(entropy),
                    "certainty01": float(certainty01),
                    "halt_prob": float(halt_prob) if halt_prob is not None else None,
                    # Model context metadata for downstream processing
                    "n_recursions": int(meta.get("cfg", {}).get("n_recursions", 6)),
                    "use_attention": bool(meta.get("cfg", {}).get("use_attention", False)),
                    "dropout": float(meta.get("cfg", {}).get("dropout", 0.1)),
                }

                # Step 5: Add diagnostic attributes based on verbosity level
                if self.attr_level in ("standard", "full"):
                    _logger.debug("Adding standard diagnostic attributes")
                    attrs.update(_extract_standard_aux(aux))
                    
                    # Include optional bridge heads if available
                    if "agree01" in aux and isinstance(aux["agree01"], torch.Tensor):
                        attrs["agree01"] = float(_tf(aux["agree01"]))
                        _logger.debug("Added agree01 diagnostic")
                    if "sens01" in aux and isinstance(aux["sens01"], torch.Tensor):
                        attrs["sens01"] = float(_tf(aux["sens01"]))
                        _logger.debug("Added sens01 diagnostic")

                if self.attr_level == "full":
                    _logger.debug("Adding full diagnostic attributes")
                    attrs.update(_extract_full_aux(aux))
                    
                    # Add raw signal summaries for deep debugging
                    if "score_logit" in aux:
                        attrs["score_logit_mean"] = float(_tf(aux["score_logit"]))
                    if "aux3_logits" in aux and isinstance(aux["aux3_logits"], torch.Tensor):
                        al = aux["aux3_logits"]
                        attrs["aux3_logits_l1_mean"] = float(al.abs().mean().item())

                # Step 6: Convert to Shared Core Metrics format
                _logger.debug("Converting to SCM format for cross-model alignment")
                scm = _build_scm_from_tiny_attrs(attrs)
                attrs.update(scm)
                _logger.debug(f"SCM conversion complete - aggregate: {scm.get('scm.aggregate01', 0):.3f}")

                # Step 7: Mirror dimension scores for PHOS compatibility
                _logger.debug("Mirroring dimension scores for PHOS alignment")
                for dname in ("reasoning", "knowledge", "clarity", "faithfulness", "coverage"):
                    key = f"scm.{dname}.score01"
                    if key in scm:
                        v01 = float(scm[key])
                        attrs[f"tiny.{dname}.score01"]  = v01
                        attrs[f"tiny.{dname}.score100"] = round(v01 * 100.0, 4)
                        attrs[f"tiny.{dname}"] = float(scm[key])
                _logger.debug("Dimension score mirroring complete")

                # Step 8: Generate scoring rationale
                rationale = (
                    f"tiny[{dim}] raw01={float(raw01):.4f}, "
                    f"H={float(entropy):.3f}, C={float(certainty01):.3f}, "
                    f"halt_p={float(halt_prob) if halt_prob is not None else -1:.3f}"
                )
                _logger.debug(f"Generated rationale: {rationale}")

                # Step 9: Build aligned vector for topological analysis
                _logger.debug("Building aligned vector for GAP analysis")
                vector = _tiny_build_vector(attrs)

                # Step 10: Create final ScoreResult
                results[dim] = ScoreResult(
                    dimension=dim,
                    score=tiny_score01,
                    source=self.model_type,
                    rationale=rationale,
                    weight=1.0,
                    attributes={
                        **attrs, 
                        "vector": vector["vector"], 
                        "columns": vector["columns"], 
                        "values": vector["values"]
                    },
                )
                _logger.debug(f"Successfully created ScoreResult for {dim}")

            except Exception as e:
                _logger.error(f"Scoring error for dimension {dim}: {e}")
                self.logger.log("TinyScoreError", {"dimension": dim, "error": str(e)})

        _logger.info(f"Scoring completed for {len(results)} dimensions")
        return ScoreBundle(results=results)

    # -------------------------
    # Utility Methods
    # -------------------------
    @staticmethod
    def _get(d: Dict[str, Any], key: str):
        """
        Safe dictionary access with exception handling.
        
        Args:
            d: Dictionary to access
            key: Key to retrieve
            
        Returns:
            Value if present and accessible, None otherwise
        """
        try:
            return d.get(key)
        except Exception:
            return None

    def __repr__(self):
        """String representation showing loaded models."""
        loaded = {k: (v is not None) for k, v in self.models.items()}
        return f"<TinyScorer(model_type={self.model_type}, loaded={loaded})>"


def _take_scalar(t):
    """
    Extract scalar value from tensor or return float directly.
    
    Args:
        t: Input tensor or scalar
        
    Returns:
        Extracted scalar value as float
    """
    # works with tensor or float
    if isinstance(t, torch.Tensor):
        return float(t.detach().mean().cpu().item())
    return float(t)


# -------------------------
# Helper Functions
# -------------------------

def _tf(v):
    """
    Tensor/array/number → scalar float with safe fallback.
    
    Handles various input types and extracts mean value from tensors.
    Provides safe defaults for None or invalid inputs.
    
    Args:
        v: Input value (tensor, array, or scalar)
        
    Returns:
        Extracted scalar float value
    """
    if v is None:
        _logger.debug("Received None value, returning 0.0")
        return 0.0
    if isinstance(v, torch.Tensor):
        # handle both scalar and vector tensors - use mean for vectors
        result = v.detach().float().mean().item()
        _logger.debug(f"Converted tensor to scalar: {result}")
        return result
    try:
        result = float(v)
        _logger.debug(f"Converted value to float: {result}")
        return result
    except Exception:
        _logger.debug(f"Failed to convert value: {v}, returning 0.0")
        return 0.0


def _sigmoid_mean(v):
    """
    Apply sigmoid and compute mean for halting logits.
    
    Args:
        v: Input tensor or value
        
    Returns:
        Mean sigmoid probability, or None if input is None
    """
    if v is None:
        _logger.debug("Received None for sigmoid_mean")
        return None
    if isinstance(v, torch.Tensor):
        result = torch.sigmoid(v.detach()).mean().item()
        _logger.debug(f"Computed sigmoid mean: {result}")
        return result
    result = float(v)
    _logger.debug(f"Returning float value: {result}")
    return result


def _safe_scale_0_100(raw: float, meta: dict | None) -> float:
    """
    Scale raw [0,1] score to [0,100] range with metadata awareness.
    
    Uses metadata min/max values if available, otherwise uses default 0-100 scaling.
    
    Args:
        raw: Raw score in [0,1] range
        meta: Model metadata containing scaling parameters
        
    Returns:
        Scaled score in appropriate range
    """
    if not meta:
        result = float(max(0.0, min(1.0, raw)) * 100.0)
        _logger.debug(f"Scaled without meta: {raw} -> {result}")
        return result
    
    lo = float(meta.get("min_value", 0.0))
    hi = float(meta.get("max_value", 100.0))
    result = float(max(lo, min(hi, lo + (hi - lo) * max(0.0, min(1.0, raw)))))
    _logger.debug(f"Scaled with meta: {raw} -> {result} (range: {lo}-{hi})")
    return result


def _tiny_build_vector(attrs: Dict[str, Any]) -> Dict[str, Any]:
    """
    Build aligned vector representation for GAP analysis.
    
    Creates a deterministic vector structure that enables cross-model
    alignment in topological analysis. Includes both raw TRM statistics
    and SCM-formatted metrics.
    
    Args:
        attrs: Dictionary of attributes from TRM scoring
        
    Returns:
        Dictionary containing vector, columns, and values for alignment
    """
    _logger.debug("Building aligned vector from attributes")
    vec: Dict[str, float] = {}
    
    # Core TRM statistics for direct access
    vec["tiny.score01"]        = float(attrs.get("tiny.score01", 0.0))
    vec["tiny.score100"]       = float(attrs.get("tiny.score100", 0.0))
    vec["tiny.certainty01"]    = float(attrs.get("certainty01", 0.5))
    vec["tiny.entropy_mean"]   = float(attrs.get("entropy", 0.0))
    
    if "halt_prob" in attrs and attrs["halt_prob"] is not None:
        vec["tiny.halt_prob"] = float(attrs["halt_prob"])
    
    _logger.debug(f"Added {len(vec)} core TRM statistics to vector")

    # SCM-formatted metrics for cross-model alignment
    scm_keys = [
        "scm.reasoning.score01", "scm.knowledge.score01", "scm.clarity.score01",
        "scm.faithfulness.score01", "scm.coverage.score01", "scm.aggregate01",
        "scm.uncertainty01", "scm.ood_hat01", "scm.consistency01",
        "scm.length_norm01", "scm.temp01", "scm.agree_hat01",
    ]
    
    scm_count = 0
    for k in scm_keys:
        if k in attrs:
            vec[k] = float(attrs[k])
            scm_count += 1
    
    _logger.debug(f"Added {scm_count} SCM metrics to vector")

    # Mirror dimension scores for PHOS visualization compatibility
    mirror_count = 0
    for d in ("reasoning", "knowledge", "clarity", "faithfulness", "coverage"):
        k = f"scm.{d}.score01"
        if k in attrs:
            v01 = float(attrs[k])
            vec[f"tiny.{d}.score01"]  = v01
            vec[f"tiny.{d}.score100"] = round(v01 * 100.0, 4)
            vec[f"tiny.{d}"]          = v01
            mirror_count += 1
    
    _logger.debug(f"Mirrored {mirror_count} dimension scores")

    # Create final aligned structure
    cols = list(vec.keys())
    vals = [vec[c] for c in cols]
    _logger.debug(f"Vector construction complete: {len(cols)} columns, {len(vals)} values")
    
    return {"vector": vec, "columns": cols, "values": vals}


def _extract_standard_aux(aux: Dict[str, Any]) -> Dict[str, float]:
    """
    Extract standard diagnostic attributes from TRM auxiliary outputs.
    
    Provides balanced diagnostic coverage including confidence estimates,
    calibration signals, robustness measures, and sensitivity analysis.
    All outputs are normalized to [0,1] range.
    
    Args:
        aux: TRM auxiliary outputs dictionary
        
    Returns:
        Dictionary of standardized diagnostic attributes
    """
    _logger.debug("Extracting standard diagnostic attributes")
    out: Dict[str, float] = {}

    # Confidence triplet from 3-class auxiliary head
    if "aux3_probs" in aux and isinstance(aux["aux3_probs"], torch.Tensor):
        p = aux["aux3_probs"].detach().float()
        out["aux3_p_bad"]  = float(p[..., 0].mean().item())
        out["aux3_p_mid"]  = float(p[..., 1].mean().item())
        out["aux3_p_good"] = float(p[..., 2].mean().item())
        _logger.debug("Extracted aux3 probability triplet")

    # Calibration and temperature signals
    out["temp01"] = float(_tf(aux.get("temp01")))

    # Out-of-distribution detection (prefer newer ood_hat01 format)
    if "ood_hat01" in aux:
        out["ood_hat01"] = float(_tf(aux["ood_hat01"]))
        _logger.debug("Using ood_hat01 for OOD detection")
    elif "ood_hat" in aux:  # backward compatibility
        out["ood_hat01"] = float(_tf(aux["ood_hat"]))
        _logger.debug("Using ood_hat (legacy) for OOD detection")

    # Robustness and sensitivity measures
    out["consistency_hat"] = float(_tf(aux.get("consistency_hat")))
    out["jacobian_fd"]     = float(_tf(aux.get("jacobian_fd")))
    _logger.debug("Extracted robustness and sensitivity measures")

    # Reconstruction quality and disagreement prediction
    out["recon_sim"]  = float(_tf(aux.get("recon_sim")))
    out["disagree_hat"] = float(_tf(aux.get("disagree_hat")))
    _logger.debug("Extracted reconstruction and disagreement signals")

    # Length normalization (prefer 0..1 normalized version)
    if "length_norm01" in aux:
        out["length_norm01"] = float(_tf(aux["length_norm01"]))
        _logger.debug("Using length_norm01 for length effect")
    else:
        # Derive from tanh-normalized len_effect if available
        if "len_effect" in aux:
            le = float(_tf(aux["len_effect"]))
            out["length_norm01"] = float(max(0.0, min(1.0, (le + 1.0) * 0.5)))
            _logger.debug("Derived length_norm01 from len_effect")
        else:
            out["length_norm01"] = 0.0
            _logger.debug("Using default length_norm01")

    # Sparse Autoencoder concept sparsity
    out["concept_sparsity"] = float(_tf(aux.get("concept_sparsity")))
    _logger.debug("Extracted concept sparsity measure")

    _logger.debug(f"Standard diagnostics extraction complete: {len(out)} attributes")
    return out


def _extract_full_aux(aux: Dict[str, Any]) -> Dict[str, float]:
    """
    Extract full diagnostic attributes including raw signal summaries.
    
    Provides maximum detail for debugging and analysis, including
    raw logit summaries and internal representation statistics.
    Use only when deep inspection is required.
    
    Args:
        aux: TRM auxiliary outputs dictionary
        
    Returns:
        Dictionary of detailed diagnostic attributes
    """
    _logger.debug("Extracting full diagnostic attributes")
    out: Dict[str, float] = {}

    # Summaries of raw head outputs for debugging
    for k in ("log_var", "consistency_logit", "disagree_logit"):
        if k in aux and isinstance(aux[k], torch.Tensor):
            t = aux[k].detach()
            out[f"{k}_mean"] = float(t.mean().item())
            _logger.debug(f"Added {k}_mean to full diagnostics")

    # Reconstruction detail analysis
    if "y_recon" in aux and isinstance(aux["y_recon"], torch.Tensor):
        yr = aux["y_recon"].detach()
        out["y_recon_norm_mean"] = float(yr.norm(dim=-1).mean().item())
        _logger.debug("Added y_recon_norm_mean to full diagnostics")

    # Sparse Autoencoder concept analysis
    if "concept_vec" in aux and isinstance(aux["concept_vec"], torch.Tensor):
        c = aux["concept_vec"].detach()
        out["concept_vec_l2_mean"] = float((c.pow(2).sum(-1).sqrt()).mean().item())
        _logger.debug("Added concept_vec_l2_mean to full diagnostics")

    _logger.debug(f"Full diagnostics extraction complete: {len(out)} attributes")
    return out


# === SCM mapping from Tiny aux → aligned scm.* columns =======================

def _build_scm_from_tiny_attrs(attrs: Dict[str, Any]) -> Dict[str, float]:
    """
    Convert TRM attributes to Shared Core Metrics format.
    
    Maps TRM's internal diagnostic signals to the standardized 5-dimensional
    SCM format using learned weighting patterns. This enables direct
    comparison with HRM and other evaluation systems.
    
    The mapping uses TRM's diagnostic patterns to infer dimension scores:
    - Reasoning: Emphasizes consistency, low uncertainty, agreement
    - Knowledge: Focuses on in-distribution signals and reconstruction
    - Clarity: Uses token quality and length normalization
    - Faithfulness: Based on reconstruction and consistency
    - Coverage: Considers concept activity and distribution alignment
    
    Args:
        attrs: TRM attributes dictionary
        
    Returns:
        Dictionary of SCM-formatted scores in [0,1] range
    """
    _logger.debug("Building SCM from TRM attributes")
    
    # Extract and clamp core diagnostic signals
    certainty = float(attrs.get("certainty01", 0.5))
    unc01     = 1.0 - max(0.0, min(1.0, certainty))
    cons01    = max(0.0, min(1.0, float(attrs.get("consistency_hat", 0.5))))
    ood01     = max(0.0, min(1.0, float(attrs.get("ood_hat", 0.0))))
    len01     = max(0.0, min(1.0, float(attrs.get("len_effect", 0.0))))
    temp01    = max(0.0, min(1.0, float(attrs.get("temp01", 0.0))))
    agree01   = max(0.0, min(1.0, float(attrs.get("agree01", 0.5))))

    # Extract additional diagnostic signals
    recon_sim      = max(0.0, min(1.0, float(attrs.get("recon_sim", 0.5))))
    concept_sparse = max(0.0, min(1.0, float(attrs.get("concept_sparsity", 0.5))))
    p_bad          = max(0.0, min(1.0, float(attrs.get("aux3_p_bad", 0.5))))
    token_ok       = 1.0 - p_bad  # clarity proxy: lower bad probability → clearer

    _logger.debug(f"Core signals - uncertainty: {unc01:.3f}, consistency: {cons01:.3f}, OOD: {ood01:.3f}")

    # Dimension-specific scoring using diagnostic patterns
    dim_scores: Dict[str, float] = {}
    
    # Reasoning: weighted toward stability, consistency, and confidence
    dim_scores["reasoning"] = 0.60*cons01 + 0.30*(1.0-unc01) + 0.10*agree01
    _logger.debug(f"Reasoning score: {dim_scores['reasoning']:.3f}")
    
    # Knowledge: emphasizes distribution alignment and comprehension
    dim_scores["knowledge"] = 0.50*(1.0-ood01) + 0.30*recon_sim + 0.20*(1.0-unc01)
    _logger.debug(f"Knowledge score: {dim_scores['knowledge']:.3f}")
    
    # Clarity: based on token quality and brevity
    dim_scores["clarity"] = 0.50*token_ok + 0.30*(1.0-len01) + 0.20*cons01
    _logger.debug(f"Clarity score: {dim_scores['clarity']:.3f}")
    
    # Faithfulness: reconstruction quality and stability
    dim_scores["faithfulness"] = 0.50*recon_sim + 0.30*cons01 + 0.20*(1.0-unc01)
    _logger.debug(f"Faithfulness score: {dim_scores['faithfulness']:.3f}")
    
    # Coverage: concept activity and confidence
    dim_scores["coverage"] = 0.40*concept_sparse + 0.40*(1.0-unc01) + 0.20*(1.0-ood01)
    _logger.debug(f"Coverage score: {dim_scores['coverage']:.3f}")

    # Ensure all scores are in valid [0,1] range
    for k in dim_scores:
        v = dim_scores[k]
        dim_scores[k] = float(min(1.0, max(0.0, v)))
    _logger.debug("Applied score clamping to [0,1] range")

    # Build final SCM dictionary
    scm: Dict[str, float] = {
        f"scm.{k}.score01": dim_scores[k]
        for k in ("reasoning", "knowledge", "clarity", "faithfulness", "coverage")
    }
    scm["scm.aggregate01"]   = float(sum(dim_scores.values())/5.0)
    scm["scm.uncertainty01"] = float(unc01)
    scm["scm.ood_hat01"]     = float(ood01)
    scm["scm.consistency01"] = float(cons01)
    scm["scm.length_norm01"] = float(len01)
    scm["scm.temp01"]        = float(temp01)
    scm["scm.agree_hat01"]   = float(agree01)

    _logger.debug(f"SCM construction complete - aggregate: {scm['scm.aggregate01']:.3f}")
    return scm