#!/usr/bin/env python3
"""Entropy generator experiments against the STRIX capture sketch (arduino/strix_capture).

Subcommands:
  freq     measure ring frequencies in every generator
  rings    estimate each ring's relative predictability at k=1 (sizing input)
  sweep    capture sources at increasing decimation k and fit the accumulated jitter
  output   check every generator's spatial XOR and processed output for bias and correlation
  pairs    capture the same ring index in two generators simultaneously and test for correlation

Entropy per bit is estimated from conditional block entropies and converted to a
predictability P; 2P-1 is the relative predictability, which multiplies when independent
sources are XORed. Statistics of XOR-combined streams overstate their randomness, because
short contexts cannot see the long-period structure of mixed rings, so size the XOR
combining from single-ring values (`rings`), not from `output`.

The jitter fit models a square wave sampled with Gaussian phase noise whose standard
deviation grows with sqrt(k). On synthetic captures it overestimates by a factor of about
1.2-1.3, so divide by 1.3 when sizing conservatively.
"""

import argparse
import math
import sys
import time
from itertools import combinations
from pathlib import Path

import numpy as np
import serial

CLOCK_HZ = 120e6

# Capture sources for generator g
SOURCE_OUTPUT, SOURCE_SPATIAL, SOURCE_RING, SOURCE_TAP = 0x00, 0x04, 0x08, 0x0C

REG_FREQ_GATE, REG_GENERATORS, REG_RINGS = 0x05, 0x06, 0x07
REG_SOURCE_A, REG_DECIM_HIGH, REG_DECIM_LOW = 0x08, 0x09, 0x0A
REG_CONTROL, REG_STATUS, REG_SOURCE_B, REG_MODE = 0x0B, 0x0C, 0x0D, 0x0E
REG_RING_SELECT, REG_FREQ_COUNT3 = 0x11, 0x12
REG_GEN_ENABLE, REG_TEMPORAL, REG_FREQ_GEN = 0x16, 0x17, 0x18
CONTROL_START_FREQ = 0x04
STATUS_FREQ_VALID = 0x08


def ring_length(ring: int) -> int:
    return 9 + 3 * (ring % 3)


def source_name(code: int, ring: int) -> str:
    generator, kind = code & 0x03, code & 0x0C
    names = {SOURCE_OUTPUT: "output", SOURCE_SPATIAL: "spatial",
             SOURCE_RING: f"ring {ring}", SOURCE_TAP: f"tap {ring}"}
    return f"g{generator} {names[kind]}"


# Board communication

class Board:
    def __init__(self, port: str):
        # The sketch also prints its identity when the port opens, possibly after our "id"
        # reply if the board is still booting. Take the first identity line, then drain
        # until the port is quiet so replies stay aligned with commands.
        self.serial = serial.Serial(port, 115200, timeout=5)
        self.serial.write(b"id\n")
        identity = ""
        while not identity.startswith("STRX"):
            identity = self.serial.readline().decode(errors="replace").strip()
            if not identity:
                sys.exit("no answer from the capture sketch")
        self.serial.timeout = 0.5
        while self.serial.readline():
            pass
        self.serial.timeout = 200
        if not identity.startswith("STRX v3"):
            sys.exit(f"unexpected board identity: {identity!r} (needs the v3 FPGA image)")
        self.generators = self.read(REG_GENERATORS)
        self.rings = self.read(REG_RINGS)
        print(f"board: {identity}, {self.generators} generators x {self.rings} rings")

    def command(self, line: str) -> str:
        self.serial.write(f"{line}\n".encode())
        reply = self.serial.readline().decode(errors="replace").strip()
        if not reply or reply.startswith("ERR"):
            sys.exit(f"board: {reply or 'timeout'} (command {line!r})")
        return reply

    def read(self, register: int) -> int:
        return int(self.command(f"r {register}"))

    def write(self, register: int, value: int) -> None:
        self.command(f"w {register} {value}")

    def capture(self, source_a: int, decimation: int, source_b: int | None = None) -> bytes:
        self.write(REG_SOURCE_A, source_a)
        self.write(REG_SOURCE_B, source_b if source_b is not None else 0)
        self.write(REG_MODE, 1 if source_b is not None else 0)
        self.write(REG_DECIM_HIGH, decimation >> 8)
        self.write(REG_DECIM_LOW, decimation & 0xFF)
        header = self.command("cap")
        if not header.startswith("DATA "):
            sys.exit(f"board: {header}")
        length = int(header.split()[1])
        data = self.serial.read(length)
        if len(data) != length:
            sys.exit(f"short capture: {len(data)} of {length} bytes")
        return data

    def frequency(self, generator: int, ring: int) -> float:
        self.write(REG_FREQ_GEN, generator)
        self.write(REG_RING_SELECT, ring)
        self.write(REG_CONTROL, CONTROL_START_FREQ)
        while not self.read(REG_STATUS) & STATUS_FREQ_VALID:
            pass
        count = 0
        for i in range(4):
            count = (count << 8) | self.read(REG_FREQ_COUNT3 + i)
        gate_seconds = (1 << self.read(REG_FREQ_GATE)) / CLOCK_HZ
        return count * 256 / gate_seconds


def cached_capture(board: Board | None, path: Path, *args) -> bytes | None:
    if board:
        path.write_bytes(board.capture(*args))
    return path.read_bytes() if path.exists() else None


# Statistics

def to_bits(data: bytes) -> np.ndarray:
    return np.unpackbits(np.frombuffer(data, dtype=np.uint8)).astype(np.int64)


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(bits: np.ndarray, order: int) -> dict:
    signed = 2.0 * bits - 1.0
    entropy = block_entropy(bits, order + 1) - block_entropy(bits, order)
    return {
        "ones": float(bits.mean()),
        "autocorr": float(np.mean(signed[:-1] * signed[1:]) - np.mean(signed) ** 2),
        "entropy": entropy,
        "relative": 2 * predictability_from_entropy(entropy) - 1,
    }


def correlation(a: np.ndarray, b: np.ndarray) -> float:
    a = a - a.mean()
    b = b - b.mean()
    return float(np.dot(a, b) / math.sqrt(np.dot(a, a) * np.dot(b, b)))


def print_header() -> None:
    print(f"{'source':<14} {'k':>5} {'ones':>7} {'autocorr':>9} {'H/bit':>7} {'2P-1':>8}")


def print_row(name: str, k: int, r: dict) -> None:
    print(f"{name:<14} {k:>5} {r['ones']:7.4f} {r['autocorr']:9.4f} {r['entropy']:7.4f} {r['relative']:8.5f}")


# 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))))
    return float((cdf(x + 0.5) - cdf(x)).sum(axis=1).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))])


# Subcommands

def run_freq(board: Board | None, args) -> None:
    if board is None:
        sys.exit("freq needs the board")
    print(f"{'gen':>3} {'length':>6} {'rings':>5} {'min MHz':>8} {'mean MHz':>9} {'max MHz':>8}")
    for g in range(board.generators):
        by_length: dict[int, list[float]] = {}
        for ring in range(board.rings):
            by_length.setdefault(ring_length(ring), []).append(board.frequency(g, ring) / 1e6)
        for length, freqs in sorted(by_length.items()):
            print(f"{g:>3} {length:>6} {len(freqs):>5} {min(freqs):8.2f} {np.mean(freqs):9.2f} {max(freqs):8.2f}")


def run_rings(board: Board | None, args) -> None:
    values = []
    print(f"{'ring':>4} {'length':>6} {'ones':>7} {'2P-1':>8}")
    for ring in args.rings:
        path = args.out / f"ring_g{args.generator}_r{ring}_k1.bin"
        if board:
            board.write(REG_RING_SELECT, ring)
        data = cached_capture(board, path, SOURCE_RING + args.generator, 1)
        if data is None:
            continue
        r = analyze(to_bits(data), args.order)
        values.append(r["relative"])
        print(f"{ring:>4} {ring_length(ring):>6} {r['ones']:7.4f} {r['relative']:8.5f}")
    if values:
        # Independent sources multiply: the generator's spatial XOR has 2P-1 = product over rings.
        log_spatial = float(np.sum(np.log(values)))
        print(f"\n2P-1 per ring at k=1 (order {args.order}): mean {np.mean(values):.4f}, worst {max(values):.4f}")
        print(f"spatial XOR of these {len(values)} rings: 2P-1 = {math.exp(log_spatial):.3e}")
        for temporal in range(1, 9):
            print(f"  temporal XOR {temporal}: output 2P-1 = {math.exp(temporal * log_spatial):.3e}, "
                  f"{CLOCK_HZ / temporal / 1e6:5.1f} Mbps per generator")


def run_sweep(board: Board | None, args) -> None:
    if board:
        board.write(REG_RING_SELECT, args.ring)
    results = {}
    print_header()
    for source in args.sources:
        for k in args.decimations:
            data = cached_capture(board, args.out / f"sweep_s{source:02x}_r{args.ring}_k{k}.bin", source, k)
            if data is None:
                continue
            results[(source, k)] = r = analyze(to_bits(data), args.order)
            print_row(source_name(source, args.ring), k, r)

    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"]) for k in args.decimations
                  if (source, k) in results and results[(source, k)]["relative"] > 0.1]
        if len(points) >= 3:
            ks, rel = zip(*points)
            print(f"  {source_name(source, args.ring):<14} {fit_jitter(list(ks), list(rel)):.5f}")


def run_output(board: Board | None, args) -> None:
    # Longer contexts are undersampled at 256 Kbit and read falsely low on random data.
    order = min(args.order, 8)
    temporal = board.read(REG_TEMPORAL) if board else args.temporal
    generators = board.generators if board else 4
    print(f"temporal XOR = {temporal}; output captured with decimation {temporal}")
    print("ones should be within about +-0.003 of 0.5 (3 sigma for 256 Kbit)\n")
    print_header()
    for g in range(generators):
        for source, k in [(SOURCE_SPATIAL + g, 1), (SOURCE_OUTPUT + g, temporal)]:
            data = cached_capture(board, args.out / f"output_s{source:02x}_k{k}.bin", source, k)
            if data is not None:
                print_row(source_name(source, 0), k, analyze(to_bits(data), order))


def run_pairs(board: Board | None, args) -> None:
    generators = board.generators if board else 4
    rng = np.random.default_rng(0)
    print(f"{'ring':>4} {'gens':>5} {'corr':>8} {'null sd':>8} {'z':>7}")
    flagged = total = 0
    for ring in args.rings:
        if board:
            board.write(REG_RING_SELECT, ring)
        for a, b in combinations(range(generators), 2):
            data = cached_capture(board, args.out / f"pair_r{ring}_g{a}{b}_k{args.k}.bin",
                                  SOURCE_RING + a, args.k, SOURCE_RING + b)
            if data is None:
                continue
            bits = to_bits(data).astype(float)
            x, y = bits[0::2], bits[1::2]
            # Null distribution: the same two streams misaligned by large random offsets.
            null = [correlation(x, np.roll(y, s)) for s in rng.integers(len(y) // 4, 3 * len(y) // 4, 20)]
            c, sd = correlation(x, y), float(np.std(null))
            z = c / sd if sd > 0 else float("inf")
            mark = "  <-- dependent?" if abs(z) > 5 else ""
            flagged += bool(mark)
            total += 1
            print(f"{ring:>4} {a}-{b:<3} {c:8.4f} {sd:8.4f} {z:7.2f}{mark}")
    print(f"\n{flagged} of {total} pairs exceed |z| > 5")


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("--order", type=int, default=16, help="context bits for the entropy estimate")
    parser.add_argument("--analyze-only", action="store_true", help="reuse captures in --out")
    sub = parser.add_subparsers(dest="command", required=True)

    sub.add_parser("freq")
    rings = sub.add_parser("rings")
    rings.add_argument("--generator", type=int, default=0)
    rings.add_argument("--rings", type=int, nargs="+", default=list(range(90)))
    sweep = sub.add_parser("sweep")
    sweep.add_argument("--ring", type=int, default=0, help="ring for the ring/tap sources")
    sweep.add_argument("--sources", type=lambda s: int(s, 0), nargs="+",
                       default=[SOURCE_TAP, SOURCE_RING, SOURCE_SPATIAL])
    sweep.add_argument("--decimations", type=int, nargs="+",
                       default=[1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096])
    output = sub.add_parser("output")
    output.add_argument("--temporal", type=int, default=5, help="for --analyze-only")
    pairs = sub.add_parser("pairs")
    pairs.add_argument("--rings", type=int, nargs="+", default=list(range(0, 90, 9)))
    pairs.add_argument("-k", type=int, default=64, help="decimation")

    args = parser.parse_args()
    args.out.mkdir(parents=True, exist_ok=True)
    board = None if args.analyze_only else Board(args.port)
    handlers = {"freq": run_freq, "rings": run_rings, "sweep": run_sweep,
                "output": run_output, "pairs": run_pairs}
    handlers[args.command](board, args)


if __name__ == "__main__":
    main()