#!/usr/bin/env python3
"""Exhaustive encoding sweep on the 868 MHz data region.

For every plausible encoding, measure:
  * max |autocorrelation| over lags 1..12  (an interleaved multi-channel stream correlates
    at lag = number of channels, not at lag 1 -- the earlier test only looked at lags 1-2)
  * value entropy (a real sample stream is concentrated, random data is not)
  * share of the single most common value

Encodings: bit offset 0..15, width 4..16, MSB/LSB first, plus differential and Manchester
decoding of the bit stream before packing.
"""
import json, numpy as np, collections, math, itertools, sys

BLOB = []


def load(path):
    rows = [json.loads(l) for l in open(path)]
    rows.sort(key=lambda r: r["t"])
    regions = []
    for r in rows:
        bits = r["payload"]
        n = len(bits) // 8 * 8
        b = bytes(int(bits[i:i + 8], 2) for i in range(0, n, 8))
        if len(b) >= 54 and b[:4] == bytes.fromhex("1aa30918"):
            regions.append(b[7:54])          # header + counter/flags stripped
    return regions


def bits_of(regions, transform):
    out = []
    for r in regions:
        b = np.unpackbits(np.frombuffer(r, dtype=np.uint8))
        if transform == "diff":             # XOR-inverse of differential encoding
            b = np.concatenate(([b[0]], b[1:] ^ b[:-1]))
        elif transform == "manchester":
            # take the first chip of each pair; drop if pairs aren't complementary
            pairs = b[:len(b) // 2 * 2].reshape(-1, 2)
            ok = pairs[:, 0] != pairs[:, 1]
            if ok.mean() < 0.9:
                return None
            b = pairs[:, 0] ^ (pairs[:, 0] > pairs[:, 1])   # normalise polarity
        out.append(b)
    return out


def pack(bits, w, order):
    n = len(bits) // w
    m = bits[:n * w].reshape(n, w)
    if order == "msb":
        weights = 1 << np.arange(w - 1, -1, -1)
    else:
        weights = 1 << np.arange(0, w, 1)
    return (m * weights).sum(axis=1)


def score(vals, maxlag=12):
    v = np.asarray(vals, dtype=np.float64)
    if len(v) < 200 or v.std() == 0:
        return None
    v = (v - v.mean()) / v.std()
    rs = []
    for k in range(1, maxlag + 1):
        r = np.corrcoef(v[:-k], v[k:])[0, 1]
        rs.append(r if np.isfinite(r) else 0.0)
    c = collections.Counter(vals)
    n = len(vals)
    ent = -sum(q / n * math.log2(q / n) for q in c.values())
    top, cnt = c.most_common(1)[0]
    return {"max_r": float(np.max(np.abs(rs))), "argmax_lag": int(np.argmax(np.abs(rs))) + 1,
            "r1": rs[0], "entropy": ent, "top": top, "top_share": cnt / n}


def main(path):
    regions = load(path)
    print(f"{len(regions)} full 47-byte data regions (header stripped)\n")
    results = []
    for transform in ("none", "diff", "manchester"):
        bl = bits_of(regions, transform)
        if bl is None:
            print(f"{transform}: rejected (pairs not complementary)")
            continue
        joined = np.concatenate(bl)
        for w in (4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16):
            for off in range(0, 16):
                seg = joined[off:]
                for order in ("msb", "lsb"):
                    vals = pack(seg, w, order)
                    s = score(vals)
                    if s:
                        s.update({"transform": transform, "w": w, "off": off, "order": order})
                        results.append(s)
    results.sort(key=lambda r: -r["max_r"])
    print(f"{len(results)} encodings scored.  Top 12 by |autocorrelation| at any lag 1-12:\n")
    print(f"{'|r|':>6} {'lag':>4} {'r1':>7} {'H':>5} {'top%':>6}  {'width':>5} {'off':>4} {'order':>4} transform")
    for r in results[:12]:
        print(f"{r['max_r']:6.3f} {r['argmax_lag']:4d} {r['r1']:7.3f} {r['entropy']:5.2f} "
              f"{r['top_share']*100:5.1f}% {r['w']:5d} {r['off']:4d} {r['order']:>4} {r['transform']}")
    # for reference: what does pure random data score?
    rng = np.random.default_rng(0)
    ref = [score(rng.integers(0, 256, len(regions) * 47 // 2)) for _ in range(20)]
    print(f"\nreference, uniform random bytes: max|r| {np.mean([x['max_r'] for x in ref]):.3f} "
          f"(95th pct {np.percentile([x['max_r'] for x in ref],95):.3f}), entropy {np.mean([x['entropy'] for x in ref]):.2f}")


if __name__ == "__main__":
    main(sys.argv[1] if len(sys.argv) > 1 else "/home/tmco/sdr/deep/bulk4.jsonl")