#!/usr/bin/env python3
"""Ring-oscillator experiments against the STRIX capture sketch (arduino/strix_capture).

Subcommands:
  sweep   capture sources at increasing decimation k and fit the accumulated jitter
  freq    measure every ring's frequency
  xor     capture the XOR of rings 0..n-1 for n = 1..16 and compare with the product law
  pairs   capture ring pairs simultaneously and test them 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. 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 the XOR combining conservatively.
"""

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

import numpy as np
import serial

RINGS = 16
CLOCK_HZ = 120e6
SOURCE_ENHANCED, SOURCE_TAP, SOURCE_XOR = 0x00, 0x10, 0x20

REG_FREQ_GATE = 0x05
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_FREQ_SELECT, REG_FREQ_COUNT3 = 0x11, 0x12
CONTROL_START_FREQ = 0x04
STATUS_FREQ_VALID = 0x08


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


def source_name(code: int) -> str:
    if code >= SOURCE_XOR:
        return f"xor 0-{code - SOURCE_XOR}"
    if code >= SOURCE_TAP:
        return f"tap {code - SOURCE_TAP}"
    return f"enhanced {code}"


# Board communication

class Board:
    def __init__(self, port: str):
        self.serial = serial.Serial(port, 115200, timeout=200)
        # The sketch prints its identity once when the port opens; drain it so replies stay
        # aligned with commands.
        time.sleep(0.5)
        self.serial.reset_input_buffer()
        identity = self.command("id")
        if not identity.startswith("STRX v2"):
            sys.exit(f"unexpected board identity: {identity!r} (needs the v2 FPGA image)")
        print(f"board: {identity}")

    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, ring: int) -> float:
        self.write(REG_FREQ_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 print_header() -> None:
    print(f"{'source':<12} {'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:<12} {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_sweep(board: Board | None, args) -> None:
    results = {}
    print_header()
    for source in args.sources:
        for k in args.decimations:
            data = cached_capture(board, args.out / f"sweep_s{source:02x}_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), 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):<12} {fit_jitter(list(ks), list(rel)):.5f}")


def run_freq(board: Board | None, args) -> None:
    if board is None:
        sys.exit("freq needs the board")
    print(f"{'ring':>4} {'length':>6} {'MHz':>9}")
    for ring in range(RINGS):
        print(f"{ring:>4} {ring_length(ring):>6} {board.frequency(ring) / 1e6:9.2f}")


def run_xor(board: Board | None, args) -> None:
    single = {}
    for ring in range(RINGS):
        data = cached_capture(board, args.out / f"single_s{ring:02x}_k{args.k}.bin", SOURCE_ENHANCED + ring, args.k)
        if data is not None:
            single[ring] = analyze(to_bits(data), args.order)["relative"]

    print(f"{'n':>3} {'ones':>7} {'H/bit':>7} {'2P-1':>8} {'product':>8}")
    product = 1.0
    for n in range(1, RINGS + 1):
        product *= single.get(n - 1, float("nan"))
        data = cached_capture(board, args.out / f"xor_n{n}_k{args.k}.bin", SOURCE_XOR + n - 1, args.k)
        if data is None:
            continue
        r = analyze(to_bits(data), args.order)
        print(f"{n:>3} {r['ones']:7.4f} {r['entropy']:7.4f} {r['relative']:8.5f} {product:8.5f}")
    print("\n'product' is the prediction for independent rings: the product of the single rings' 2P-1.")
    print("Measured values sit on the estimator's noise floor (about 0.03 for 256 Kbit) once it is reached.")


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 run_pairs(board: Board | None, args) -> None:
    rng = np.random.default_rng(0)
    pairs = list(combinations(range(RINGS), 2))
    print(f"{'pair':>7} {'corr':>8} {'null sd':>8} {'z':>7}")
    flagged = 0
    for i, j in pairs:
        data = cached_capture(board, args.out / f"pair_{i}_{j}_k{args.k}.bin",
                              SOURCE_ENHANCED + i, args.k, SOURCE_ENHANCED + j)
        if data is None:
            continue
        bits = to_bits(data).astype(float)
        a, b = bits[0::2], bits[1::2]
        # Null distribution: the same two streams misaligned by large random offsets.
        null = [correlation(a, np.roll(b, s)) for s in rng.integers(len(b) // 4, 3 * len(b) // 4, 20)]
        c, sd = correlation(a, b), float(np.std(null))
        z = c / sd if sd > 0 else float("inf")
        mark = "  <-- dependent?" if abs(z) > 5 else ""
        flagged += bool(mark)
        print(f"{i:>3}-{j:<3} {c:8.4f} {sd:8.4f} {z:7.2f}{mark}")
    print(f"\n{flagged} of {len(pairs)} 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=8, 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)

    sweep = sub.add_parser("sweep")
    sweep.add_argument("--sources", type=lambda s: int(s, 0), nargs="+",
                       default=[SOURCE_TAP, SOURCE_ENHANCED, SOURCE_ENHANCED + 1, SOURCE_XOR + 1])
    sweep.add_argument("--decimations", type=int, nargs="+",
                       default=[1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096])
    sub.add_parser("freq")
    xor = sub.add_parser("xor")
    xor.add_argument("-k", type=int, default=1, help="decimation")
    pairs = sub.add_parser("pairs")
    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)
    {"sweep": run_sweep, "freq": run_freq, "xor": run_xor, "pairs": run_pairs}[args.command](board, args)


if __name__ == "__main__":
    main()