import json, time
import numpy as np

# BBS-C(2), Table 1 in article.md. table[q, s] = (next_q, output_s)
TABLE = np.array([[[0, 0], [1, 0]],
                  [[0, 1], [2, 0]],
                  [[1, 1], [2, 1]]], dtype=np.int64)
WEIGHTS = np.arange(3, dtype=np.int64)


def check_table():
    pairs = [(q, s) for q in range(3) for s in range(2)]
    images = [tuple(TABLE[q, s]) for q, s in pairs]
    errors = []
    for q, s in pairs:
        nq, y = TABLE[q, s]
        errors.append(int(WEIGHTS[q] + s - WEIGHTS[nq] - y))
    return len(set(images)) == 6, errors


def scan(bits, q0=0):
    q = int(q0)
    out = np.empty(len(bits), dtype=np.int64)
    for i, s in enumerate(np.asarray(bits, dtype=np.int64)):
        q, out[i] = TABLE[q, int(s)]
    return out, q


def inverse_scan(out, q_final):
    inverse = {tuple(TABLE[q, s]): (q, s)
               for q in range(3) for s in range(2)}
    q, recovered = int(q_final), []
    for y in np.asarray(out)[::-1]:
        q, s = inverse[(q, int(y))]
        recovered.append(s)
    return np.asarray(recovered[::-1], dtype=np.int64), q


def pair_freq(bits):
    b = np.asarray(bits, dtype=np.int64)
    c = np.zeros(4, dtype=np.float64)
    for a, z in zip(b[:-1], b[1:]):
        c[2 * int(a) + int(z)] += 1
    return c / max(1, len(b) - 1)


def exhaustive_inverse_check(max_len=12):
    # Checks every binary word up to max_len, not merely a random sample.
    count = 0
    for n in range(max_len + 1):
        for code in range(1 << n):
            x = np.array([(code >> i) & 1 for i in range(n)], dtype=np.int64)
            y, qn = scan(x)
            xr, q0 = inverse_scan(y, qn)
            if not np.array_equal(x, xr) or q0 != 0:
                return False, count
            count += 1
    return True, count


def run_mixer(n=4096, seed=7):
    rng = np.random.default_rng(seed)
    x = rng.integers(0, 2, size=n, dtype=np.int64)
    t0 = time.perf_counter()
    y, qn = scan(x)
    elapsed = time.perf_counter() - t0
    recovered, q0 = inverse_scan(y, qn)
    conservation = int(WEIGHTS[0] + x.sum() - WEIGHTS[qn] - y.sum())
    return {
        "n": n, "inverse_exact": bool(np.array_equal(x, recovered) and q0 == 0),
        "conservation_total_error": conservation,
        "input_mean": float(x.mean()), "output_mean": float(y.mean()),
        "pair_l2_drift": float(np.linalg.norm(pair_freq(y) - pair_freq(x))),
        "forward_seconds_numpy": elapsed,
        "carrier_values_saved_for_reverse": 1,
        "carrier_state_values": 3,
    }


def gru_control(lengths=(128, 512, 2048, 8192), hidden=3, seed=7):
    try:
        import torch
        torch.manual_seed(seed)
        device = "cuda" if torch.cuda.is_available() else "cpu"
        def run(dev):
            rows = []
            for n in lengths:
                x = torch.randint(0, 2, (1, n, 1), device=dev, dtype=torch.float32)
                model = torch.nn.GRU(1, hidden, batch_first=True).to(dev)
                t0 = time.perf_counter(); h, _ = model(x)
                if dev == "cuda": torch.cuda.synchronize()
                elapsed = time.perf_counter() - t0
                # h is the per-token recurrent activation required by a generic
                # implementation for downstream training/backpropagation.
                rows.append({"n": n, "seconds": elapsed,
                             "activation_elements": int(h.numel()),
                             "activation_bytes_fp32": int(h.numel() * 4),
                             "device": dev})
            return rows
        try:
            return run(device)
        except Exception:
            return run("cpu")
    except Exception as e:
        return [{"error": "torch control unavailable: " + repr(e)}]


def main():
    bijective, errors = check_table()
    exhaustive, checked = exhaustive_inverse_check()
    result = {
        "local_bijective": bijective,
        "local_conservation_errors": errors,
        "exhaustive_inverse_check": exhaustive,
        "words_checked": checked,
        "mixer": run_mixer(),
        "gru": gru_control(),
    }
    with open("results.json", "w") as f:
        json.dump(result, f, indent=2)
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()