chimera/model.py

"""
Chimera 5.2 — full causal LM with FUNCTIONAL self-evolution.

Key changes for auto-evolution:
* SelfEvolutionEngine is called at EVERY layer during forward pass
* Semantic memory modulation is added to hidden states
* TTT updates target MLP weights in-place during forward
* Evolution loss is added to causal LM loss during training
* Contrastive evaluation tracks memory usefulness
* Loop depth classifier sets compute budget per sequence
"""

from __future__ import annotations

import json
from typing import Any, List, Optional, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint

from .quantization import BitLinear, RMSNorm
from .layers import (GatedDeltaNetLayer, MLSTMLayer, TitansMACLayer,
                     TSPSpanKnotLayer, SwiGLUMLP)
from .moe import MoELayer
from .looping import ParcaeLoopController
from .inference import (SpanInferenceEngine, GrammarFST, EntropyValve,
                        DebtLedger, BraidState)
from .evolution import SelfEvolutionEngine
from .multimodal import VisionEncoder, AudioEncoder


class CausalLMOutput(dict):
    """Light HF-compatible output dict supporting tuple unpacking."""

    def __init__(self, loss: Optional[torch.Tensor] = None,
                 logits: Optional[torch.Tensor] = None,
                 hidden_states: Optional[torch.Tensor] = None,
                 caches: Optional[list] = None,
                 evolution_metrics: Optional[dict] = None):
        super().__init__(loss=loss, logits=logits,
                         hidden_states=hidden_states, caches=caches,
                         evolution_metrics=evolution_metrics)
        self.loss = loss
        self.logits = logits
        self.hidden_states = hidden_states
        self.caches = caches
        self.evolution_metrics = evolution_metrics or {}

    def __iter__(self):
        yield self.loss
        yield self.logits


def expand_layer_pattern(config: dict) -> List[str]:
    """Expand the layer-pattern shorthand into a list."""
    backbone = config.get("backbone", {})
    pattern_str = backbone.get("layer_pattern", "GD XM GD TM GD XM GD SK")
    aliases = backbone.get("layer_aliases", {
        "GD": "gated_deltanet", "XM": "xlstm_m",
        "TM": "titans_mac", "SK": "tsp_span_knot",
    })
    pattern = pattern_str.split()
    n_layers = int(config.get("num_hidden_layers", 28))
    full = (pattern * (n_layers // len(pattern) + 1))[:n_layers]
    return [aliases.get(p, p) for p in full]


class Chimera51Block(nn.Module):
    """One block with evolution-aware forward."""

    _RECURRENT = {"gated_deltanet", "xlstm_m", "titans_mac", "tsp_span_knot"}

    def __init__(self, config: dict, layer_type: str, layer_idx: int,
                 use_moe: bool = False):
        super().__init__()
        h = int(config["hidden_size"])
        eps = float(config.get("rms_norm_eps", 1e-6))
        heads = int(config["num_heads"])
        head_dim = int(config["head_dim"])
        ternary = bool(config.get("use_ternary", True))
        chunk_sz = int(config.get("gated_deltanet", {}).get("chunk_size", 64))

        self.layer_idx = layer_idx
        self.layer_type = layer_type
        self.attn_norm = RMSNorm(h, eps=eps)

        if layer_type == "gated_deltanet":
            self.attn = GatedDeltaNetLayer(h, heads, head_dim, norm_eps=eps,
                                           chunk_size=chunk_sz, use_ternary=ternary)
        elif layer_type == "xlstm_m":
            mem_h = config.get("xlstm", {}).get("memory_size_per_head", [head_dim, head_dim])
            self.attn = MLSTMLayer(h, heads, int(mem_h[0]), norm_eps=eps,
                                   use_ternary=ternary)
        elif layer_type == "titans_mac":
            tc = config.get("titans", {})
            self.attn = TitansMACLayer(h, heads, head_dim,
                                       memory_depth=int(tc.get("memory_depth", 2)),
                                       persistent_slots=int(tc.get("persistent_memory_slots", 64)),
                                       local_window=int(tc.get("local_window_size", 1024)),
                                       norm_eps=eps, use_ternary=ternary)
        elif layer_type == "tsp_span_knot":
            self.attn = TSPSpanKnotLayer(h, heads, head_dim, norm_eps=eps,
                                         chunk_size=chunk_sz, use_ternary=ternary)
        else:
            raise ValueError(f"Unknown layer type: {layer_type}")

        self.mlp_norm = RMSNorm(h, eps=eps)
        self.use_moe = bool(use_moe)
        if self.use_moe:
            moe_cfg = config.get("backbone", {}).get("moe", {})
            self.mlp = MoELayer(
                hidden_size=h,
                moe_intermediate_size=int(moe_cfg.get("moe_intermediate_size", h * 2)),
                n_routed_experts=int(moe_cfg.get("n_routed_experts", 16)),
                n_shared_experts=int(moe_cfg.get("n_shared_experts", 1)),
                num_experts_per_tok=int(moe_cfg.get("num_experts_per_tok", 2)),
                use_ternary=ternary,
            )
        else:
            inter = int(config.get("intermediate_size", int(h * 8 / 3)))
            inter = 256 * ((inter + 255) // 256)
            self.mlp = SwiGLUMLP(h, inter, use_ternary=ternary)

        # Evolution modulation projection (learnable scale)
        self.evo_gate = nn.Linear(h, h, bias=False)
        nn.init.zeros_(self.evo_gate.weight)

    def forward(self, x: torch.Tensor, cache: Optional[dict] = None,
                evo_modulation: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, dict]:
        # Apply attention with pre-norm
        normed = self.attn_norm(x)
        attn_out, new_cache = self.attn(normed, cache=cache)
        x = x + attn_out

        # Apply MLP with pre-norm
        x = x + self.mlp(self.mlp_norm(x))

        # Apply evolution modulation (gated residual)
        if evo_modulation is not None:
            gate = torch.sigmoid(self.evo_gate(x))
            x = x + gate * evo_modulation

        return x, new_cache


class Chimera51ForCausalLM(nn.Module):
    """Chimera 5.x causal language model with functional self-evolution."""

    def __init__(self, config: dict):
        super().__init__()
        self.config = config
        h = int(config["hidden_size"])
        vocab = int(config["vocab_size"])
        n_layers = int(config["num_hidden_layers"])
        eps = float(config.get("rms_norm_eps", 1e-6))

        self.embed = nn.Embedding(vocab, h)
        layer_types = expand_layer_pattern(config)
        moe_layers = set(int(i) for i in config.get("backbone", {}).get("moe", {}).get("layers", []))

        self.layers = nn.ModuleList([
            Chimera51Block(config, layer_types[i], i, use_moe=(i in moe_layers))
            for i in range(n_layers)
        ])

        self.norm = RMSNorm(h, eps=eps)
        self.lm_head = nn.Linear(h, vocab, bias=False)

        if config.get("tie_word_embeddings", True):
            self.lm_head.weight = self.embed.weight

        # Parcae looping controller
        loop_cfg = config.get("looping", {})
        self.looping_enabled = bool(loop_cfg.get("enabled", True)) and n_layers >= 3
        if self.looping_enabled:
            self.prelude_start, self.prelude_end = loop_cfg.get("prelude", [0, min(3, n_layers - 1)])
            self.loop_start, self.loop_end = loop_cfg.get("loop", [min(4, n_layers - 1), max(4, n_layers - 4)])
            self.coda_start, self.coda_end = loop_cfg.get("coda", [max(0, n_layers - 4), n_layers - 1])
            self.loop_controller = ParcaeLoopController(
                h, loop_range=tuple(loop_cfg.get("loop_range", [1, 6])),
                loop_default=int(loop_cfg.get("loop_default", 2)),
                adaptive_exit_threshold=float(loop_cfg.get("adaptive_exit_threshold", 0.01)),
            )

        # Inference systems
        si_cfg = config.get("span_inference", {})
        self.span_engine = SpanInferenceEngine(h, si_cfg) if si_cfg.get("enabled", True) else None
        self.grammar = GrammarFST(config.get("grammar", {}))
        self.entropy_valve = EntropyValve(config.get("entropy_valve", {}))
        self.debt_ledger = DebtLedger(config.get("debt_ledger", {}))

        # Self-evolution — FUNCTIONAL
        evo_cfg = dict(config.get("self_evolution", {}))
        evo_cfg["_semantic_memory_config"] = config.get("semantic_memory", {})
        self.evolution = SelfEvolutionEngine(evo_cfg, h)
        self.evo_weight = float(config.get("evolution_loss_weight", 0.01))
        self.evo_every_n_layers = int(config.get("evolution_every_n_layers", 4))

        # Multimodal
        mm_cfg = dict(config.get("multimodal", {}))
        mm_cfg["hidden_size"] = h
        if mm_cfg.get("enabled", False):
            self.vision_encoder = VisionEncoder(mm_cfg)
            self.audio_encoder = AudioEncoder(mm_cfg)
        else:
            self.vision_encoder = None
            self.audio_encoder = None

        self.gradient_checkpointing = False
        self._init_weights()
        self._wire_semantic_memory()

    def enable_gradient_checkpointing(self) -> None:
        self.gradient_checkpointing = True

    def disable_gradient_checkpointing(self) -> None:
        self.gradient_checkpointing = False

    def _wire_semantic_memory(self) -> None:
        mem = self.evolution.semantic_memory
        for layer in self.layers:
            if hasattr(layer.attn, "set_semantic_memory"):
                layer.attn.set_semantic_memory(mem)

    def _init_weights(self) -> None:
        init_range = float(self.config.get("initializer_range", 0.006))
        for module in self.modules():
            if isinstance(module, (nn.Linear, BitLinear)):
                if module.weight is not None:
                    nn.init.normal_(module.weight, mean=0.0, std=init_range)
                if getattr(module, "bias", None) is not None:
                    nn.init.zeros_(module.bias)
            elif isinstance(module, nn.Embedding):
                nn.init.normal_(module.weight, mean=0.0, std=init_range)
        for module in self.modules():
            if isinstance(module, BitLinear):
                module.invalidate_packed()

    def _run_layers(self, x: torch.Tensor, start: int, end: int,
                    caches: Optional[list],
                    compute_logits: bool = False,
                    labels: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, Optional[torch.Tensor], list]:
        """Run layers with evolution hooks. Returns (x, logits_if_computed, caches)."""
        all_metrics = []
        logits = None
        evolution_loss = torch.tensor(0.0, device=x.device)

        for i in range(start, min(end + 1, len(self.layers))):
            layer = self.layers[i]
            cache = caches[i] if caches is not None else None

            # Evolution modulation every N layers (lightweight)
            evo_mod = None
            if i % self.evo_every_n_layers == 0 and self.evolution is not None:
                # Compute modulation from semantic memory
                # Note: loss parameter requires a scalar loss tensor for TTT/surprise;
                #       pass None during standard forward, compute explicitly for TTT
                evo_result = self.evolution(
                    hidden_states=x.detach() if not x.requires_grad else x,
                    layer_idx=i,
                    loss=None
                )
                evo_mod = evo_result['modulation']
                if evo_result['evolution_loss'] is not None:
                    evolution_loss = evolution_loss + evo_result['evolution_loss']
                all_metrics.append(evo_result.get('metrics', {}))

                # TTT update for target layers (only in training, no backprop)
                if self.training and evo_result.get('ttt_delta') is not None:
                    with torch.no_grad():
                        # Apply TTT to MLP down-projection if this is a target layer
                        if hasattr(layer.mlp, 'w_down'):
                            layer.mlp.w_down.data.add_(evo_result['ttt_delta'] * self.evolution.ttt.inner_lr)

            if self.gradient_checkpointing and self.training:
                def _ckpt_fn(x_in, layer=layer, cache=cache, evo=evo_mod):
                    out, _ = layer(x_in, cache=cache, evo_modulation=evo)
                    return out
                x = checkpoint(_ckpt_fn, x, use_reentrant=False)
            else:
                x, new_cache = layer(x, cache=cache, evo_modulation=evo_mod)
                if caches is not None:
                    caches[i] = new_cache

            # Compute probe logits for entropy valve (every few layers)
            if compute_logits and i == end:
                logits = self.lm_head(self.norm(x[:, -1:, :]))

        return x, logits, caches, evolution_loss, all_metrics

    def forward(self, input_ids: torch.Tensor,
                labels: Optional[torch.Tensor] = None,
                pixel_values: Optional[torch.Tensor] = None,
                mel_features: Optional[torch.Tensor] = None,
                num_loops: Optional[int] = None,
                caches: Optional[list] = None,
                use_cache: bool = False,
                logits_to_keep: int = 0,
                return_evolution_metrics: bool = False):
        x = self.embed(input_ids)

        # Multimodal prepend
        if pixel_values is not None and self.vision_encoder is not None:
            v = self.vision_encoder(pixel_values)
            if v is not None:
                x = torch.cat([v, x], dim=1)
        if mel_features is not None and self.audio_encoder is not None:
            a = self.audio_encoder(mel_features)
            if a is not None:
                x = torch.cat([a, x], dim=1)

        if caches is None and use_cache:
            caches = [None] * len(self.layers)

        total_evo_loss = torch.tensor(0.0, device=x.device)
        all_evo_metrics = []

        # Prelude + Loop + Coda with evolution
        if self.looping_enabled and hasattr(self, "loop_controller"):
            # Prelude
            x, probe_logits, caches, evo_loss, metrics = self._run_layers(
                x, self.prelude_start, self.prelude_end, caches,
                compute_logits=not self.training, labels=labels)
            total_evo_loss = total_evo_loss + evo_loss
            all_evo_metrics.extend(metrics)

            # Determine loop depth
            effective = num_loops
            if effective is None and not self.training and probe_logits is not None:
                effective = self.entropy_valve.get_loop_count(probe_logits)
            elif effective is None and self.evolution is not None:
                # Use loop classifier from evolution
                last_hidden = x[:, -1, :].mean(dim=0, keepdim=True)  # Average over batch
                effective = self.evolution.loop_classifier(last_hidden).item()
                effective = max(1, min(effective, 6))

            # Loop body
            loop_fn = lambda inp: self._run_layers(
                inp, self.loop_start, self.loop_end, caches, labels=labels)[0]
            x = self.loop_controller(x, loop_fn, num_loops=effective)

            # Coda
            x, _, caches, evo_loss, metrics = self._run_layers(
                x, self.coda_start, self.coda_end, caches, labels=labels)
            total_evo_loss = total_evo_loss + evo_loss
            all_evo_metrics.extend(metrics)
        else:
            x, _, caches, evo_loss, metrics = self._run_layers(
                x, 0, len(self.layers) - 1, caches,
                compute_logits=not self.training, labels=labels)
            total_evo_loss = total_evo_loss + evo_loss
            all_evo_metrics.extend(metrics)

        # Final norm and logits
        if logits_to_keep and labels is None:
            keep = int(logits_to_keep)
            tail = x[:, -keep:, :]
            tail = self.norm(tail)
            if self.span_engine is not None:
                tail = self.span_engine(tail)
            logits = self.lm_head(tail)
        else:
            x = self.norm(x)
            if self.span_engine is not None:
                x = self.span_engine(x)
            logits = self.lm_head(x)

        logits = self.grammar(logits)
        logits = self.debt_ledger(logits)

        # Self-feedback refinement check (inference only)
        if not self.training and self.evolution is not None:
            should_refine = self.evolution.self_feedback.should_refine(logits)
            if should_refine:
                all_evo_metrics.append({'refinement_triggered': True})

        # Compute loss
        loss = None
        if labels is not None:
            seq_len = min(logits.size(1), labels.size(1))
            shift_logits = logits[:, :seq_len, :].contiguous()
            shift_labels = labels[:, :seq_len].contiguous()
            ce_loss = F.cross_entropy(
                shift_logits.view(-1, shift_logits.size(-1)),
                shift_labels.view(-1),
                ignore_index=-100,
            )
            # Add evolution loss (contrastive memory evaluation)
            loss = ce_loss + self.evo_weight * total_evo_loss
        else:
            ce_loss = None

        # Store episodic case after forward (for inference mode)
        if not self.training and self.evolution is not None:
            last_hidden = x[:, -1, :].detach()
            # Schedule episodic storage for end of sequence
            # (In real use, call model.evolution.store_episodic() explicitly)

        return CausalLMOutput(
            loss=loss,
            logits=logits,
            hidden_states=x,
            caches=caches if use_cache else None,
            evolution_metrics={
                'ce_loss': ce_loss.item() if ce_loss is not None else None,
                'evo_loss': total_evo_loss.item(),
                'layer_metrics': all_evo_metrics,
            } if return_evolution_metrics else None
        )

    @torch.no_grad()
    def prepare_for_inference(self) -> None:
        """Pre-pack every BitLinear so the first generation step is fast."""
        for module in self.modules():
            if isinstance(module, BitLinear):
                module.prepare_for_inference()

    def get_mode_config(self, mode: str = "balanced") -> dict:
        modes = self.config.get("modes", {})
        return modes.get(mode, modes.get("balanced", {}))

    def count_parameters(self) -> dict:
        total = sum(p.numel() for p in self.parameters())
        ternary = sum(p.numel() for _, m in self.named_modules()
                      if isinstance(m, BitLinear) for p in m.parameters())
        return {"total": total, "ternary": ternary, "fp32": total - ternary}

    @classmethod
    def from_config_file(cls, path: str) -> "Chimera51ForCausalLM":
        with open(path, "r", encoding="utf-8") as fh:
            config = json.load(fh)
        return cls(config)


__all__ = ["Chimera51ForCausalLM", "Chimera51Block", "CausalLMOutput",
           "expand_layer_pattern"]