"""llm_toolkit.py - Week 2 lab: LLMs for builders. Standard library only.

Four small tools, each with a mock model so no API key is needed:
  1. Decoding: greedy, temperature, top-k, top-p over a toy distribution
  2. Attention: scaled dot-product, causal, and linear attention (as an RNN)
  3. Constrained decoding: token masking, plus a validate-and-retry loop
  4. Context budget: count, clear old tool output, compact, reorder

Run `python llm_toolkit.py` for a demo. Run `python test_llm_toolkit.py` for tests.
"""
import json
import math
import random

# ---------------------------------------------------------------- 1. decoding

def softmax(logits, temperature=1.0):
    """Turn scores into probabilities. Lower temperature = sharper."""
    if temperature <= 0:
        raise ValueError("temperature must be > 0 (use greedy for 0)")
    scaled = [x / temperature for x in logits]
    m = max(scaled)
    exps = [math.exp(x - m) for x in scaled]
    total = sum(exps)
    return [e / total for e in exps]


def entropy(probs):
    """Average surprise in bits. 0 means the model is certain."""
    return -sum(p * math.log2(p) for p in probs if p > 0)


def greedy(probs):
    """Index of the most likely token. Ties go to the lowest index."""
    return max(range(len(probs)), key=lambda i: (probs[i], -i))


def top_k_filter(probs, k):
    """Keep the k most likely tokens, zero the rest, renormalize."""
    if k < 1:
        raise ValueError("k must be >= 1")
    keep = set(sorted(range(len(probs)), key=lambda i: (-probs[i], i))[:k])
    return _renorm([p if i in keep else 0.0 for i, p in enumerate(probs)])


def top_p_filter(probs, p):
    """Nucleus: keep the smallest set of top tokens whose mass reaches p."""
    if not 0 < p <= 1:
        raise ValueError("p must be in (0, 1]")
    order = sorted(range(len(probs)), key=lambda i: (-probs[i], i))
    keep, mass = set(), 0.0
    for i in order:
        keep.add(i)
        mass += probs[i]
        if mass >= p - 1e-12:
            break
    return _renorm([q if i in keep else 0.0 for i, q in enumerate(probs)])


def _renorm(probs):
    total = sum(probs)
    return [q / total for q in probs]


def sample_index(probs, rng):
    """Draw one index from a probability list."""
    r, acc = rng.random(), 0.0
    for i, q in enumerate(probs):
        acc += q
        if r < acc:
            return i
    return max(i for i, q in enumerate(probs) if q > 0)


def decode_step(logits, rng, temperature=1.0, top_k=None, top_p=None):
    """One decoding step. temperature=0 means greedy."""
    if temperature == 0:
        return greedy(logits)
    probs = softmax(logits, temperature)
    if top_k is not None:
        probs = top_k_filter(probs, top_k)
    if top_p is not None:
        probs = top_p_filter(probs, top_p)
    return sample_index(probs, rng)

# ---------------------------------------------------------------- 2. attention

def matmul(a, b):
    bt = list(zip(*b))
    return [[sum(x * y for x, y in zip(row, col)) for col in bt] for row in a]


def transpose(a):
    return [list(r) for r in zip(*a)]


def attention(Q, K, V, causal=True):
    """Scaled dot-product attention. Returns (output, weights).
    Q, K: n x d lists. V: n x dv. Row i of weights sums to 1."""
    d = len(Q[0])
    scores = matmul(Q, transpose(K))
    n = len(scores)
    weights = []
    for i in range(n):
        row = [s / math.sqrt(d) for s in scores[i]]
        if causal:
            row = [x if j <= i else float("-inf") for j, x in enumerate(row)]
        m = max(row)
        e = [0.0 if x == float("-inf") else math.exp(x - m) for x in row]
        t = sum(e)
        weights.append([x / t for x in e])
    return matmul(weights, V), weights


def feature_map(x):
    """elu(x) + 1, the positive feature map used for linear attention."""
    return [v + 1.0 if v > 0 else math.exp(v) for v in x]


def linear_attention_parallel(Q, K, V):
    """Causal linear attention written like softmax attention: O(n^2)."""
    n = len(Q)
    out = []
    for i in range(n):
        qf = feature_map(Q[i])
        ws = [sum(a * b for a, b in zip(qf, feature_map(K[j]))) for j in range(i + 1)]
        z = sum(ws)
        out.append([sum(ws[j] * V[j][c] for j in range(i + 1)) / z for c in range(len(V[0]))])
    return out


def linear_attention_recurrent(Q, K, V):
    """The same result as an RNN. Keeps a fixed-size state: O(n) time, O(1) memory."""
    d, dv = len(Q[0]), len(V[0])
    S = [[0.0] * dv for _ in range(d)]   # sum of phi(k) v^T
    z = [0.0] * d                        # sum of phi(k)
    out = []
    for q, k, v in zip(Q, K, V):
        kf, qf = feature_map(k), feature_map(q)
        for a in range(d):
            z[a] += kf[a]
            for c in range(dv):
                S[a][c] += kf[a] * v[c]
        den = sum(qf[a] * z[a] for a in range(d))
        out.append([sum(qf[a] * S[a][c] for a in range(d)) / den for c in range(dv)])
    return out


def kv_cache_bytes(layers, kv_heads, head_dim, seq_len, bytes_per_value=2, batch=1):
    """Size of the key-value cache: 2 (K and V) x layers x heads x dim x tokens."""
    return 2 * layers * kv_heads * head_dim * seq_len * bytes_per_value * batch

# ------------------------------------------------- 3. constrained decoding

class MockJSONModel:
    """Pretends to be an LLM that is bad at JSON at first.
    It replays `replies` in order. Once it runs out, it repeats the last one."""
    def __init__(self, replies):
        self.replies, self.calls, self.prompts = list(replies), 0, []

    def complete(self, prompt):
        self.prompts.append(prompt)
        r = self.replies[min(self.calls, len(self.replies) - 1)]
        self.calls += 1
        return r


_TYPES = {"string": str, "integer": int, "number": (int, float), "boolean": bool,
          "array": list, "object": dict, "null": type(None)}


def validate(value, schema, path="$"):
    """Tiny JSON-Schema subset: type, enum, required, properties, items.
    Returns a list of error strings. Empty list means valid."""
    errs = []
    t = schema.get("type")
    if t:
        ok = isinstance(value, _TYPES[t]) and not (t in ("integer", "number") and isinstance(value, bool))
        if t == "integer" and isinstance(value, float):
            ok = False
        if not ok:
            return [f"{path}: expected {t}, got {type(value).__name__}"]
    if "enum" in schema and value not in schema["enum"]:
        errs.append(f"{path}: {value!r} is not one of {schema['enum']}")
    if isinstance(value, dict):
        for key in schema.get("required", []):
            if key not in value:
                errs.append(f"{path}: missing required key '{key}'")
        for key, sub in schema.get("properties", {}).items():
            if key in value:
                errs += validate(value[key], sub, f"{path}.{key}")
    if isinstance(value, list) and "items" in schema:
        for i, item in enumerate(value):
            errs += validate(item, schema["items"], f"{path}[{i}]")
    return errs


def extract_json(text):
    """Pull the first {...} block out of a reply. Models often add chatter or code fences."""
    start = text.find("{")
    if start < 0:
        raise ValueError("no JSON object found")
    depth, in_str, esc = 0, False, False
    for i in range(start, len(text)):
        c = text[i]
        if in_str:
            esc = (c == "\\") and not esc
            if c == '"' and not esc:
                in_str = False
        elif c == '"':
            in_str = True
        elif c == "{":
            depth += 1
        elif c == "}":
            depth -= 1
            if depth == 0:
                return json.loads(text[start:i + 1])
    raise ValueError("unterminated JSON object")


def generate_json(model, task, schema, max_retries=3):
    """Ask, parse, validate. On failure, send the errors back and retry.
    Returns (value, attempts). Raises RuntimeError when retries run out."""
    prompt = f"{task}\nReply with JSON only. Schema: {json.dumps(schema)}"
    for attempt in range(1, max_retries + 2):
        reply = model.complete(prompt)
        try:
            value = extract_json(reply)
            errs = validate(value, schema)
        except ValueError as e:      # json.JSONDecodeError is a ValueError
            errs = [f"not valid JSON: {e}"]
        if not errs:
            return value, attempt
        prompt = (f"{task}\nReply with JSON only. Schema: {json.dumps(schema)}\n"
                  f"Your last reply was rejected:\n- " + "\n- ".join(errs) + "\nTry again.")
    raise RuntimeError(f"no valid JSON after {max_retries + 1} attempts; last errors: {errs}")


class TrieConstraint:
    """Allow only token sequences that spell one of `allowed` strings.
    This is the idea behind constrained decoding: at each step, only legal
    next tokens keep their score. Real systems build such masks from a regex
    or grammar over the model's real vocabulary [Willard and Louf 2023]."""
    def __init__(self, allowed, tokenize):
        self.paths = {tuple(tokenize(s)) for s in allowed}

    def legal_tokens(self, prefix):
        n = len(prefix)
        return {p[n] for p in self.paths if len(p) > n and p[:n] == tuple(prefix)}

    def is_complete(self, prefix):
        return tuple(prefix) in self.paths


def constrained_decode(score_fn, vocab, constraint, rng, temperature=1.0, max_len=20):
    """Sample tokens, but set every illegal token's score to -inf first.
    score_fn(prefix) -> list of logits aligned with vocab (a mock 'model')."""
    prefix = []
    while not constraint.is_complete(prefix) and len(prefix) < max_len:
        legal = constraint.legal_tokens(prefix)
        if not legal:
            raise RuntimeError("dead end: no legal token")
        logits = score_fn(prefix)
        masked = [x if vocab[i] in legal else float("-inf") for i, x in enumerate(logits)]
        probs = softmax(masked, temperature)
        prefix.append(vocab[sample_index(probs, rng)])
    return "".join(prefix)

# ------------------------------------------------------ 4. context budget

def count_tokens(text):
    """Toy token counter: 1 word = 1 token. Real tokenizers differ. Swap yours in."""
    return len(text.split())


def total_tokens(messages, counter=count_tokens):
    return sum(counter(m["text"]) for m in messages)


def clear_old_tool_results(messages, keep_last=2):
    """Lightest compaction: replace old tool outputs with a short stub."""
    tool_idx = [i for i, m in enumerate(messages) if m["role"] == "tool"]
    old = set(tool_idx[:-keep_last]) if keep_last else set(tool_idx)
    return [dict(m, text="[tool result cleared]") if i in old and not m.get("pinned") else m
            for i, m in enumerate(messages)]


def mock_summarize(messages):
    """Stand-in for an LLM summary: first six words of each message."""
    return "Summary: " + " | ".join(" ".join(m["text"].split()[:6]) for m in messages)


def compact(messages, limit, reserve_output=0, summarize=mock_summarize,
            keep_recent=3, counter=count_tokens):
    """Fit messages into limit - reserve_output tokens.
    Step 1: clear old tool results. Step 2: summarize the oldest unpinned
    messages into one. Pinned messages and the last `keep_recent` stay.
    Returns (messages, report)."""
    room = limit - reserve_output
    report = {"before": total_tokens(messages, counter), "steps": []}
    cur = messages
    if total_tokens(cur, counter) > room:
        cur = clear_old_tool_results(cur)
        report["steps"].append("cleared old tool results")
    if total_tokens(cur, counter) > room:
        tail = cur[-keep_recent:] if keep_recent else []
        head = cur[:len(cur) - len(tail)]
        pinned = [m for m in head if m.get("pinned")]
        loose = [m for m in head if not m.get("pinned")]
        if loose:
            cur = pinned + [{"role": "system", "text": summarize(loose)}] + tail
            report["steps"].append(f"summarized {len(loose)} messages")
    report["after"] = total_tokens(cur, counter)
    report["fits"] = report["after"] <= room
    return cur, report


def edges_first(items_by_relevance):
    """Reorder so the most relevant items sit at the start and the end,
    and the least relevant sit in the middle (see 'Lost in the Middle')."""
    front, back = [], []
    for i, item in enumerate(items_by_relevance):
        (front if i % 2 == 0 else back).append(item)
    return front + back[::-1]

# ---------------------------------------------------------------- 5. extras

def coverage(p_single, n):
    """Chance that at least one of n independent samples is right."""
    return 1 - (1 - p_single) ** n


def majority_vote(answers):
    """Most common answer. Ties go to the one seen first."""
    counts = {}
    for a in answers:
        counts[a] = counts.get(a, 0) + 1
    return max(counts, key=lambda a: counts[a])

# ---------------------------------------------------------------- demo

def _demo():
    rng = random.Random(7)
    words = ["to", "as", "the", "a", "two", "how"]
    logits = [2.0, 0.5, 0.3, 0.2, -0.5, -1.0]
    print("== 1. decoding (toy distribution) ==")
    for T in (0.5, 1.0, 2.0):
        p = softmax(logits, T)
        print(f"T={T}: " + " ".join(f"{w}={q:.2f}" for w, q in zip(words, p)) + f"  entropy={entropy(p):.2f} bits")
    p1 = softmax(logits)
    print("top_k=2 :", [round(x, 2) for x in top_k_filter(p1, 2)])
    print("top_p=.8:", [round(x, 2) for x in top_p_filter(p1, 0.8)])
    print("greedy  :", words[greedy(p1)])
    draws = [words[decode_step(logits, rng, 1.0, top_p=0.9)] for _ in range(8)]
    print("8 draws, T=1, top_p=0.9:", draws)

    print("\n== 2. attention ==")
    Q = [[1.0, 0.0], [0.0, 1.0], [2.0, 0.0]]
    K = [[1.0, 0.0], [0.0, 1.0], [1.0, 0.0]]
    V = [[1.0, 0.0], [0.0, 1.0], [0.5, 0.5]]
    out, w = attention(Q, K, V)
    print("causal weights row 3:", [round(x, 3) for x in w[2]])
    a = linear_attention_parallel(Q, K, V)
    b = linear_attention_recurrent(Q, K, V)
    print("linear: parallel == recurrent:", all(abs(x - y) < 1e-9 for r, s in zip(a, b) for x, y in zip(r, s)))
    kv = kv_cache_bytes(32, 8, 128, 32768)
    print(f"KV cache for a made-up 32-layer model at 32k tokens: {kv / 2**30:.1f} GiB")

    print("\n== 3. constrained decoding ==")
    schema = {"type": "object", "required": ["name", "intent"], "properties": {
        "name": {"type": "string"}, "intent": {"type": "string", "enum": ["meeting", "intro", "follow-up"]}}}
    model = MockJSONModel(['Sure! {"name": "Sarah", "intent": "chat"}',
                           '{"name": "Sarah"}',
                           '```json\n{"name": "Sarah", "intent": "meeting"}\n```'])
    value, attempts = generate_json(model, "Extract contact info: I'm Sarah. Meet Thursday?", schema)
    print(f"valid after {attempts} attempts:", value)
    vocab = ["meet", "ing", "in", "tro", "follow", "-up", "chat", "xyz"]
    tok = lambda s: {"meeting": ["meet", "ing"], "intro": ["in", "tro"], "follow-up": ["follow", "-up"]}[s]
    cons = TrieConstraint(["meeting", "intro", "follow-up"], tok)
    bad_model = lambda prefix: [0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 5.0, 5.0]  # loves illegal tokens
    outs = {constrained_decode(bad_model, vocab, cons, rng) for _ in range(20)}
    print("20 constrained samples, all legal:", sorted(outs))

    print("\n== 4. context budget ==")
    msgs = [{"role": "system", "text": "You are a careful coding agent.", "pinned": True},
            {"role": "user", "text": "Fix the failing test in the billing module please"},
            {"role": "tool", "text": "log line " * 40},
            {"role": "assistant", "text": "I will open the test file and read it now"},
            {"role": "tool", "text": "file contents " * 30},
            {"role": "assistant", "text": "The bug is a wrong rounding mode"},
            {"role": "tool", "text": "diff output " * 10},
            {"role": "user", "text": "Great. Now add a regression test."}]
    new, rep = compact(msgs, limit=120, reserve_output=20)
    print(rep)
    print("order for 5 docs ranked 1..5:", edges_first([1, 2, 3, 4, 5]))


if __name__ == "__main__":
    _demo()