"""Smoke test for the primitive the entropy sampler depends on:
fork the KV cache at the current position, decode one candidate token
into the fork, read back its resulting next-token distribution, then
discard the fork. If this works, the scoring loop (s(w) = p(w|c)^a * H(w)^b)
is just bookkeeping on top of this.
"""
import math
import sys

import numpy as np

import llama_capi as C
from entropy_sampler import MODEL_PATH

PROMPT = "The old lighthouse keeper had seen many storms, but none like"
TOP_K_BASE = 8       # candidates to show from the base distribution
TOP_N_FUTURE = 20    # tokens used to compute the future entropy of a fork


def softmax(x):
    x = x - np.max(x)
    e = np.exp(x)
    return e / e.sum()


def normalized_entropy(probs_top_n):
    p = probs_top_n / probs_top_n.sum()
    p = p[p > 0]
    h = -np.sum(p * np.log(p))
    return h / math.log(len(probs_top_n))


def decode_single_seq_batch(ctx, tokens, start_pos, seq_id, want_logits_last_only=True):
    n = len(tokens)
    batch = C.lib.llama_batch_init(n, 0, 1)
    batch.n_tokens = n
    for i, tok in enumerate(tokens):
        batch.token[i] = tok
        batch.pos[i] = start_pos + i
        batch.n_seq_id[i] = 1
        batch.seq_id[i][0] = seq_id
        batch.logits[i] = 0
    if want_logits_last_only:
        batch.logits[n - 1] = 1
    else:
        for i in range(n):
            batch.logits[i] = 1
    rc = C.lib.llama_decode(ctx, batch)
    C.lib.llama_batch_free(batch)
    if rc != 0:
        raise RuntimeError(f"llama_decode failed rc={rc}")


def main():
    if not MODEL_PATH:
        sys.exit("Set ENTROPY_SAMPLER_MODEL to a .gguf file path first.")
    print(f"loading model: {MODEL_PATH}")
    C.lib.llama_backend_init()

    mparams = C.lib.llama_model_default_params()
    mparams.n_gpu_layers = 99

    model = C.lib.llama_model_load_from_file(MODEL_PATH.encode(), mparams)
    if not model:
        sys.exit("model load failed")

    vocab = C.lib.llama_model_get_vocab(model)
    n_vocab = C.lib.llama_vocab_n_tokens(vocab)
    print(f"n_vocab = {n_vocab}")

    cparams = C.lib.llama_context_default_params()
    cparams.n_ctx = 4096
    cparams.n_seq_max = 4  # seq 0 = real context, seq 1..3 available for forks
    cparams.n_threads = 8
    cparams.n_threads_batch = 8

    ctx = C.lib.llama_init_from_model(model, cparams)
    if not ctx:
        sys.exit("context init failed")

    mem = C.lib.llama_get_memory(ctx)

    prompt_tokens = C.tokenize(vocab, PROMPT)
    print(f"prompt: {PROMPT!r} -> {len(prompt_tokens)} tokens")

    # --- base pass: decode the prompt on seq 0, get p(w | c) ---
    decode_single_seq_batch(ctx, prompt_tokens, start_pos=0, seq_id=0)
    n_ctx_pos = len(prompt_tokens)  # next free position on seq 0

    logits_ptr = C.lib.llama_get_logits_ith(ctx, -1)
    base_logits = np.ctypeslib.as_array(logits_ptr, shape=(n_vocab,)).copy()
    base_probs = softmax(base_logits)

    top_idx = np.argsort(-base_probs)[:TOP_K_BASE]
    print("\ntop base candidates p(w | c):")
    for i in top_idx:
        print(f"  {base_probs[i]:.4f}  {C.token_to_str(vocab, int(i))!r}")

    # --- fork + peek + discard for the single most likely candidate ---
    w = int(top_idx[0])
    w_str = C.token_to_str(vocab, w)
    fork_seq = 1

    C.lib.llama_memory_seq_cp(mem, 0, fork_seq, -1, -1)  # fork seq 0 -> seq 1
    decode_single_seq_batch(ctx, [w], start_pos=n_ctx_pos, seq_id=fork_seq)

    fork_logits_ptr = C.lib.llama_get_logits_ith(ctx, -1)
    fork_logits = np.ctypeslib.as_array(fork_logits_ptr, shape=(n_vocab,)).copy()
    fork_probs = softmax(fork_logits)

    top_future_idx = np.argsort(-fork_probs)[:TOP_N_FUTURE]
    h_hat = normalized_entropy(fork_probs[top_future_idx])

    C.lib.llama_memory_seq_rm(mem, fork_seq, -1, -1)  # discard the fork

    print(f"\nforked on candidate {w_str!r} (p={base_probs[w]:.4f})")
    print(f"normalized future entropy H_hat(w) over top-{TOP_N_FUTURE}: {h_hat:.4f}")
    print("top continuations after that candidate:")
    for i in top_future_idx[:8]:
        print(f"  {fork_probs[i]:.4f}  {C.token_to_str(vocab, int(i))!r}")

    # sanity check: seq 0 should be untouched by the fork - decoding the
    # same candidate directly on seq 0 must reproduce identical logits
    decode_single_seq_batch(ctx, [w], start_pos=n_ctx_pos, seq_id=0)
    check_ptr = C.lib.llama_get_logits_ith(ctx, -1)
    check_logits = np.ctypeslib.as_array(check_ptr, shape=(n_vocab,)).copy()
    max_diff = float(np.max(np.abs(check_logits - fork_logits)))
    print(f"\nseq0-vs-fork logit max abs diff (expect ~0): {max_diff:.6f}")

    C.lib.llama_free(ctx)
    C.lib.llama_model_free(model)
    print("\nOK: fork -> peek -> discard round-trip works.")


if __name__ == "__main__":
    main()