#!/usr/bin/env python3
"""Collect raw ring-oscillator captures from the STRIX capture sketch and estimate jitter.

For each source and decimation k (keep every k-th sample at 120 MHz) the script stores the
capture, estimates the per-bit entropy from conditional block entropies, and converts it to a
predictability P and relative predictability 2P-1. Accumulated phase jitter grows with sqrt(k),
so fitting a jittered square-wave sampling model to P(k) gives the effective jitter per sample
interval, in ring cycles. The estimate is approximate: the entropy is measured from past
samples, while the model assumes the ideal phase is known. On synthetic captures with known
random-walk jitter it overestimates by a factor of about 1.2-1.3, so divide by 1.3 when sizing
the XOR combining conservatively.
"""

import argparse
import math
import sys
from pathlib import Path

import numpy as np
import serial

SOURCES = {0: "tap A", 1: "enhanced A", 2: "enhanced B", 3: "pair XOR"}
DEFAULT_DECIMATIONS = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096]


# Board communication

def open_board(port: str) -> serial.Serial:
    board = serial.Serial(port, 115200, timeout=200)
    board.reset_input_buffer()
    board.write(b"id\n")
    while True:
        line = board.readline().decode(errors="replace").strip()
        if not line:
            sys.exit("no answer from the capture sketch")
        if line.startswith("STRX"):
            print(f"board: {line}")
            return board
        if line.startswith("ERR"):
            sys.exit(f"board: {line}")


def capture(board: serial.Serial, source: int, decimation: int) -> bytes:
    board.write(f"cap {source} {decimation}\n".encode())
    header = board.readline().decode(errors="replace").strip()
    if not header.startswith("DATA "):
        sys.exit(f"board: {header or 'timeout'}")
    length = int(header.split()[1])
    data = board.read(length)
    if len(data) != length:
        sys.exit(f"short capture: {len(data)} of {length} bytes")
    return data


# Statistics

def binary_entropy(p: float) -> float:
    if p <= 0.0 or p >= 1.0:
        return 0.0
    return -(p * math.log2(p) + (1 - p) * math.log2(1 - p))


def predictability_from_entropy(h: float) -> float:
    """Inverse of the binary entropy on P >= 0.5, by bisection."""
    low, high = 0.5, 1.0
    for _ in range(60):
        mid = (low + high) / 2
        if binary_entropy(mid) > h:
            low = mid
        else:
            high = mid
    return (low + high) / 2


def block_entropy(bits: np.ndarray, length: int) -> float:
    weights = 1 << np.arange(length)
    windows = np.lib.stride_tricks.sliding_window_view(bits, length)
    counts = np.bincount(windows @ weights, minlength=1 << length)
    p = counts[counts > 0] / counts.sum()
    return float(-(p * np.log2(p)).sum())


def analyze(data: bytes, order: int) -> dict:
    bits = np.unpackbits(np.frombuffer(data, dtype=np.uint8)).astype(np.int64)
    signed = 2.0 * bits - 1.0
    entropy = block_entropy(bits, order + 1) - block_entropy(bits, order)
    predictability = predictability_from_entropy(entropy)
    return {
        "ones": float(bits.mean()),
        "autocorr": float(np.mean(signed[:-1] * signed[1:]) - np.mean(signed) ** 2),
        "entropy": entropy,
        "relative_predictability": 2 * predictability - 1,
    }


# Jitter model: a square wave sampled at a nominal phase mu in the high half-cycle, with
# Gaussian phase noise of standard deviation sigma (in cycles). Averaging the probability of
# reading the expected level over mu gives the predictability of a sample.

_erf = np.vectorize(math.erf)


def model_predictability(sigma: float) -> float:
    mu = np.linspace(0.0, 0.5, 201)
    x = np.arange(-math.ceil(6 * sigma) - 1, math.ceil(6 * sigma) + 2)
    cdf = lambda edge: 0.5 * (1 + _erf((edge[None, :] - mu[:, None]) / (sigma * math.sqrt(2))))
    p_high = (cdf(x + 0.5) - cdf(x)).sum(axis=1)
    return float(p_high.mean())


def fit_jitter(decimations: list[int], relative: list[float]) -> float:
    """Least-squares fit of sigma_1 in sigma_k = sigma_1 * sqrt(k), via a lookup table."""
    table_sigma = np.logspace(-5, 1, 600)
    table_relative = np.array([2 * model_predictability(s) - 1 for s in table_sigma])
    ks = np.array(decimations, dtype=float)
    measured = np.array(relative)

    def model(sigma_1: float) -> np.ndarray:
        return np.interp(np.log(sigma_1 * np.sqrt(ks)), np.log(table_sigma), table_relative)

    candidates = np.logspace(-5, 0, 1000)
    errors = [np.sum((model(s) - measured) ** 2) for s in candidates]
    return float(candidates[int(np.argmin(errors))])


# Main

def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--port", default="/dev/ttyACM0")
    parser.add_argument("--out", type=Path, default=Path("captures"))
    parser.add_argument("--sources", type=int, nargs="+", default=list(SOURCES))
    parser.add_argument("--decimations", type=int, nargs="+", default=DEFAULT_DECIMATIONS)
    parser.add_argument("--order", type=int, default=8, help="context bits for the entropy estimate")
    parser.add_argument("--analyze-only", action="store_true", help="reuse captures in --out")
    args = parser.parse_args()

    args.out.mkdir(parents=True, exist_ok=True)
    board = None if args.analyze_only else open_board(args.port)

    results = {}
    print(f"{'source':<11} {'k':>5} {'ones':>7} {'autocorr':>9} {'H/bit':>7} {'2P-1':>8}")
    for source in args.sources:
        for k in args.decimations:
            path = args.out / f"src{source}_dec{k}.bin"
            if board:
                path.write_bytes(capture(board, source, k))
            if not path.exists():
                continue
            r = analyze(path.read_bytes(), args.order)
            results[(source, k)] = r
            print(f"{SOURCES[source]:<11} {k:>5} {r['ones']:7.4f} {r['autocorr']:9.4f} "
                  f"{r['entropy']:7.4f} {r['relative_predictability']:8.5f}")

    print("\nfitted jitter per 120 MHz sample interval (ring cycles, rms):")
    for source in args.sources:
        # Near-random points sit on the estimator's noise floor and would bias the fit.
        points = [(k, results[(source, k)]["relative_predictability"])
                  for k in args.decimations
                  if (source, k) in results and results[(source, k)]["relative_predictability"] > 0.1]
        if len(points) >= 3:
            ks, rel = zip(*points)
            print(f"  {SOURCES[source]:<11} {fit_jitter(list(ks), list(rel)):.5f}")

    print("\ncheck: enhanced A should match (tap A)^3 in relative predictability")
    for k in args.decimations:
        if (0, k) in results and (1, k) in results:
            tap = results[(0, k)]["relative_predictability"]
            enhanced = results[(1, k)]["relative_predictability"]
            print(f"  k={k:<5} tap^3={tap ** 3:8.5f}  enhanced={enhanced:8.5f}")


if __name__ == "__main__":
    main()