"""EPT — Entropy-Patch Tokenization.

Segments a byte stream into variable-length patches such that each patch
carries approximately ``budget_bits`` of accumulated surprisal under an
order-1 byte-level Markov model ("constant-information quanta").

Consequences:
  * low-entropy regions (repeated keywords, indentation, boilerplate) produce
    LONG patches, high-entropy regions produce SHORT ones — no merge table,
    no vocabulary to train;
  * patches are content-addressed with FNV-1a-64: the same byte substring is
    always the same token, wherever it occurs, with zero learning;
  * strictly streaming: one pass, O(n) with tiny state.
"""
from __future__ import annotations

import numpy as np

from .hashing import fnv1a64


class EPT:
    def __init__(
        self,
        budget_bits: float = 16.0,
        min_len: int = 2,
        max_len: int = 32,
        alpha: float = 0.5,
    ) -> None:
        self.budget_bits = float(budget_bits)
        self.min_len = int(min_len)
        self.max_len = int(max_len)
        self.alpha = float(alpha)
        self.log2p: np.ndarray | None = None  # (256, 256) float32

    # ------------------------------------------------------------------ fit
    def fit(self, data: bytes) -> "EPT":
        """Estimate the order-1 byte Markov model on training bytes."""
        arr = np.frombuffer(data, dtype=np.uint8).astype(np.int64)
        C = np.zeros((256, 256), dtype=np.int64)
        np.add.at(C, (arr[:-1], arr[1:]), 1)
        P = (C + self.alpha).astype(np.float64)
        P /= P.sum(axis=1, keepdims=True)
        self.log2p = np.log2(P).astype(np.float32)
        return self

    # ------------------------------------------------------------ surprise
    def surprise(self, data: bytes) -> np.ndarray:
        """Per-byte surprisal -log2 P(byte | previous byte), in bits."""
        if self.log2p is None:
            raise RuntimeError("EPT must be fit() before use")
        arr = np.frombuffer(data, dtype=np.uint8).astype(np.int64)
        s = np.empty(len(arr), dtype=np.float32)
        if len(arr) == 0:
            return s
        s[0] = 8.0  # uniform prior for the first byte
        if len(arr) > 1:
            s[1:] = -self.log2p[arr[:-1], arr[1:]]
        return s

    # ----------------------------------------------------------- segmenting
    def boundaries(self, data: bytes) -> np.ndarray:
        """Greedy constant-surprise segmentation.

        Returns inclusive boundary positions ``b_0=0 < b_1 < ... < b_k = n``;
        patch i is ``data[b_i : b_{i+1}]``. A patch closes when its
        accumulated surprisal reaches ``budget_bits`` (or ``max_len`` /
        end-of-stream forces it).
        """
        if len(data) == 0:
            return np.zeros(0, dtype=np.int64)
        s = self.surprise(data)
        n = len(s)
        cum = np.concatenate([[0.0], np.cumsum(s, dtype=np.float64)])
        bounds = [0]
        start = 0
        searchsorted = np.searchsorted
        min_len, max_len, budget = self.min_len, self.max_len, self.budget_bits
        while start < n:
            target = cum[start] + budget
            b = int(searchsorted(cum, target, side="left"))
            if b < start + min_len:
                b = start + min_len
            elif b > start + max_len:
                b = start + max_len
            if b > n:
                b = n
            bounds.append(b)
            start = b
        # drop a degenerate trailing boundary at exactly n produced twice
        if len(bounds) > 1 and bounds[-1] == bounds[-2]:
            bounds.pop()
        return np.asarray(bounds, dtype=np.int64)

    def segment(self, data: bytes) -> list[bytes]:
        b = self.boundaries(data)
        return [data[b[i] : b[i + 1]] for i in range(len(b) - 1)]

    def segment_ids(self, data: bytes) -> list[int]:
        return [fnv1a64(p) for p in self.segment(data)]

    # ---------------------------------------------------------- calibration
    def calibrate(self, sample: bytes, target_len: float = 4.5, iters: int = 10) -> float:
        """Binary-search the surprisal budget to hit a target average patch
        length on a calibration sample. Returns the chosen budget."""
        lo, hi = 1.0, 64.0
        for _ in range(iters):
            mid = 0.5 * (lo + hi)
            self.budget_bits = mid
            b = self.boundaries(sample)
            avg = len(sample) / max(len(b) - 1, 1)
            if avg < target_len:
                lo = mid
            else:
                hi = mid
        self.budget_bits = 0.5 * (lo + hi)
        return self.budget_bits

    # ------------------------------------------------------------- metrics
    def stats(self, patches: list[bytes]) -> dict:
        lens = np.array([len(p) for p in patches], dtype=np.float64)
        return {
            "n_patches": int(len(patches)),
            "avg_len": float(lens.mean()) if len(lens) else 0.0,
            "min_len": int(lens.min()) if len(lens) else 0,
            "max_len": int(lens.max()) if len(lens) else 0,
            "p50_len": float(np.percentile(lens, 50)) if len(lens) else 0.0,
            "p90_len": float(np.percentile(lens, 90)) if len(lens) else 0.0,
            "bytes": int(lens.sum()),
        }

    # ---------------------------------------------------------- (de)hydrate
    def state_dict(self) -> dict:
        return {
            "budget_bits": self.budget_bits,
            "min_len": self.min_len,
            "max_len": self.max_len,
            "alpha": self.alpha,
            "log2p": self.log2p.astype(np.float32),
        }

    def load_state_dict(self, sd: dict) -> None:
        self.budget_bits = float(sd["budget_bits"])
        self.min_len = int(sd["min_len"])
        self.max_len = int(sd["max_len"])
        self.alpha = float(sd["alpha"])
        self.log2p = np.asarray(sd["log2p"], dtype=np.float32)