"""Future-Entropy Sampler.

At each step, instead of sampling straight from p(w | c), fork the KV cache
once per top-k candidate, decode one token into each fork to see what
distribution q_w follows it, score candidates by how much p(w|c) they carry
*and* how much future choice they preserve, then commit to one and discard
the rest of the forks.

  s(w) = p(w|c)^a * H_hat(w)^b,   a = 1 - alpha,  b = 1 + alpha,  alpha in [-1, 1]

alpha = -1 -> a=2,b=0 (pure probability, squared)
alpha =  0 -> a=1,b=1 (balanced - matches the base s(w) = p(w)*H_hat(w) form)
alpha = +1 -> a=0,b=2 (pure future-entropy, squared)
This linear a/b mapping isn't specified by the source article beyond "derived
from alpha" - it's the simplest crossfade that recovers the balanced form at
alpha=0 and degenerates cleanly at the extremes.
"""
import math
import os

import numpy as np

import llama_capi as C

MODEL_PATH = os.environ.get("ENTROPY_SAMPLER_MODEL")


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


def normalized_entropy(probs):
    p = probs / probs.sum()
    p = p[p > 0]
    if len(p) <= 1:
        return 0.0
    h = -np.sum(p * np.log(p))
    return float(h / math.log(len(p)))


def alpha_to_exponents(alpha):
    alpha = max(-1.0, min(1.0, alpha))
    return 1.0 - alpha, 1.0 + alpha


def sine_alpha_schedule(period_tokens=12, amplitude=1.0, phase=0.0, offset=0.0):
    """offset shifts the wave's center - e.g. offset=-0.2 dips deeper into
    safe/coherent territory than it peaks into entropy-chasing territory,
    biasing toward coherence while still spiking for occasional surprise."""
    def schedule(step):
        return offset + amplitude * math.sin(2 * math.pi * step / period_tokens + phase)
    return schedule


class EntropySampler:
    def __init__(self, model_path=MODEL_PATH, n_ctx=8192, top_k=12,
                 top_n_future=20, n_threads=8, seed=0, think_ban=True,
                 n_batch=8192, n_ubatch=512):
        if not model_path:
            raise RuntimeError(
                "No model path given. Set ENTROPY_SAMPLER_MODEL to a .gguf file, "
                "or pass model_path= / --model explicitly."
            )
        self.top_k = top_k
        self.top_n_future = top_n_future
        self.rng = np.random.default_rng(seed)

        C.lib.llama_backend_init()

        mparams = C.lib.llama_model_default_params()
        mparams.n_gpu_layers = 99
        self.model = C.lib.llama_model_load_from_file(model_path.encode(), mparams)
        if not self.model:
            raise RuntimeError(f"model load failed for path: {model_path}")

        self.vocab = C.lib.llama_model_get_vocab(self.model)
        self.n_vocab = C.lib.llama_vocab_n_tokens(self.vocab)

        # This model only offers enable_thinking through its chat template's
        # assistant-turn preamble - we do raw completion, not chat-formatted
        # prompting, so that switch is never in play. Ban <think>/</think>
        # directly instead: both are single dedicated tokens in this vocab.
        # (Known limited: suppresses literal leakage, not the underlying
        # register-break the model reroutes around it with.)
        self.banned_token_ids = []
        if think_ban:
            for special in ("<think>", "</think>"):
                toks = C.tokenize(self.vocab, special, add_special=False, parse_special=True)
                if len(toks) == 1:
                    self.banned_token_ids.append(toks[0])

        cparams = C.lib.llama_context_default_params()
        cparams.n_ctx = n_ctx
        cparams.n_seq_max = top_k + 1  # seq 0 = real context, 1..top_k = forks
        cparams.n_threads = n_threads
        cparams.n_threads_batch = n_threads
        # n_batch is the hard ceiling on tokens submitted in one llama_decode()
        # call - prime() decodes the whole prompt in a single call, so this
        # must cover the largest prompt (i.e. rendered chat history) expected,
        # not just n_ubatch's internal per-step chunk size. Left at the
        # llama.cpp default (2048) this hits GGML_ASSERT(n_tokens_all <=
        # cparams.n_batch), which aborts the whole process - not a
        # recoverable Python exception.
        cparams.n_batch = n_batch
        cparams.n_ubatch = n_ubatch

        self.ctx = C.lib.llama_init_from_model(self.model, cparams)
        if not self.ctx:
            raise RuntimeError("context init failed")

        self.mem = C.lib.llama_get_memory(self.ctx)
        self.n_past = 0

    def close(self):
        C.lib.llama_free(self.ctx)
        C.lib.llama_model_free(self.model)

    def _decode(self, tokens, seq_ids, positions, want_logits):
        """seq_ids/positions/want_logits are per-token, same length as tokens."""
        n = len(tokens)
        batch = C.lib.llama_batch_init(n, 0, 1)
        batch.n_tokens = n
        for i in range(n):
            batch.token[i] = tokens[i]
            batch.pos[i] = positions[i]
            batch.n_seq_id[i] = 1
            batch.seq_id[i][0] = seq_ids[i]
            batch.logits[i] = 1 if want_logits[i] else 0
        rc = C.lib.llama_decode(self.ctx, batch)
        C.lib.llama_batch_free(batch)
        if rc != 0:
            raise RuntimeError(f"llama_decode failed rc={rc}")

    def _logits_row(self, i):
        ptr = C.lib.llama_get_logits_ith(self.ctx, i)
        return np.ctypeslib.as_array(ptr, shape=(self.n_vocab,)).copy()

    def prime(self, prompt):
        """Decode the prompt on seq 0. Returns base logits for the first step."""
        C.lib.llama_memory_seq_rm(self.mem, 0, -1, -1)  # clear any prior run's cache
        tokens = C.tokenize(self.vocab, prompt)
        positions = list(range(len(tokens)))
        want = [False] * (len(tokens) - 1) + [True]
        self._decode(tokens, [0] * len(tokens), positions, want)
        self.n_past = len(tokens)
        return self._logits_row(-1)

    def _mask_banned(self, logits):
        if self.banned_token_ids:
            logits = logits.copy()
            logits[self.banned_token_ids] = -np.inf
        return logits

    def _top_candidates(self, base_logits):
        probs = softmax(self._mask_banned(base_logits))
        top_idx = np.argsort(-probs)[: self.top_k]
        return top_idx, probs[top_idx]

    def _fork_and_score(self, top_idx):
        """Batched fork+peek+discard: one candidate token per fork sequence,
        all decoded in a single llama_decode call."""
        fork_seqs = list(range(1, self.top_k + 1))
        for fseq in fork_seqs:
            C.lib.llama_memory_seq_cp(self.mem, 0, fseq, -1, -1)

        self._decode(
            tokens=[int(t) for t in top_idx],
            seq_ids=fork_seqs,
            positions=[self.n_past] * self.top_k,
            want_logits=[True] * self.top_k,
        )

        h_hat = np.empty(self.top_k)
        for j in range(self.top_k):
            fork_probs = softmax(self._mask_banned(self._logits_row(j)))
            top_future = np.argsort(-fork_probs)[: self.top_n_future]
            h_hat[j] = normalized_entropy(fork_probs[top_future])

        for fseq in fork_seqs:
            C.lib.llama_memory_seq_rm(self.mem, fseq, -1, -1)

        return h_hat

    def _score_candidates(self, base_logits):
        top_idx, top_p = self._top_candidates(base_logits)
        h_hat = self._fork_and_score(top_idx)
        return top_idx, top_p, h_hat

    def step(self, base_logits, alpha, confidence_threshold=None):
        """confidence_threshold: if the top base candidate's p(w|c) is at or
        above this, skip the fork+peek lookahead entirely and sample straight
        from p(w|c) among the top-k - most tokens in fluent prose are forced
        (punctuation, function words) and the entropy signal adds nothing
        there. None disables skipping (always fork, the original behavior)."""
        top_idx, top_p = self._top_candidates(base_logits)
        forked = confidence_threshold is None or top_p[0] < confidence_threshold

        if forked:
            h_hat = self._fork_and_score(top_idx)
            a, b = alpha_to_exponents(alpha)
            scores = (top_p ** a) * (h_hat ** b)
            total = scores.sum()
            if total <= 0:
                scores = top_p  # degenerate fallback: every candidate closes off the future
                total = scores.sum()
        else:
            h_hat = np.full(self.top_k, np.nan)
            scores = top_p
            total = scores.sum()

        dist = scores / total
        choice = self.rng.choice(len(top_idx), p=dist)
        chosen_token = int(top_idx[choice])

        diagnostics = {
            "candidates": [
                {
                    "token": C.token_to_str(self.vocab, int(top_idx[j])),
                    "p": float(top_p[j]),
                    "h_hat": None if not forked else float(h_hat[j]),
                    "score": float(dist[j]),
                }
                for j in range(len(top_idx))
            ],
            "chosen": choice,
            "alpha": alpha,
            "forked": forked,
        }

        self._decode([chosen_token], [0], [self.n_past], [True])
        self.n_past += 1
        next_base_logits = self._logits_row(-1)

        return chosen_token, next_base_logits, diagnostics

    def generate_stream(self, prompt, max_new_tokens=60, alpha=0.0, confidence_threshold=None):
        """Yields (piece, diag) one token at a time - the shared engine behind
        generate() (batch) and the HTTP server (streaming)."""
        alpha_fn = alpha if callable(alpha) else (lambda step: alpha)

        base_logits = self.prime(prompt)
        for step_i in range(max_new_tokens):
            a = alpha_fn(step_i)
            token, base_logits, diag = self.step(base_logits, a, confidence_threshold)

            if C.lib.llama_vocab_is_eog(self.vocab, token):
                return

            piece = C.token_to_str(self.vocab, token)
            yield piece, diag

    def generate(self, prompt, max_new_tokens=60, alpha=0.0, verbose=True,
                 confidence_threshold=None):
        text = prompt
        for step_i, (piece, diag) in enumerate(
            self.generate_stream(prompt, max_new_tokens, alpha, confidence_threshold)
        ):
            text += piece

            if verbose:
                chosen = diag["candidates"][diag["chosen"]]
                h_str = f"{chosen['h_hat']:.3f}" if chosen["h_hat"] is not None else "  -- "
                print(
                    f"[{step_i:3d}] alpha={diag['alpha']:+.2f} {'fork' if diag['forked'] else 'skip'}"
                    f" -> {piece!r:12s} p={chosen['p']:.3f} H_hat={h_str} "
                    f"score={chosen['score']:.3f}"
                )

        return text


if __name__ == "__main__":
    sampler = EntropySampler(top_k=12, top_n_future=20, n_ctx=4096)
    try:
        prompt = "The old lighthouse keeper had seen many storms, but none like"

        print("=== alpha = -1.0 (pure probability) ===")
        print(sampler.generate(prompt, max_new_tokens=40, alpha=-1.0))

        print("\n=== alpha = 0.0 (balanced) ===")
        print(sampler.generate(prompt, max_new_tokens=40, alpha=0.0))

        print("\n=== alpha = +1.0 (pure future-entropy) ===")
        print(sampler.generate(prompt, max_new_tokens=40, alpha=1.0))

        print("\n=== alpha-wave (sine oscillation, period=10 tokens) ===")
        schedule = sine_alpha_schedule(period_tokens=10, amplitude=1.0)
        print(sampler.generate(prompt, max_new_tokens=40, alpha=schedule))
    finally:
        sampler.close()