"""Merlin-Agent modeling: Qwen3_5 text causal LM with per-layer quantum injection.

At each full-attention layer, a fixed quantum-derived direction (buffer q_k) is
added to that layer's output hidden state with an RMS-matched, alpha-scaled
magnitude. Injection lives in the model (buffers + hooks re-installed in
__init__), so it survives save/reload — not a bare, non-persisted hook.

alpha == 0 -> exact base model.
"""
import torch
try:
    from transformers import Qwen3_5ForCausalLM
except ImportError:  # not exported at top level in some transformers versions
    from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5ForCausalLM

try:
    from .configuration_merlin_agent import MerlinAgentConfig
except ImportError:  # when loaded as flat remote code (trust_remote_code)
    from configuration_merlin_agent import MerlinAgentConfig


def _inject(h: torch.Tensor, q: torch.Tensor, alpha: float) -> torch.Tensor:
    if alpha == 0:
        return h
    q = q.to(dtype=h.dtype, device=h.device)
    rms_h = h.pow(2).mean(dim=-1, keepdim=True).sqrt()
    rms_q = q.pow(2).mean().sqrt().clamp_min(1e-6)
    return h + alpha * (rms_h / rms_q) * q


class MerlinAgentForCausalLM(Qwen3_5ForCausalLM):
    config_class = MerlinAgentConfig

    def __init__(self, config):
        super().__init__(config)
        self._inj_layers = list(config.full_attention_layer_indices)
        for k in range(len(self._inj_layers)):
            self.register_buffer(f"q_{k}", torch.zeros(config.hidden_size), persistent=True)
        self._install_injection_hooks()

    def _install_injection_hooks(self):
        for k, li in enumerate(self._inj_layers):
            self.model.layers[li].register_forward_hook(self._make_hook(f"q_{k}"))

    def _make_hook(self, qk: str):
        def hook(module, args, output):
            q = getattr(self, qk)
            a = float(self.config.quantum_injection_alpha)
            if isinstance(output, tuple):
                return (_inject(output[0], q, a),) + tuple(output[1:])
            return _inject(output, q, a)
        return hook

    def set_quantum_signatures(self, q_vectors):
        """Load the 8 projected quantum direction vectors (each shape [hidden_size])."""
        assert len(q_vectors) == len(self._inj_layers)
        for k, v in enumerate(q_vectors):
            getattr(self, f"q_{k}").copy_(torch.as_tensor(v, dtype=torch.float32))