"""E1 (PLAN_SELFPACED): learned per-prompt halting over the workspace loop.

HaltingMergeAdapter = frozen-recipe MergeAdapter + a halting head read off
the workspace state at the last prompt position after each iteration:

    p_i = sigmoid(w · [e_last ; s_hat_i,last] + b)        (halt after iter i)
    q_i = p_i * prod_{j<i}(1 - p_j),  q_kmax += remainder (ACT-style)

Training (soft, no RL): s_mix = sum_i q_i * s_i  -> suffix -> CE, plus a
compute penalty  lambda * E[iters] = lambda * sum_i q_i * i.
The halting head's weight is zero-init and its bias starts at -6
(p ~ 0.0025), so at init >=98% of the mass sits on k_max: the model is
approximately the fixed-k merge (documented tolerance, not bit-exact).

Deploy: iterate until cumulative halt mass crosses 0.5; per-item k*.
"""

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

from loop_common import ALPHA, MergeAdapter


class HaltingMergeAdapter(MergeAdapter):

    def __init__(self, d=1536, hidden=512, alpha=ALPHA, bias0=-6.0):
        super().__init__(d=d, hidden=hidden, alpha=alpha)
        self.halt = nn.Linear(2 * d, 1)
        nn.init.zeros_(self.halt.weight)
        nn.init.constant_(self.halt.bias, bias0)

    def halt_prob(self, e_last, s_last):
        """p(halt | state) from the last prompt position. (B,) in (0,1)."""
        e32, s32 = e_last.float(), s_last.float()
        s_hat = s32 * (e32.norm(dim=-1, keepdim=True)
                       / (s32.norm(dim=-1, keepdim=True) + 1e-6))
        return torch.sigmoid(self.halt(torch.cat([e32, s_hat], -1))).squeeze(-1)


def halting_loop_logits(looper, adapter, input_ids, k_max, attention_mask,
                        prompt_last_idx, loop_mask=None, use_checkpoint=True):
    """Soft-halting teacher-forced logits.

    prompt_last_idx: (B,) index of each item's last prompt position.
    Returns (logits_of_mixed_state, q, exp_iters):
      q (B, k_max) halting distribution over iterations 1..k_max.
    """
    calls, _ = looper.capture(input_ids, attention_mask, logits_to_keep=1)
    e = looper._hin[looper.l0].detach()
    B = e.shape[0]
    bidx = torch.arange(B, device=e.device)
    with torch.no_grad():
        s = looper.band(e, calls)
    s_mix = torch.zeros_like(e)
    keep = torch.ones(B, device=e.device)          # prob of not-yet-halted
    qs = []
    for i in range(k_max):
        x = adapter(e, s)
        if loop_mask is not None:
            x = torch.where(loop_mask[..., None], x, e)
        if use_checkpoint:
            s = checkpoint(lambda x_: looper.band(x_, calls), x,
                           use_reentrant=False)
        else:
            s = looper.band(x, calls)
        p = adapter.halt_prob(e[bidx, prompt_last_idx],
                              s[bidx, prompt_last_idx])
        q_i = keep * p if i < k_max - 1 else keep   # remainder to k_max
        keep = keep * (1 - p)
        s_mix = s_mix + q_i[:, None, None] * s.float()
        qs.append(q_i)
    q = torch.stack(qs, -1)                          # (B, k_max)
    exp_iters = (q * torch.arange(1, k_max + 1, device=q.device)).sum(-1)
    logits = looper.suffix_logits(s_mix.to(e.dtype), calls)
    return logits, q, exp_iters


@torch.no_grad()
def halted_k_per_item(looper, adapter, input_ids, k_max, attention_mask,
                      prompt_last_idx, thresh=0.5, allow_k0=True):
    """Deploy-time halting: smallest k where cumulative halt mass >= thresh
    (k_max if never). With allow_k0, the head is also consulted on the
    pre-loop state s_0 — items it flags there get k*=0 (no looping at
    all; the E0 lesson: easy items are safest untouched).
    Returns (B,) ints in 0..k_max."""
    calls, _ = looper.capture(input_ids, attention_mask, logits_to_keep=1)
    e = looper._hin[looper.l0].detach()
    B = e.shape[0]
    bidx = torch.arange(B, device=e.device)
    s = looper.band(e, calls)
    kstar = torch.full((B,), k_max, dtype=torch.long, device=e.device)
    cum = torch.zeros(B, device=e.device)
    keep = torch.ones(B, device=e.device)
    if allow_k0:
        p0 = adapter.halt_prob(e[bidx, prompt_last_idx],
                               s[bidx, prompt_last_idx])
        kstar[p0 >= thresh] = 0
        cum = cum + p0
        keep = keep * (1 - p0)
    for i in range(k_max):
        x = adapter(e, s)
        s = looper.band(x, calls)
        p = adapter.halt_prob(e[bidx, prompt_last_idx],
                              s[bidx, prompt_last_idx])
        cum = cum + keep * p
        keep = keep * (1 - p)
        newly = (cum >= thresh) & (kstar == k_max) & (i < k_max - 1)
        kstar[newly] = i + 1
    return kstar