"""Graft: splice connector soft tokens into the frozen Jais-3-0.9B (K2HorizonForCausalLM) through inputs_embeds.
Prompt (same chat format as serve_mix2.prompt / train_lora.encode, fastest thinking mode):
  <|begin_of_text|><|ifm|im_start|>system\n{SYSTEM}<|ifm|im_end|><|ifm|im_start|>user\n <image> [soft x144] </image> {question} <|ifm|im_end|><|ifm|im_start|>assistant\n<ifm|think_faster>\n</ifm|think_faster> {answer}<|ifm|im_end|>
The two marker rows (64031/64032) are replaced at splice time by two learned vectors (initialised from the real rows); the embedding matrix itself is never unfrozen.
Labels only on the answer (+ end-of-turn).  Right padding for training, left padding (with correct position ids) for generation."""
import os, sys, torch, torch.nn as nn, torch.nn.functional as F
HERE = os.path.dirname(os.path.abspath(__file__)); BRAIN = os.path.dirname(HERE)
if HERE not in sys.path: sys.path.insert(0, HERE)
from connector import Connector

# Local Jais copy if JAIS_PATH is set and exists; otherwise the public Hub repo (auto-downloaded).
BASE = os.environ.get("JAIS_PATH") or os.path.join(BRAIN, "jais3_09b")
if not os.path.isdir(BASE): BASE = os.environ.get("JAIS_REPO") or "inception42/Jais-3-0.9B"
IMG, IMG_END = 64031, 64032
S, EN = "<|ifm|im_start|>", "<|ifm|im_end|>"
SYSTEM = "أنت «تصحيح»، مساعد عربي للكتابة واللغة. تتكلم بطبيعية واختصار مثل صديق، بلهجة المستخدم نفسها، وبالعربية فقط."   # same default as train_lora.py
HEADER = "<ifm|think_faster>\n</ifm|think_faster>"


def load_trunk(device="cpu", dtype=None, layers=None):
    """TRUNK_LAYERS / TRUNK_DTYPE=fp32 are TEST-ONLY knobs (a truncated real-weight trunk in fp32 makes CPU tests fast; the CPU has no bf16 units)"""
    from transformers import AutoTokenizer, AutoModelForCausalLM
    dtype = dtype or (torch.float32 if os.environ.get("TRUNK_DTYPE") == "fp32" else torch.bfloat16); layers = layers or int(os.environ.get("TRUNK_LAYERS", 0)) or None
    tok = AutoTokenizer.from_pretrained(BASE, trust_remote_code=True)
    m = AutoModelForCausalLM.from_pretrained(BASE, trust_remote_code=True, dtype=dtype).to(device).eval()
    if layers: m.model.layers = m.model.layers[:layers]; m.config.num_hidden_layers = layers
    for p in m.parameters(): p.requires_grad_(False)
    return tok, m


class Graft(nn.Module):
    def __init__(self, trunk, tok, connector=None, system=SYSTEM):
        super().__init__()
        object.__setattr__(self, "trunk", trunk)               # held by reference, NOT a submodule: Graft.parameters() = connector + markers only
        self.tok = tok; self.system = system
        dev = trunk.get_input_embeddings().weight.device
        self.connector = (connector or Connector()).to(dev)
        E = trunk.get_input_embeddings().weight
        self.markers = nn.Parameter(E[[IMG, IMG_END]].detach().float().clone())      # [2,1536] learned <image> and </image> vectors
        self.end_id = tok.convert_tokens_to_ids(EN)
        self._ids = lambda s: tok(s, add_special_tokens=False).input_ids

    # ------------------------------------------------------------------ prompt pieces
    def pieces(self, question, answer=None, has_image=True):
        """-> (pre ids, post ids, answer ids).  pre ends right before <image>; post starts right after </image>."""
        ids = self._ids
        head = ids("<|begin_of_text|>" + S + "system\n" + self.system + EN + S + "user\n")
        tail = ids(("" if not has_image else "") + question + EN + S + "assistant\n" + HEADER)
        ans = ids(answer + EN) if answer is not None else []
        return head, tail, ans

    def text_ids(self, question):
        """the ids of the SAME prompt with no image (used by the G0 parity test)"""
        head, tail, _ = self.pieces(question, has_image=False)
        return head + tail

    # ------------------------------------------------------------------ embeddings
    def _embed(self, ids):
        E = self.trunk.get_input_embeddings()
        with torch.no_grad(): return E(torch.tensor(ids, device=E.weight.device))

    def build(self, feats, questions, answers=None, low=False, left_pad=False):
        """feats [B,576,768] (or None -> text only); returns dict(inputs_embeds, attention_mask, position_ids, labels, n_img)
        dtype of inputs_embeds = trunk dtype.  Gradients flow to connector + markers only."""
        dt = self.trunk.get_input_embeddings().weight.dtype; dev = self.markers.device
        soft = self.connector(feats.to(dev), low=low) if feats is not None else None
        n_img = soft.shape[1] if soft is not None else 0
        rows, labs = [], []
        for b, q in enumerate(questions):
            head, tail, ans = self.pieces(q, None if answers is None else answers[b], has_image=feats is not None)
            parts = [self._embed(head)]
            if soft is not None: parts += [self.markers[0:1].to(dt), soft[b].to(dt), self.markers[1:2].to(dt)]
            parts += [self._embed(tail)]
            if ans: parts += [self._embed(ans)]
            e = torch.cat(parts, 0); rows.append(e)
            L = e.shape[0]; y = torch.full((L,), -100, dtype=torch.long);
            if ans: y[L - len(ans):] = torch.tensor(ans)
            labs.append(y)
        T = max(r.shape[0] for r in rows); B = len(rows); H = rows[0].shape[1]
        emb = torch.zeros(B, T, H, dtype=dt, device=dev); att = torch.zeros(B, T, dtype=torch.long); lab = torch.full((B, T), -100, dtype=torch.long)
        for b, (r, y) in enumerate(zip(rows, labs)):
            L = r.shape[0]; s = T - L if left_pad else 0
            emb[b, s:s + L] = r; att[b, s:s + L] = 1; lab[b, s:s + L] = y
        att = att.to(dev); lab = lab.to(dev)
        pos = (att.cumsum(1) - 1).clamp_min(0)
        return dict(inputs_embeds=emb, attention_mask=att, position_ids=pos, labels=lab, n_img=n_img)

    # ------------------------------------------------------------------ loss (logits only at answer positions)
    def loss(self, feats, questions, answers, low=False, return_tokens=False):
        d = self.build(feats, questions, answers, low=low)
        h = self.trunk.model(inputs_embeds=d["inputs_embeds"], attention_mask=d["attention_mask"], position_ids=d["position_ids"], use_cache=False).last_hidden_state
        lab = d["labels"][:, 1:]; hs = h[:, :-1]
        m = lab != -100
        logits = self.trunk.lm_head(hs[m]).float()
        ce = F.cross_entropy(logits, lab[m], reduction="sum")
        n = int(m.sum())
        return (ce / max(n, 1), ce, n) if return_tokens else ce / max(n, 1)

    # ------------------------------------------------------------------ greedy generation (own loop: inputs_embeds + KV cache + left padding)
    def iter_tokens(self, feats, questions, max_new=48, low=False, use_cache=True, stop_ids=None, step_ctx=None):
        """generator: yields a list with the next token id of every batch row at each step (rows that already stopped repeat their stop token).
        step_ctx() -> context manager entered around each forward only (so no_grad / autocast / adapter-off never leak to the caller while suspended)."""
        import contextlib
        from transformers import DynamicCache
        step_ctx = step_ctx or contextlib.nullcontext
        with torch.no_grad(): d = self.build(feats, questions, None, low=low, left_pad=True)
        emb, att, pos = d["inputs_embeds"], d["attention_mask"], d["position_ids"]
        B = emb.shape[0]; stop = {self.end_id, self.tok.eos_token_id, *(stop_ids or [])}
        E = self.trunk.get_input_embeddings(); done = [False] * B
        cache = DynamicCache(config=self.trunk.config) if use_cache else None
        cur_emb, cur_att, cur_pos = emb, att, pos; last = pos[:, -1:]
        for step in range(max_new):
            with torch.no_grad(), step_ctx():
                o = self.trunk.model(inputs_embeds=cur_emb, attention_mask=cur_att, position_ids=cur_pos, past_key_values=cache, use_cache=use_cache)
                nxt = self.trunk.lm_head(o.last_hidden_state[:, -1]).float().argmax(-1)
                ne = E(nxt).unsqueeze(1)
            toks = [int(t) for t in nxt]; yield toks
            for b, t in enumerate(toks): done[b] = done[b] or t in stop
            if all(done): return
            last = last + 1
            cur_att = torch.cat([cur_att, torch.ones(B, 1, dtype=cur_att.dtype, device=cur_att.device)], 1)
            if use_cache: cur_emb, cur_pos = ne, last
            else: emb = torch.cat([emb, ne], 1); pos = torch.cat([pos, last], 1); cur_emb, cur_pos = emb, pos

    def generate(self, feats, questions, max_new=48, low=False, use_cache=True, return_ids=False, stop_ids=None, step_ctx=None):
        stop = {self.end_id, self.tok.eos_token_id, *(stop_ids or [])}; out = [[] for _ in questions]; fin = [False] * len(questions)
        for toks in self.iter_tokens(feats, questions, max_new, low, use_cache, stop_ids, step_ctx):
            for b, t in enumerate(toks):
                if not fin[b]: out[b].append(t); fin[b] = t in stop
        if return_ids: return out
        return [self.tok.decode([i for i in ids if i not in stop], skip_special_tokens=True).strip() for ids in out]

    # ------------------------------------------------------------------ trainable state
    def trainable(self): return [p for p in self.parameters() if p.requires_grad]

    def state(self): return {"connector": self.connector.state_dict(), "markers": self.markers.detach().cpu()}

    def load_state(self, st):
        self.connector.load_state_dict(st["connector"]); self.markers.data.copy_(st["markers"].to(self.markers.device))