###############################################################################
# CogniThread Model v2 — with vocabulary head for text generation
###############################################################################
class CogniThreadV2(nn.Module):
    """Full CogniThread + vocab head for next-token prediction."""
    def __init__(self, d_model, d_state=16, n_heuristics=64, d_z=32,
                 T_diff=10, d_conv=4, expand=2, vocab_size=256, *,
                 rngs: nn.Rngs):
        super().__init__()
        self.d_model = d_model
        self.d_state = d_state
        self.d_z = d_z
        self.vocab_size = vocab_size
        # Byte-level embedding
        self.embed = nn.Embed(vocab_size, d_model, rngs=rngs)
        # Core CogniThread
        self.ssm = CogniThreadSSM(d_model, d_state, d_conv, expand, rngs=rngs)
        self.state_proj = nn.Linear(d_model, d_state, rngs=rngs)
        self.fuzzy = FuzzyHeuristicRetriever(d_state, n_heuristics, 1.0, 5, rngs=rngs)
        self.diffusion = LatentDiffusionSynthesizer(d_z, d_state, T_diff, rngs=rngs)
        self.diff_proj = nn.Linear(d_z, d_model, rngs=rngs)
        self.fuzzy_proj = nn.Linear(d_state, d_model, rngs=rngs)
        self.gate = AdaptiveGate(d_model, rngs=rngs)
        # Vocabulary head
        self.lm_head = nn.Linear(d_model, vocab_size, rngs=rngs)

    def __call__(self, input_ids, key=None):
        """input_ids: (batch, seq_len) int tokens. Returns logits."""
        if key is None: key = jax.random.key(0)
        x = self.embed(input_ids)  # (batch, seq_len, d_model)
        h_ssm = self.ssm(x)
        h_t = self.state_proj(h_ssm)
        retrieved, w_f, mu = self.fuzzy(h_t)
        retrieved = self.fuzzy_proj(retrieved)
        z_0 = self.diffusion(h_t, key=key)
        generated = self.diff_proj(z_0)
        y_final, var_y = self.gate(retrieved, generated, w_f)
        logits = self.lm_head(y_final)  # (batch, seq_len, vocab_size)
        return logits

    def generate(self, input_ids, max_new_tokens=64, key=None, temperature=0.8):
        """Autoregressive text generation."""
        if key is None: key = jax.random.key(0)
        batch, seq_len = input_ids.shape
        generated = list(input_ids.T)  # list of (batch,) arrays
        cur_ids = input_ids
        for i in range(max_new_tokens):
            key, sk = jax.random.split(key)
            logits = self(cur_ids, key=sk)  # (batch, seq, vocab)
            next_logits = logits[:, -1, :] / max(temperature, 0.01)
            next_token = jax.random.categorical(sk, next_logits, axis=-1)  # (batch,)
            generated.append(next_token)
            cur_ids = jnp.concatenate([cur_ids[:, 1:], next_token[:, None]], axis=1)
        return jnp.stack(generated, axis=1)  # (batch, seq_len + max_new_tokens)