"""Structured SOP decoding for the analyst persona.

The scratchpad is generated freely; the verdict and confidence sections are
decoded greedily under a controlled vocabulary, so the model cannot emit an
unstructured or off-protocol answer for the final judgment.

Usage:
  from research.structured import analyst_report
  report = analyst_report(model, tok, doc_text, max_scratch=140)
"""

import torch

VERDICTS = [
    "true", "false", "mostly true", "partially true", "mixed",
    "unsubstantiated", "unsupported", "overclaim", "misleading",
    "inaccurate", "unverifiable", "cannot confirm",
    "not a discrepancy", "conflict", "no meaningful pattern",
    # canonical set (eval_labels.CANON) — the tiny head must be able to emit these
    "refutes", "not enough information", "not a contradiction",
    "contradiction", "low confidence", "abstain", "cannot provide",
]
CONFIDENCES = ["HIGH", "MEDIUM", "LOW", "cannot assess"]


def _decode_phrase(model, tok, ids, persona_id, allowed_phrases, max_new=16, fallback=None):
    """Constrained greedy decode over the allowed phrases (longest-prefix).

    Returns (exact_phrase_token_ids, ok). Never emits partial/garbage tokens:
    - if a full allowed phrase is matched, its EXACT tokens are returned,
    - otherwise the caller falls back (default: explicit abstain text),
      so the model can never present an off-protocol verdict.
    """
    targets = [tok.encode(p).ids for p in allowed_phrases]
    targets = [t for t in targets if t]
    if not targets:
        return [], False
    progress = [0] * len(targets)
    seq = torch.tensor([ids], dtype=torch.long)
    for _ in range(max_new):
        window = seq[:, -model.cfg.max_seq_len:]
        logits = model(window, persona_ids=torch.tensor([persona_id]) if persona_id else None)
        logits = logits[:, -1, :]
        allowed = {}
        for i, t in enumerate(targets):
            if progress[i] < len(t):
                allowed.setdefault(t[progress[i]], i)
        if not allowed:
            break
        allowed_ids = torch.tensor(list(allowed.keys()), dtype=torch.long)
        mask = torch.full_like(logits, -float("inf"))
        mask[:, allowed_ids] = logits[:, allowed_ids]
        nxt = int(mask.argmax().item())
        seq = torch.cat([seq, torch.tensor([[nxt]], dtype=torch.long)], dim=1)
        for i, t in enumerate(targets):
            if progress[i] < len(t) and t[progress[i]] == nxt:
                progress[i] += 1
                if progress[i] == len(t):
                    return t, True  # exact full-phrase tokens
        # if we advanced no target further, stop (no clean path)
        if not any(progress[i] > 0 and progress[i] < len(targets[i]) for i in range(len(targets))):
            break
    return [], False


@torch.no_grad()
def analyst_report(model, tok, doc, persona_id=1, max_scratch=140, max_reason=80):
    """Two-pass structured analysis: scratchpad -> verdict -> confidence -> reasoning.

    The verdict and confidence are CONSTRAINT-DECODED with a hard fallback: if the
    head does not cleanly emit an allowed phrase, we return an explicit abstention
    ("not enough information" / "cannot assess") and flag decode_failed=True.
    Garbage is never presented as a verdict.
    """
    persona_tok = "<|analyst|>"
    prompt = persona_tok + "<|user|>" + doc + "<|assistant|>"
    ids = tok.encode(prompt).ids

    # free-form scratchpad
    scratch = model.generate(tok, ids, persona_id=persona_id, max_new=max_scratch,
                             temperature=0.6, top_k=40, repetition_penalty=1.4,
                             no_repeat_ngram_size=4)
    ids = scratch

    # constrained verdict (hard fallback -> abstention)
    vids, vok = _decode_phrase(model, tok, ids, persona_id, VERDICTS)
    if vok:
        verdict = tok.decode(vids).strip()
    else:
        verdict = "not enough information"
        vids = tok.encode(verdict).ids
    ids = ids + vids

    # constrained confidence
    ids = ids + tok.encode(" Confidence: ").ids
    cids, cok = _decode_phrase(model, tok, ids, persona_id, CONFIDENCES)
    if cok:
        confidence = tok.decode(cids).strip()
    else:
        confidence = "cannot assess"
        cids = tok.encode(confidence).ids
    ids = ids + cids

    # free-form reasoning
    ids = ids + tok.encode(" Reasoning: ").ids
    reason = model.generate(tok, ids, persona_id=persona_id, max_new=max_reason,
                            temperature=0.6, top_k=40, repetition_penalty=1.4,
                            no_repeat_ngram_size=4)
    reasoning = tok.decode(reason[len(ids):]).strip()

    return {
        "scratchpad": tok.decode(scratch[len(tok.encode(prompt).ids):]).strip(),
        "verdict": verdict,
        "confidence": confidence,
        "reasoning": reasoning,
        "decode_failed": not (vok and cok),
    }