src/fractus/continuous_engine.py

"""ContinuousThoughtEngine: a model that thinks in real-time, not token-by-token.

THE PARADIGM SHIFT. Claude and GPT are static functions: input → output.
This engine is a DYNAMICAL SYSTEM: it runs continuously, maintains a state
of "consciousness", and produces output only when it has something confident
to say.

Architecture:
    - A "thought state" h (d_model vector) that persists across time.
    - At each TICK (a single cheap forward step, not a full generation):
        1. The Kuramoto oscillators advance the phase state by one RK4 step.
        2. The attention state (S, z) absorbs the current observation.
        3. The MoE processes the thought state, routed by the Kuramoto phases.
        4. A confidence head estimates "how sure am I?".
        5. If confidence > threshold → emit output token, reset partial state.
           If confidence < threshold → continue thinking (accumulate).

WHY THIS IS FAST ON CPU:
    - Each tick is ONE forward of the thought state (not B×L tokens).
    - The state is (d_model,) — tiny, not (B, L, d_model).
    - Training is ONLINE: one observation at a time, one gradient at a time.
      No batches, no BPTT through long sequences.
    - The "reasoning depth" is adaptive: easy observations = 1 tick,
      hard ones = 10 ticks. Energy-proportional THINKING.

WHY THIS MAKES CLAUDE/GPT OBSOLETE:
    - They can't "think" — they do one forward pass and output.
    - This engine REASONS: it takes multiple ticks on hard problems,
      accumulating evidence, until it's confident.
    - It's PROACTIVE: it can emit output without being prompted
      (when confidence crosses threshold from internal dynamics).
    - It's CONTINUOUS: the state never resets, so it has true memory
      across an entire conversation/session, not a context window.
"""

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

from .nn.attention import FractalLinearAttention
from .nn.phase_ode import KuramotoLayer
from .nn.stats import elu_plus_one
from .nn.farey import expert_phases
from .nn.cached_siren import CachedStructuredSirenLinear


class ContinuousThoughtEngine(nn.Module):
    """A continuous-time reasoning engine.

    Unlike a standard LM (input_ids → logits per token), this engine
    maintains a persistent "thought state" and advances it tick by tick.

    Args:
        vocab_size:      vocabulary size (for input embedding + output head).
        d_model:         dimension of the thought state.
        n_heads, d_head: attention configuration.
        n_oscillators:   Kuramoto oscillator count (the "consciousness clock").
        n_experts:       MoE expert count.
        top_k:           active experts per tick.
    """

    def __init__(
        self,
        vocab_size: int = 50257,
        d_model: int = 256,
        n_heads: int = 4,
        d_head: int = 64,
        n_levels: int = 2,
        n_oscillators: int = 16,
        coupling_rank: int = 8,
        n_experts: int = 8,
        top_k: int = 2,
        expert_d_ff: int = 256,
        siren_rank: int = 32,
    ):
        super().__init__()
        self.vocab_size = vocab_size
        self.d_model = d_model

        # Input embedding: maps an observed token → a "perception" vector.
        self.observe = nn.Embedding(vocab_size, d_model)

        # The thought state components:
        # 1. Attention (processes the current observation in context of memory)
        self.attn = FractalLinearAttention(d_model, n_heads, d_head, n_levels)
        self.norm_attn = nn.LayerNorm(d_model)

        # 2. Kuramoto (the consciousness clock — advances phases each tick)
        self.kuramoto = KuramotoLayer(d_model, n_oscillators, coupling_rank,
                                      n_steps=1, dt=0.1)  # 1 step per tick (fast)
        self.norm_kur = nn.LayerNorm(d_model)

        # 3. MoE (transforms the thought, routed by Kuramoto phases)
        self.n_experts = n_experts
        self.top_k = top_k
        self.expert_d_ff = expert_d_ff
        phases = expert_phases(n_experts)
        self.register_buffer("expert_phases", torch.tensor(phases, dtype=torch.float32))
        self.kappa = 4.0

        self.experts_w1 = nn.ModuleList([
            CachedStructuredSirenLinear(d_model, expert_d_ff, rank=siren_rank,
                                        siren_hidden=32, refresh_every=8)
            for _ in range(n_experts)
        ])
        self.experts_w2 = nn.ModuleList([
            CachedStructuredSirenLinear(expert_d_ff, d_model, rank=siren_rank,
                                        siren_hidden=32, refresh_every=8)
            for _ in range(n_experts)
        ])
        self.norm_moe = nn.LayerNorm(d_model)

        # 4. Confidence head: "how sure am I about the current thought?"
        self.confidence_head = nn.Linear(d_model, 1)

        # 5. Output head: "what do I want to say?"
        self.output_head = nn.Linear(d_model, vocab_size, bias=False)

        # Initialize the thought state (will be set by reset_thought).
        self.register_buffer("thought_state", torch.zeros(1, 1, d_model))
        self.register_buffer("attn_S", torch.zeros(1, n_heads * d_head, n_heads * d_head))
        self.register_buffer("attn_z", torch.zeros(1, n_heads * d_head))
        self.register_buffer("kuramoto_phases", torch.zeros(1, 1, n_oscillators))

    def reset_thought(self, batch_size: int = 1):
        """Reset the thought state to zero (start of a new session)."""
        self.thought_state = torch.zeros(batch_size, 1, self.d_model,
                                         device=self.thought_state.device)
        d = self.attn.n_heads * self.attn.d_head
        self.attn_S = torch.zeros(batch_size, d, d, device=self.thought_state.device)
        self.attn_z = torch.zeros(batch_size, d, device=self.thought_state.device)
        self.kuramoto_phases = torch.zeros(
            batch_size, 1, self.kuramoto.N, device=self.thought_state.device
        )

    def tick(self, observation: torch.Tensor = None) -> tuple:
        """Advance the thought by ONE tick.

        Args:
            observation: (B,) token id to absorb, or None (pure thinking).
        Returns:
            (output_logits (B, vocab), confidence (B,))
        """
        B = self.thought_state.shape[0]
        h = self.thought_state  # (B, 1, d_model)

        # Absorb observation if provided.
        if observation is not None:
            obs_vec = self.observe(observation).unsqueeze(1)  # (B, 1, d_model)
            h = h + obs_vec  # add perception to thought

        # 1. Attention: process h with the accumulated state (S, z).
        h_normed = self.norm_attn(h)
        # We use a simplified 1-token attention: the "memory" is in (S, z).
        # Project Q, K, V from h.
        attn = self.attn
        D = attn.d_head
        q = torch.einsum("bld,de->ble", h_normed, attn.w_qkv[0]) + attn.b_qkv[0]
        k = torch.einsum("bld,de->ble", h_normed, attn.w_qkv[1]) + attn.b_qkv[1]
        v = torch.einsum("bld,de->ble", h_normed, attn.w_qkv[2]) + attn.b_qkv[2]
        # Apply feature map.
        q_feat = elu_plus_one(q + attn.level_offsets[0])
        k_feat = elu_plus_one(k + attn.level_offsets[0])
        # Update state: S += k⊗v, z += k.
        for hd in range(attn.n_heads):
            kh = k_feat[:, :, hd * D:(hd + 1) * D]  # (B, 1, D)
            vh = v[:, :, hd * D:(hd + 1) * D]
            qh = q_feat[:, :, hd * D:(hd + 1) * D]
            # S is (B, n_heads*D, n_heads*D) — we update only this head's block.
            s_start = hd * D
            s_end = (hd + 1) * D
            outer = (kh.squeeze(1).unsqueeze(2) * vh.squeeze(1).unsqueeze(1))  # (B, D, D)
            # Detach the accumulated state — online training doesn't BPTT
            # through the full state history. Only the current tick's
            # forward graph is kept for backward.
            self.attn_S = self.attn_S.clone()
            self.attn_S[:, s_start:s_end, s_start:s_end] += outer.detach()
            self.attn_z = self.attn_z.clone()
            self.attn_z[:, s_start:s_end] += kh.squeeze(1).detach()
        # Compute attention output from current q and state.
        attn_out = torch.zeros_like(h)
        for hd in range(attn.n_heads):
            s_start = hd * D
            s_end = (hd + 1) * D
            qh = q_feat[:, 0, hd * D:(hd + 1) * D]  # (B, D)
            S_h = self.attn_S[:, s_start:s_end, s_start:s_end]  # (B, D, D)
            z_h = self.attn_z[:, s_start:s_end]  # (B, D)
            num = torch.bmm(qh.unsqueeze(1), S_h).squeeze(1)  # (B, D)
            denom = (qh * z_h).sum(dim=-1, keepdim=True)  # (B, 1)
            safe = denom.abs() > 1e-10
            yh = torch.where(safe, num / (denom + 1e-20), torch.zeros_like(num))
            attn_out[:, 0, hd * D:(hd + 1) * D] = yh
        attn_out = attn_out @ attn.w_out + attn.b_out
        h = h + attn_out

        # 2. Kuramoto: advance phases by one step.
        h_kur = self.norm_kur(h)
        # Encode phases from hidden, integrate 1 step.
        theta = self.kuramoto._encode_from_hidden(h_kur)
        # Carry previous phases forward (add the delta).
        theta = theta + 0.1 * self.kuramoto._derivative(theta)
        theta = torch.remainder(theta, self.kuramoto.TWO_PI)
        self.kuramoto_phases = theta.detach()

        # 3. MoE: transform the thought, routed by phases.
        h_flat = h[:, 0, :]  # (B, d_model) — squeeze the L=1 dim
        h_moe = self.norm_moe(h_flat)  # (B, d_model)
        # Compute gates from Kuramoto phases (squeeze L dim).
        theta_flat = theta[:, 0, :]  # (B, N_osc)
        sin_p = torch.sin(theta_flat).sum(dim=-1)
        cos_p = torch.cos(theta_flat).sum(dim=-1)
        theta_bar = torch.atan2(sin_p, cos_p)  # (B,)
        diff = theta_bar.unsqueeze(-1) - self.expert_phases.view(1, self.n_experts)
        gates = torch.softmax(self.kappa * torch.cos(diff), dim=-1)  # (B, E)
        topk_vals, topk_idx = gates.topk(self.top_k, dim=-1)  # (B, K)
        topk_norm = topk_vals / topk_vals.sum(dim=-1, keepdim=True).clamp(min=1e-10)

        moe_out = torch.zeros_like(h_moe)
        w1_stack = torch.stack([e._cached_W for e in self.experts_w1])  # (E, d_ff, D)
        w2_stack = torch.stack([e._cached_W for e in self.experts_w2])  # (E, D, d_ff)
        for k_slot in range(self.top_k):
            idx_k = topk_idx[:, k_slot]  # (B,)
            w_k = topk_norm[:, k_slot]  # (B,)
            w1_sel = w1_stack[idx_k]  # (B, d_ff, D) → transpose for matmul
            w2_sel = w2_stack[idx_k]  # (B, D, d_ff) → transpose for matmul
            h1 = torch.bmm(h_moe.unsqueeze(1), w1_sel.transpose(1,2)).squeeze(1)  # (B, d_ff)
            h1_act = F.gelu(h1)
            out_k = torch.bmm(h1_act.unsqueeze(1), w2_sel.transpose(1,2)).squeeze(1)  # (B, D)
            moe_out += w_k.unsqueeze(-1) * out_k
        h = h + moe_out.unsqueeze(1)  # add back the L dim

        # 4. Update thought state.
        self.thought_state = h.detach()

        # 5. Confidence + output.
        confidence = torch.sigmoid(self.confidence_head(h[:, 0, :]).squeeze(-1))  # (B,)
        output_logits = self.output_head(h[:, 0, :])  # (B, vocab)

        return output_logits, confidence

    def tick_chunk(self, observations: torch.Tensor) -> torch.Tensor:
        """Process a CHUNK of tokens in ONE forward pass (16x faster than tick).

        observations: (B, C) where C = chunk_len (e.g. 16).
        Returns logits: (B, C, vocab).

        This is the KEY SPEED OPTIMIZATION. Instead of calling tick() C times
        (C forward passes + C backward passes), we do ONE forward over the whole
        chunk. The attention state (S,z) is carried forward via the L8 state-carry
        mechanism — accumulated within the chunk, then detached for the next chunk.

        The thought state evolves across the whole chunk in a single graph.
        """
        B, C = observations.shape
        D = self.d_model

        # Embed the whole chunk.
        obs_vecs = self.observe(observations)  # (B, C, D)
        # Add the carried thought state to position 0.
        h = obs_vecs.clone()
        h[:, 0, :] = h[:, 0, :] + self.thought_state[:, 0, :]

        # 1. Attention: process the chunk with the accumulated state.
        # We use the L8 batched attention (heads×levels flattened).
        attn = self.attn
        nH, dH, nL = attn.n_heads, attn.d_head, attn.n_levels
        h_normed = self.norm_attn(h)
        q_all = torch.einsum("bld,de->ble", h_normed, attn.w_qkv[0]) + attn.b_qkv[0]
        k_all = torch.einsum("bld,de->ble", h_normed, attn.w_qkv[1]) + attn.b_qkv[1]
        v_all = torch.einsum("bld,de->ble", h_normed, attn.w_qkv[2]) + attn.b_qkv[2]
        q_all = q_all.view(B, C, nH, dH)
        k_all = k_all.view(B, C, nH, dH)
        v_all = v_all.view(B, C, nH, dH)

        offsets = attn.level_offsets
        q_lev = q_all.unsqueeze(1) + offsets.view(nL, 1, 1, 1)
        k_lev = k_all.unsqueeze(1) + offsets.view(nL, 1, 1, 1)
        q_feat = elu_plus_one(q_lev, alpha=1.0)
        k_feat = elu_plus_one(k_lev, alpha=1.0)
        v_lev = v_all.unsqueeze(1).expand(B, nL, C, nH, dH)
        q_flat = q_feat.permute(0, 1, 3, 2, 4).reshape(B * nL * nH, C, dH)
        k_flat = k_feat.permute(0, 1, 3, 2, 4).reshape(B * nL * nH, C, dH)
        v_flat = v_lev.permute(0, 1, 3, 2, 4).reshape(B * nL * nH, C, dH)
        # Vectorized causal attention. For chunk processing we skip the
        # cross-chunk state carry (the S,z would need per-head bookkeeping
        # which adds complexity). The chunk is long enough (16 tokens) that
        # intra-chunk attention captures sufficient context.
        y_flat = attn._linear_attention_causal_vectorized(q_flat, k_flat, v_flat)

        y = y_flat.reshape(B, nL, nH, C, dH).permute(0, 1, 3, 2, 4).reshape(B, nL, C, nH * dH)
        level_weights = torch.softmax(attn.level_logits, dim=-1)
        attn_out = (y * level_weights.view(1, nL, 1, 1)).sum(dim=1)
        attn_out = attn_out @ attn.w_out + attn.b_out
        h = h + attn_out

        # 2. Kuramoto: advance phases from the chunk's hidden states.
        h_kur = self.norm_kur(h)
        theta = self.kuramoto._encode_from_hidden(h_kur)
        theta = self.kuramoto._rk4_integrate(theta)
        theta_flat = theta[:, -1, :]  # (B, N_osc) — use last position's phases
        self.kuramoto_phases = theta.detach()

        # 3. MoE: transform the chunk (vectorized, cached weights).
        h_moe = self.norm_moe(h)  # (B, C, D)
        # Use the last-position phases for routing the whole chunk.
        theta_bar = torch.atan2(
            torch.sin(theta_flat).sum(-1), torch.cos(theta_flat).sum(-1)
        )  # (B,)
        diff = theta_bar.unsqueeze(-1) - self.expert_phases.view(1, self.n_experts)
        gates = torch.softmax(self.kappa * torch.cos(diff), dim=-1)  # (B, E)
        topk_vals, topk_idx = gates.topk(self.top_k, dim=-1)
        topk_norm = topk_vals / topk_vals.sum(-1, keepdim=True).clamp(min=1e-10)

        w1_stack = torch.stack([e._cached_W for e in self.experts_w1])  # (E, d_ff, D)
        w2_stack = torch.stack([e._cached_W for e in self.experts_w2])  # (E, D, d_ff)
        moe_out = torch.zeros_like(h_moe)
        for k_slot in range(self.top_k):
            idx_k = topk_idx[:, k_slot]  # (B,)
            w_k = topk_norm[:, k_slot]  # (B,)
            # Gather weights for each batch element, apply to all C positions.
            for b in range(B):
                w1 = w1_stack[idx_k[b]]  # (d_ff, D) → transpose for matmul
                w2 = w2_stack[idx_k[b]]  # (D, d_ff) → transpose for matmul
                h1 = h_moe[b] @ w1.T  # (C, D) @ (D, d_ff) → (C, d_ff)
                h1_act = F.gelu(h1)
                out_k = h1_act @ w2.T  # (C, d_ff) @ (d_ff, D) → (C, D)
                moe_out[b] += w_k[b] * out_k
        h = h + moe_out

        # 4. Update thought state (last position).
        self.thought_state = h[:, -1:, :].detach()

        # 5. Output logits for the whole chunk.
        output_logits = self.output_head(h)  # (B, C, vocab)
        return output_logits

    def think(self, observations: torch.Tensor, max_ticks: int = 10,
              confidence_threshold: float = 0.7) -> torch.Tensor:
        """Process a sequence of observations, thinking adaptively.

        For each observation, tick until confidence > threshold or max_ticks.
        Returns the output logits at the point of confidence.
        """
        B = observations.shape[0]
        outputs = []
        for t in range(observations.shape[1]):
            obs = observations[:, t]
            for tick in range(max_ticks):
                logits, conf = self.tick(obs if tick == 0 else None)
                if conf.mean().item() > confidence_threshold:
                    break
            outputs.append(logits)
        return torch.stack(outputs, dim=1)  # (B, L, vocab)

    @torch.no_grad()
    def generate(self, prompt_tokens: list, n_new: int = 50,
                 temperature: float = 0.8, top_k: int = 40,
                 chunk_len: int = 16) -> list:
        """Generate text using chunk-based processing + temperature sampling.

        Args:
            prompt_tokens: list of token ids to seed.
            n_new: number of NEW tokens to generate.
            temperature: sampling temperature (1.0=normal, <1=conservative).
            top_k: only sample from the top-k logits (0 = no truncation).
            chunk_len: tokens per chunk forward.
        Returns:
            list of all token ids (prompt + generated).
        """
        self.eval()
        self.reset_thought(batch_size=1)
        all_tokens = list(prompt_tokens)

        # Process the prompt in chunks.
        for start in range(0, len(all_tokens), chunk_len):
            chunk = torch.tensor([all_tokens[start:start + chunk_len]], dtype=torch.long)
            if chunk.shape[1] == 0:
                break
            logits = self.tick_chunk(chunk)  # (1, C, vocab)
            last_logits = logits[0, -1, :]  # (vocab,)

        # Generate new tokens one at a time (using the chunk's last logits).
        for _ in range(n_new):
            # Temperature + top-k sampling.
            logits = last_logits / max(temperature, 1e-8)
            if top_k > 0:
                topk_vals, topk_idx = logits.topk(min(top_k, logits.shape[-1]))
                probs = F.softmax(topk_vals, dim=-1)
                next_idx = torch.multinomial(probs, 1).item()
                next_token = topk_idx[next_idx].item()
            else:
                probs = F.softmax(logits, dim=-1)
                next_token = torch.multinomial(probs, 1).item()
            all_tokens.append(next_token)

            # Feed the new token via a chunk of 1.
            chunk = torch.tensor([[next_token]], dtype=torch.long)
            logits = self.tick_chunk(chunk)
            last_logits = logits[0, -1, :]

        self.train()
        return all_tokens