"""MiniCPM adapter for the pinned Apache-2.0 Kev pointer implementation.

No Qwen token IDs are reused. Structural markers are five existing MiniCPM
unused tokens. All user angle brackets are escaped before tokenization.
Gold labels enter only the returned training targets, never the input payload.
"""
from copy import deepcopy

from decision_data import api_payload, gold_key, question_keys
from kev.api import SystemOneRequest, to_record
from kev.model import SPECIAL, DecisionModel, encode

MARKERS = ['<unused_token_' + str(i) + '>' for i in range(5)]
# Five existing, distinct, single-token specials per vocabulary family, tried in this order. MiniCPM (2B and
# 1B share one tokenizer) has unused tokens; Qwen vocabularies carry Kev's own five specials, which the
# pinned Kev recipe was written for.
MARKER_SETS = {'minicpm': MARKERS, 'kev': list(SPECIAL)}


def choose_markers(tokenizer):
    """(family, markers, ids): the first family whose five markers exist as distinct single tokens."""
    for family, markers in MARKER_SETS.items():
        ids = [tokenizer.convert_tokens_to_ids(t) for t in markers]
        if len(set(ids)) != 5 or any(i is None or i == tokenizer.unk_token_id for i in ids):
            continue
        if all(tokenizer(t, add_special_tokens=False).input_ids == [i] for t, i in zip(markers, ids)):
            return family, list(markers), ids
    raise ValueError('No vocabulary family with five distinct existing single-token structural markers')


class StructuralTokenizer:
    def __init__(self, tokenizer):
        self.tokenizer = tokenizer
        self.family, self.markers, self.marker_ids = choose_markers(tokenizer)
        self.mapping = dict(zip(SPECIAL, self.markers))
        self.pad_token_id = tokenizer.pad_token_id

    def convert_tokens_to_ids(self, text):
        return self.tokenizer.convert_tokens_to_ids(self.mapping.get(text, text))

    def __call__(self, text, **kwargs):
        # Block the structural markers of every family and other existing controls in inputs.
        return self.tokenizer(text.replace('<', '＜').replace('>', '＞'), **kwargs)


def encode_record(tokenizer, record, max_tokens=2048, option_isolation=False):
    wire = api_payload(record)
    rec, _ = to_record(SystemOneRequest.model_validate(wire))
    result = encode(tokenizer, rec, strict=True, max_state=max_tokens,
                    max_branch=max_tokens, option_isolation=option_isolation)
    if len(result['ids']) > max_tokens:
        raise ValueError('Packed input exceeds the declared context cap')
    result['labels'] = [question_keys(q).index(gold_key(q)) for q in record['questions'].values()]
    result['question_ids'] = list(record['questions'])
    result['soft'] = [soft_target(q) for q in record['questions'].values()]
    result['ordinal'] = [q['type'] == 'score' for q in record['questions'].values()]
    # Option keys travel with the encoding so an anchor distribution can be matched by key, not by position.
    result['option_keys'] = [question_keys(q) for q in record['questions'].values()]
    return result


def ranked_probability_score(probabilities, gold):
    """Mean squared error between the predicted and the observed cumulative distribution.

    This is the objective the pinned Apache-2.0 Kev recipe records as `ranked_probability_score` for Score
    questions. Unlike cross entropy it charges for how far the predicted mass sits from the true level, which
    is what our error map needs: on the five-level replay questions 115 of our 142 errors are off by exactly
    one level, so the ordering is nearly right and only the distance is unpriced.
    """
    if len(probabilities) < 2 or not 0 <= gold < len(probabilities):
        raise ValueError('Invalid ordered scale')
    total, running = 0., 0.
    for i in range(len(probabilities) - 1):
        running += probabilities[i]
        total += (running - (1. if i >= gold else 0.)) ** 2
    return total / (len(probabilities) - 1)


def spherical_score(probabilities, target):
    """Spherical scoring rule: the target-weighted probability divided by the vector's own length.

    Taken from the open Laya implementation of this contract, whose reward is
    `log score + 0.75 * spherical - 1.0 * ranked probability score`. It is strictly proper like log loss but
    bounded, so it keeps rewarding a confident correct answer where log loss has already saturated, and it
    punishes spreading mass thinly across options. Returned as a score to maximise; the trainer subtracts it.
    """
    if len(probabilities) != len(target) or not probabilities:
        raise ValueError('Distribution and target must share a length')
    norm = sum(p * p for p in probabilities) ** .5
    if norm <= 0:
        raise ValueError('Degenerate distribution')
    return sum(t * p for t, p in zip(target, probabilities)) / norm


def ordinal_target(levels, gold, smoothing):
    """Neighbour-smoothed target for an ordered scale.

    A plain one-hot target treats every wrong level as equally wrong, which is false for Score:
    predicting 4 when the answer is 5 is not the same mistake as predicting 1. The smoothing mass is
    split between the immediate neighbours, so the ordering is expressed in the target itself and the
    loss stays an ordinary log loss.
    """
    if not 0 <= smoothing < 1:
        raise ValueError('Ordinal smoothing must be in [0, 1)')
    if levels < 2 or not 0 <= gold < levels:
        raise ValueError('Invalid ordered scale')
    if smoothing == 0:
        return None
    target = [0.] * levels
    target[gold] = 1 - smoothing
    neighbours = [i for i in (gold - 1, gold + 1) if 0 <= i < levels]
    for i in neighbours:
        target[i] = smoothing / len(neighbours)
    return target


def soft_target(question):
    """Exact target distribution for a question that carries one, aligned with the option order.

    Returns None when the question has no distribution, or when augmentation changed the option set,
    so such a question keeps the ordinary hard cross entropy.
    """
    given = question.get('probabilities')
    if not given:
        return None
    keys = question_keys(question)
    if set(keys) != set(given):
        return None
    vector = [float(given[k]) for k in keys]
    total = sum(vector)
    if abs(total - 1) > 1e-6 or any(v < 0 for v in vector):
        raise ValueError('Question probabilities are not a distribution')
    return vector


def permute_choices(record, rng):
    result = deepcopy(record)
    for question in result['questions'].values():
        # Score levels and boolean order have fixed semantics.
        if question['type'] == 'choice':
            items = list(question['criteria'].items())
            rng.shuffle(items)
            question['criteria'] = dict(items)
    return result


def permuted_pair(record, rng):
    """A copy with every choice question's options reordered, plus the mapping back to the original order.

    perms[q][j] is the original index of the option that now sits at position j, or None when the question
    was left alone. Labels are option keys, so they stay correct without translation; only the logits need
    realigning before the two predictions can be compared.
    """
    result = deepcopy(record)
    perms = []
    for question in result['questions'].values():
        if question['type'] == 'choice' and len(question['criteria']) >= 3:
            items = list(question['criteria'].items())
            order = list(range(len(items)))
            rng.shuffle(order)
            question['criteria'] = dict(items[i] for i in order)
            perms.append(order)
        else:
            perms.append(None)
    return result, perms


def load_pointer(path, tokenizer, precision='bf16', option_isolation=False, lora_targets='all', lora_rank=16,
                 revision=None, train_markers=True):
    """The pointer model. `train_markers=False` keeps the five structural-token embedding rows frozen (Kev's
    default); every checkpoint's recipe records the choice as `train_structural_embeddings`, and the loaders
    pass it back here so the adapter keys match what was saved."""
    import torch
    if precision not in ['bf16', 'fp32']:
        raise ValueError('Unsupported precision')
    model = DecisionModel(path, tokenizer, 'cuda', lora=lora_rank, head_dim=256, revision=revision,
                          dtype=torch.bfloat16 if precision == 'bf16' else torch.float32,
                          attn='sdpa', special_embeddings=train_markers, option_isolation=option_isolation,
                          lora_targets=lora_targets)
    if not train_markers:
        model.lm.config.use_cache = False
        model.lm.gradient_checkpointing_enable(gradient_checkpointing_kwargs={'use_reentrant': False})
        return model
    # Unused pretrained rows may be near zero. Repeated RMSNorm through 42
    # layers then gives enormous embedding gradients. Start the five trainable
    # rows from meaningful existing word embeddings; PEFT saves these rows.
    words = ['state', 'question', 'option', 'end', 'decision']
    with torch.no_grad():
        embedding = model.lm.get_input_embeddings()
        rows = []
        for word in words:
            ids = tokenizer.tokenizer(word, add_special_tokens=False).input_ids
            rows.append(embedding(torch.tensor(ids, device='cuda')).mean(0))
        found = [(n, p) for n, p in model.named_parameters() if n.endswith('trainable_tokens_delta.default')]
        if len(found) != 1 or found[0][1].shape != (5, model.lm.config.hidden_size):
            raise ValueError('Unexpected PEFT structural-embedding state')
        found[0][1].copy_(torch.stack(rows).to(found[0][1].dtype))
    model.lm.config.use_cache = False
    model.lm.gradient_checkpointing_enable(gradient_checkpointing_kwargs={'use_reentrant': False})
    return model