import argparse
import json
import sys
import time

import numpy as np

def stabilize(h):
    h = h.astype(np.int64, copy=True)
    n, m = h.shape
    p = np.zeros((n + 2, m + 2), dtype=np.int64)
    while True:
        t = h >> 2
        if not t.any():
            return h
        h &= 3
        p[1:-1, 1:-1] = t
        h += p[:-2, 1:-1] + p[2:, 1:-1] + p[1:-1, :-2] + p[1:-1, 2:]

def identity(n, k):
    full = np.full((n, n), k, dtype=np.int64)
    return stabilize(full - stabilize(full))

def law_checks(e, rng, trials=25):
    n = e.shape[0]
    ok = 0
    for _ in range(trials):
        c = stabilize(np.full((n, n), 3, dtype=np.int64) + rng.integers(0, 4, size=(n, n)))
        ok += bool(np.array_equal(stabilize(e + c), c))
    return ok, trials, bool(np.array_equal(stabilize(e + e), e))

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("sizes", nargs="+", type=int)
    ap.add_argument("--law-max", type=int, default=128)
    ap.add_argument("--save-dir", default=None)
    ap.add_argument("--skip-k7-above", type=int, default=10**9, help="compute the k=7 path only up to this N")
    a = ap.parse_args()
    rng = np.random.default_rng(10502)
    for n in a.sizes:
        t0 = time.time()
        e6 = identity(n, 6)
        t6 = time.time() - t0
        row = {"N": n, "density": float(e6.mean()), "grains": int(e6.sum()),
               "frac": [float((e6 == v).mean()) for v in range(4)], "secs_k6": round(t6, 1)}
        if n <= a.skip_k7_above:
            e7 = identity(n, 7)
            row["k7_equal"] = bool(np.array_equal(e6, e7))
        if n <= a.law_max:
            row["law_ok"], row["law_trials"], row["idempotent"] = law_checks(e6, rng)
        row["symmetric_d4"] = bool(np.array_equal(e6, e6.T) and np.array_equal(e6, e6[::-1, :]) and np.array_equal(e6, e6[:, ::-1]))
        if a.save_dir:
            np.save("%s/identity_N%d.npy" % (a.save_dir, n), e6.astype(np.int8))
        print(json.dumps(row), flush=True)
    return 0

if __name__ == "__main__":
    sys.exit(main())