"""Instruct mode: chat models choose every move by replying with it.

The model never sees the grid. Every legal placement is described by its outcome in a few words (lines, holes, stack
height, surface), placements with the same description become one move, and the model gets up to GROUP moves per
question in a chat conversation. It replies with a move copied from the list. With more moves the piece plays out as a
knockout: groups of up to GROUP, then the group winners.

The reply is read, not generated, so there is nothing to parse and no invalid answer. Each question is asked twice,
"Best move?" and "Worst move?", in the same conversation, and a move's score is how much more likely the model is to
reply with it to the first than to the second (average log-probability per token). That mirrors Base mode's
"good move" minus "bad move": what makes a reply likely no matter what was asked (a familiar phrase, a move that
repeats words of the conversation, its place in the list) cancels out, and what is left is whether the model picks
moves according to the question. Two pilots on the pool's models led here: letter answers (A/B/C/D) were close to
chance level, and scoring only "Best move?" replies mostly measured which phrases a model likes (a model picked "good"
moves just as often with no rules and no example in the conversation).

The conversation is multi-turn: two fixed example exchanges (the same for every model), then the model's own previous
question and reply, then the current question, which starts with what the last move did.
"""
from __future__ import annotations

import copy
import random

import torch

from players import ModelRejected

SYSTEM = (
    "You play Tetris. Each turn shows the next piece and a few moves.\n"
    "Good moves clear lines, make no holes and keep the stack low and flat.\n"
    "Reply with the move asked for, copied exactly."
)
GROUP = 8  # moves per question
BEST, WORST = "Best move?", "Worst move?"
# First exchanges of every conversation, identical for every model: one "best" and one "worst" question with obvious
# answers, so the model sees both kinds of question and the format of a reply before its first real move.
EXAMPLES = [
    ("Piece: O\n- no lines, 2 new holes, high stack, bumpy\n- clears 1 line, no new holes, low stack, flat\n" + BEST,
     "clears 1 line, no new holes, low stack, flat"),
    ("Piece: S\n- no lines, no new holes, low stack, flat\n- no lines, 1 new hole, high stack, uneven\n" + WORST,
     "no lines, 1 new hole, high stack, uneven"),
]

# Repos whose tokenizer has no chat template but whose model card documents the prompt format they were tuned on.
_ALPACA = (
    "{{ bos_token }}Below is an instruction that describes a task. Write a response that appropriately completes the "
    "request.\n\n{% set ns = namespace(sys='') %}{% for m in messages %}{% if m['role'] == 'system' %}"
    "{% set ns.sys = m['content'] + '\n\n' %}{% elif m['role'] == 'user' %}"
    "{{ '### Instruction:\n' + ns.sys + m['content'] + '\n\n### Response:\n' }}{% set ns.sys = '' %}"
    "{% else %}{{ m['content'] + eos_token + '\n\n' }}{% endif %}{% endfor %}"
)
CHAT_TEMPLATES = {
    "supralabs/supra-50m-instruct": _ALPACA,  # Alpaca format (see the repo's sft.py)
}


class NotChatModel(ModelRejected):
    """The model can't play Instruct (no chat template, or the conversation does not fit its context)."""


# ----------------------------------------------------------------------------
# Moves and questions
# ----------------------------------------------------------------------------
_LINES = {0: "no lines", 1: "clears 1 line", 2: "clears 2 lines", 3: "clears 3 lines", 4: "clears 4 lines"}


def short(c) -> str:
    """Outcome of a placement in four facts, always in the same order: lines, holes, height, surface."""
    if c.new_holes < 0:
        holes = "removes holes"
    elif c.new_holes == 0:
        holes = "no new holes"
    elif c.new_holes == 1:
        holes = "1 new hole"
    elif c.new_holes == 2:
        holes = "2 new holes"
    else:
        holes = "many new holes"
    mh = c.max_height
    height = ("very low stack" if mh <= 4 else "low stack" if mh <= 8 else "high stack" if mh <= 12
              else "very high stack" if mh <= 16 else "stack at the top")
    surface = "flat" if c.bump <= 4 else "uneven" if c.bump <= 10 else "bumpy"
    return f"{_LINES[c.lines]}, {holes}, {height}, {surface}"


def question(piece: str, moves, last: str | None = None, ask: str = BEST) -> str:
    lines = [f"Last move: {last}."] if last else []
    lines.append(f"Piece: {piece}")
    lines += [f"- {m}" for m in moves]
    lines.append(ask)
    return "\n".join(lines)


def knockout(game, seed: int, cands, decide):
    """Pick a placement for the current piece. `decide(history, piece, last, moves)` returns one score per move.

    Every random draw is seeded by (seed, piece index), so all players see the moves of a piece in the same order.
    Returns (candidate, score, reply, number of moves); the conversation state lives in game.memory."""
    memo = game.memory.setdefault("instruct", {"prev": None, "last": None})
    by_text = {}
    for c in cands:
        by_text.setdefault(short(c), []).append(c)
    moves = list(by_text)
    random.Random(f"{seed}:{game.pieces}:moves").shuffle(moves)
    history = EXAMPLES + ([memo["prev"]] if memo["prev"] else [])
    final = None  # (question, reply, score) of the last question asked for this piece
    rnd = 0
    while len(moves) > 1:
        n_groups = -(-len(moves) // GROUP)
        groups = [moves[i::n_groups] for i in range(n_groups)]  # balanced: sizes differ by one at most
        winners = []
        for gi, group in enumerate(groups):
            if len(group) == 1:
                winners.append(group[0])
                continue
            scores = decide(history, game.current_piece, memo["last"], group)
            top = max(scores)
            ties = [i for i, s in enumerate(scores) if s >= top - 1e-9]
            pick = ties[random.Random(f"{seed}:{game.pieces}:{rnd}:{gi}").randrange(len(ties))]
            winners.append(group[pick])
            final = (question(game.current_piece, group, memo["last"]), group[pick], scores[pick])
        moves = winners
        rnd += 1
    text = moves[0]
    same = by_text[text]
    cand = same[random.Random(f"{seed}:{game.pieces}").randrange(len(same))]
    if final:
        memo["prev"] = final[:2]
    memo["last"] = text
    return cand, (final[2] if final else None), (final[1] if final else ""), len(by_text)


# ----------------------------------------------------------------------------
# Talking to the model
# ----------------------------------------------------------------------------
def template_of(player):
    """The chat template a player uses (its tokenizer's, or a documented format), or None."""
    return CHAT_TEMPLATES.get(player.model_id.lower()) or getattr(player.tok, "chat_template", None)


def is_chat(player) -> bool:
    return getattr(player, "tok", None) is not None and bool(template_of(player))


def messages(history, user: str):
    """history: [(question, reply), ...] before the current question."""
    out = [{"role": "system", "content": SYSTEM}]
    for u, a in history:
        out += [{"role": "user", "content": u}, {"role": "assistant", "content": a}]
    out.append({"role": "user", "content": user})
    return out


def render(player, msgs) -> str:
    """The conversation as the model's prompt, ending where its reply starts. Templates without a system role get the
    system text at the start of the first user message."""
    tpl = template_of(player)
    try:
        text = player.tok.apply_chat_template(msgs, chat_template=tpl, tokenize=False, add_generation_prompt=True)
        if SYSTEM.splitlines()[0] in text:
            return text
    except Exception:
        pass
    merged = [dict(m) for m in msgs if m["role"] != "system"]
    merged[0]["content"] = SYSTEM + "\n\n" + merged[0]["content"]
    return player.tok.apply_chat_template(merged, chat_template=tpl, tokenize=False, add_generation_prompt=True)


def _encode(tok, text):
    return list(tok(text, add_special_tokens=False)["input_ids"])


def _reply_ids(tok, text, ids, reply):
    """Token ids of `reply` right after the prompt `text` (tokenized as ids)."""
    full = _encode(tok, text + reply)
    return full[len(ids):] if len(full) > len(ids) and full[:len(ids)] == ids else _encode(tok, reply)


def _full(player, common, rows):
    """rows: [(suffix, reply)] after the shared prefix `common`. Average log-probability per token of each reply."""
    raw = player._logprob_batch([(common + s + r, len(common) + len(s)) for s, r in rows])
    return [x / len(r) for x, (_, r) in zip(raw, rows)]


def _opening(player):
    """Token ids and KV cache of the start that every conversation shares (system prompt and examples), or False."""
    tok = player.tok
    a = _encode(tok, render(player, messages(EXAMPLES, question("T", ["x"]))))
    b = _encode(tok, render(player, messages(EXAMPLES + [(question("Z", ["y"], "w"), "y")], question("I", ["q"], "v"))))
    n = 0
    while n < min(len(a), len(b)) and a[n] == b[n]:
        n += 1
    if n < 8:
        return False
    past = player.model(input_ids=torch.tensor([a[:n]]), attention_mask=torch.ones((1, n), dtype=torch.long),
                        use_cache=True).past_key_values
    return (a[:n], past) if past is not None else False


def _start(player, common):
    """KV cache after `common`: the shared opening comes from a copy of its cache, computed once per model."""
    model = player.model
    if "_instruct_opening" not in player.__dict__:
        player._instruct_opening = _opening(player)
    opening = player._instruct_opening
    if opening and len(common) > len(opening[0]) and common[: len(opening[0])] == opening[0]:
        rest = torch.tensor([common[len(opening[0]):]])
        mask = torch.ones((1, len(common)), dtype=torch.long)
        return model(input_ids=rest, attention_mask=mask, past_key_values=copy.deepcopy(opening[1]), use_cache=True).past_key_values
    inp = torch.tensor([common])
    return model(input_ids=inp, attention_mask=torch.ones_like(inp), use_cache=True).past_key_values


def _cached(player, common, rows):
    """Same as _full, computing the shared prefix once and every suffix + reply together on its KV cache."""
    model = player.model
    past = _start(player, common)
    past.batch_repeat_interleave(len(rows))  # AttributeError on caches that can't be shared: caller falls back
    width = max(len(s) + len(r) for s, r in rows)
    ids = torch.full((len(rows), width), player.pad_id, dtype=torch.long)
    mask = torch.zeros((len(rows), len(common) + width), dtype=torch.long)
    mask[:, : len(common)] = 1
    for i, (s, r) in enumerate(rows):
        ids[i, : len(s) + len(r)] = torch.tensor(s + r)
        mask[i, len(common): len(common) + len(s) + len(r)] = 1
    lp = torch.log_softmax(model(input_ids=ids, attention_mask=mask, past_key_values=past, use_cache=False).logits.float(), dim=-1)
    out = []
    for i, (s, r) in enumerate(rows):
        out.append(sum(lp[i, len(s) + j - 1, r[j]].item() for j in range(len(r))) / len(r))
    return out


def _copied(player, common, rows):
    """Same as _full for caches that can't be batched (e.g. hybrid convolution layers): the shared prefix is computed
    once, then every suffix + reply runs on its own copy of the cache."""
    model = player.model
    past = _start(player, common)
    if past is None:
        raise ValueError("the model returns no cache")
    out = []
    for s, r in rows:
        ids = torch.tensor([s + r])
        mask = torch.ones((1, len(common) + len(s) + len(r)), dtype=torch.long)
        logits = model(input_ids=ids, attention_mask=mask, past_key_values=copy.deepcopy(past), use_cache=True).logits
        lp = torch.log_softmax(logits[0].float(), dim=-1)
        out.append(sum(lp[len(s) + j - 1, r[j]].item() for j in range(len(r))) / len(r))
    return out


_WAYS = {"batch": _cached, "copy": _copied, "full": _full}


@torch.inference_mode()
def contrast_scores(player, best_text: str, worst_text: str, moves) -> list[float]:
    """For each move: average log-probability per token of replying with it to the "best" prompt minus the same for the
    "worst" prompt. The two prompts differ only at the end, so the shared prefix is computed once."""
    tok = player.tok
    b, w = _encode(tok, best_text), _encode(tok, worst_text)
    n = 0
    while n < min(len(b), len(w)) - 1 and b[n] == w[n]:
        n += 1
    common = b[:n]
    rows = [(b[n:], _reply_ids(tok, best_text, b, m)) for m in moves] + [(w[n:], _reply_ids(tok, worst_text, w, m)) for m in moves]
    longest = max(len(b), len(w)) + max(len(r) for _, r in rows)
    if longest > player.max_len:
        raise NotChatModel(f"{player.model_id}: the conversation ({longest} tokens) exceeds its context ({player.max_len}).")
    way = player.__dict__.get("_instruct_way")
    if way is None:  # first question: use the fastest way that gives the same scores as the full computation
        scores = _full(player, common, rows)
        player._instruct_way = "full"
        for name in ("batch", "copy"):
            try:
                if max(abs(a - c) for a, c in zip(_WAYS[name](player, common, rows), scores)) < 1e-3:
                    player._instruct_way = name
                    break
            except Exception:
                continue
    else:
        scores = _WAYS[way](player, common, rows)
    k = len(moves)
    return [scores[i] - scores[k + i] for i in range(k)]


def _prompts(player, history, piece, last, moves):
    """The "best" and "worst" prompts; drops the previous exchange (never the examples) if they would not fit."""
    history = list(history)
    while True:
        best = render(player, messages(history, question(piece, moves, last, BEST)))
        worst = render(player, messages(history, question(piece, moves, last, WORST)))
        if len(_encode(player.tok, best)) + 40 <= player.max_len or len(history) <= len(EXAMPLES):
            return best, worst
        history.pop()


def ask(player, history, piece, last, moves) -> list[float]:
    with player._lock:
        return contrast_scores(player, *_prompts(player, history, piece, last, moves), list(moves))


def check(player) -> None:
    """Raise NotChatModel unless the player can hold an Instruct conversation (once per loaded model)."""
    if player.__dict__.get("_instruct_ok"):
        return
    if not is_chat(player):
        raise NotChatModel(f"`{player.model_id}` has no chat template: Instruct mode is for chat (instruct) models.")
    longest = [f"clears {i % 4 + 1} lines, many new holes, very high stack, bumpy" for i in range(GROUP)]
    prev = (question("T", longest, "no lines, removes holes, stack at the top, uneven"), longest[0])
    try:
        ask(player, EXAMPLES + [prev], "T", "no lines, removes holes, stack at the top, uneven", longest)
    except NotChatModel:
        raise
    except Exception as e:
        raise NotChatModel(f"`{player.model_id}`: the chat prompt failed ({type(e).__name__}: {str(e)[:200]}).")
    player._instruct_ok = True


def choose(game, player, seed: int, cands):
    """One Instruct move for a chat model: plays it on `game`."""
    cand, score, reply, n = knockout(game, seed, cands, lambda h, piece, last, moves: ask(player, h, piece, last, moves))
    game.apply(cand, score, n, reply)
    return True


def transcript() -> str:
    """The start of an Instruct conversation as plain text (How it works tab)."""
    q = question("T", ["no lines, no new holes, low stack, flat", "no lines, 1 new hole, low stack, bumpy",
                       "clears 2 lines, no new holes, low stack, flat", "no lines, no new holes, high stack, flat"],
                 last="clears 1 line, no new holes, low stack, flat")
    parts = [("system", SYSTEM)]
    for u, a in EXAMPLES:
        parts += [("user", u), ("assistant", a)]
    parts.append(("user", q))
    return "\n\n".join(f"[{role}]\n{text}" for role, text in parts)