"""v6 training contexts: the same item rendered as a separate sequence after no rationale (`direct`), a short stock-model rationale (`brief`) or a longer stock-model
thought (`deliberate`). The layouts are exactly those the rationales were generated and validated with (tools/render_for_gen.py, tools/score_rationales.py,
dense/think_gate2.py), so training and inference see the same tokens:

  direct       codec prompt, slot read right away (answer_slot.render_ar)
  brief        <|im_start|>system\\n{rat_s32}<|im_end|>\\n + codec prompt + "Rationale:\\n" + rationale (first `brief_cap` tokens) + "\\n\\nAnswer: " -> slot
  deliberate   thinking-template prompt + draft_d prefix + thought + "\\n</think>\\n\\n" -> slot
  per-option ('parallel') items: the generation prompt is cut after the assistant turn start, the rationale goes there, then "\\n\\n" (brief) or the close (deliberate)
  and the unchanged slot block ("1:" ... "k:") follows.

The rationale text was written for the item as displayed (identity option order, no augmentation), so the rationale contexts never shuffle or augment.
"""
import random
from typing import Any, Dict, List, Optional

import answer_slot as AS
from dense.rationale_prompts import SYSTEM
from dense.schema import codec_item

CONTEXTS = ("brief", "direct", "deliberate")
DRAFT_D = "Okay. Brief reasoning, a few words per step, no double-checking. First,"


def parse_kv(s: str) -> Dict[str, float]:
    return {k: float(v) for k, v in (p.split("=") for p in s.split(","))}


class Contexts:
    def __init__(self, tok: Any, c_off: AS.SlotCodec, prompt: str = "rat_s32", weights: str = "brief=.65,direct=.15,deliberate=.2",
                 loss_weights: str = "brief=1,direct=1,deliberate=.5", parallel_direct: float = 0.5, brief_cap: int = 64):
        self.tok, self.c_off, self.brief_cap = tok, c_off, brief_cap
        self.c_on = AS.SlotCodec(tok, state_max=None, option_max=None, template_kwargs={"enable_thinking": True}, mask_id=tok.pad_token_id or 0,
                                 strip_label_keys=True, autoregressive=True)
        enc = lambda t: tok(t, add_special_tokens=False)["input_ids"]
        self.sys_ids = enc(f"<|im_start|>system\n{SYSTEM[prompt]}<|im_end|>\n")
        self.pref, self.tail, self.sep = enc("Rationale:\n"), enc("\n\nAnswer: "), enc("\n\n")
        self.draft_d, self.close = enc(DRAFT_D), enc("\n</think>\n\n")
        self.weights, self.loss_w, self.parallel_direct = parse_kv(weights), parse_kv(loss_weights), parallel_direct

    # ---- which context an item trains in (deterministic in (seed, item id, epoch): a resumed run redraws the same contexts)
    def mode(self, ex: Dict[str, Any]) -> str:
        return AS.slot_mode(codec_item(ex)["qtype"], len(ex["target"]))

    def draw(self, ex: Dict[str, Any], seed: Any, epoch: int) -> str:
        rat = ex.get("rationale") or {}
        w = dict(self.weights)
        if self.mode(ex) == "parallel" and self.parallel_direct >= 1.0:   # v6b: per-option families never train after a rationale
            return "direct"
        if self.mode(ex) == "parallel":   # a rationale before the per-option slots hurt the stock slot: start with a heavier `direct` share
            moved = w["brief"] * self.parallel_direct
            w["brief"] -= moved; w["direct"] += moved
        if "deliberate" not in rat:
            w["brief"] += w.pop("deliberate", 0.0); w["deliberate"] = 0.0
        if "brief" not in rat:
            w["direct"] += w["brief"]; w["brief"] = 0.0
        names = [c for c in CONTEXTS if w.get(c, 0) > 0]
        return random.Random(f"{seed}:ctx:{ex['id']}:{epoch}").choices(names, weights=[w[c] for c in names])[0]

    def extra_tokens(self, ex: Dict[str, Any], ctx: str) -> int:
        """Length added by the context (for the solo-micro-batch / skip estimates)."""
        if ctx == "brief":
            return len(self.sys_ids) + len(self.pref) + len(self.tail) + min(ex["rationale"]["brief"]["n"], self.brief_cap)
        if ctx == "deliberate":
            return len(self.draft_d) + len(self.close) + ex["rationale"]["deliberate"]["n"]
        return 0

    # ---- generation side (inference): the same prompt the stored rationales were generated from, and the same clean-up of the generated text
    def gen_prompt(self, ex: Dict[str, Any]) -> List[int]:
        """Token ids of the stock-model rationale request (brief): system line + codec prompt (cut after the assistant start for per-option items) + "Rationale:\\n"."""
        r = self.c_off.render_ar(codec_item(ex))
        ids = r["ids"][:self._cut(self.c_off, r)] if r["mode"] == "parallel" else r["ids"]
        return self.sys_ids + ids + self.pref

    def clean_brief(self, g: List[int]) -> List[int]:
        """Generated tokens (stop token removed) -> stored form: a trailing bare-answer line (<= 12 chars after a newline) is dropped so the slot is not told the answer
        (as tools/gen_rationales.py)."""
        txt = self.tok.decode(g).rstrip()
        last = txt.split("\n")[-1].strip()
        if g and len(last) <= 12 and "\n" in txt:
            g = self.tok(txt[: -len(last)].rstrip() + "\n", add_special_tokens=False)["input_ids"]
        return g

    # ---- rendering
    @staticmethod
    def _cut(codec: AS.SlotCodec, r: Dict[str, Any]) -> int:
        ids, post = r["ids"], codec.post
        return next(i + len(post) for i in range(r["state_end"], len(ids) - len(post) + 1) if ids[i:i + len(post)] == post)

    def render(self, ex: Dict[str, Any], ctx: str) -> Dict[str, Any]:
        """`brief` / `deliberate` only; `direct` is rendered by the caller (it keeps the augmentations and the option shuffle)."""
        assert ctx in ("brief", "deliberate"), ctx
        it = codec_item(ex)
        rat = ex["rationale"][ctx]
        if ctx == "brief":
            codec, head, body = self.c_off, self.sys_ids, rat["tokens"][:self.brief_cap]
            pre_tail, tail_single, tail_par = self.pref, self.tail, self.sep
        else:
            codec, head, body = self.c_on, [], rat["tokens"]
            pre_tail, tail_single, tail_par = self.draft_d, self.close, self.close
        r = codec.render_ar(it)   # identity order
        if r["mode"] == "parallel":
            cut = self._cut(codec, r)
            new = head + r["ids"][:cut] + pre_tail + body + tail_par + r["ids"][cut:]
            off = len(head) + len(pre_tail) + len(body) + len(tail_par)
            return {**r, "ids": new, "slot_pos": [p + off for p in r["slot_pos"]]}
        ids = head + r["ids"] + pre_tail + body + tail_single
        return {**r, "ids": ids, "slot_pos": [len(ids) - 1]}